Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Test wrapping of response for async request

Browse Source 
The following adjustments are also made as a result:
- Use int to check if lock is held and unlock is needed, given that
for non-async requests we don't need to obtain a lock.
- Protect access methods getOutputStream and getWriter with the
same locking and state checks.

Closes gh-32340
This commit is contained in:
rstoyanchev
2024-03-05 11:41:51 +00:00
parent 822e2447a0
commit ef0717935b
4 changed files with 502 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
spring-web/src/main/java/org/springframework/web/context/request/async/StandardServletAsyncWebRequest.java
View File

@@ -21,7 +21,6 @@ import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;
@@ -189,7 +188,7 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
public void onError(AsyncEvent event) throws IOException {
this.stateLock.lock();
try {
transitionToErrorState();
this.state = State.ERROR;
Throwable ex = event.getThrowable();
this.exceptionHandlers.forEach(consumer -> consumer.accept(ex));
}
@@ -198,12 +197,6 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
}
}
private void transitionToErrorState() {
if (!isAsyncComplete()) {
this.state = State.ERROR;
}
}
@Override
public void onComplete(AsyncEvent event) throws IOException {
this.stateLock.lock();
@@ -218,8 +211,17 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
}
/**
* Package private access for testing only.
*/
ReentrantLock stateLock() {
return this.stateLock;
}
/**
* Response wrapper to wrap the output stream with {@link LifecycleServletOutputStream}.
* @since 5.3.33
*/
private static final class LifecycleHttpServletResponse extends HttpServletResponseWrapper {
@@ -242,26 +244,44 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
@Override
public ServletOutputStream getOutputStream() throws IOException {
if (this.outputStream == null) {
Assert.notNull(this.asyncWebRequest, "Not initialized");
ServletOutputStream delegate = getResponse().getOutputStream();
this.outputStream = new LifecycleServletOutputStream(delegate, this);
int level = obtainLockAndCheckState();
try {
if (this.outputStream == null) {
Assert.notNull(this.asyncWebRequest, "Not initialized");
ServletOutputStream delegate = getResponse().getOutputStream();
this.outputStream = new LifecycleServletOutputStream(delegate, this);
}
}
catch (IOException ex) {
handleIOException(ex, "Failed to get ServletResponseOutput");
}
finally {
releaseLock(level);
}
return this.outputStream;
}
@Override
public PrintWriter getWriter() throws IOException {
if (this.writer == null) {
Assert.notNull(this.asyncWebRequest, "Not initialized");
this.writer = new LifecyclePrintWriter(getResponse().getWriter(), this.asyncWebRequest);
int level = obtainLockAndCheckState();
try {
if (this.writer == null) {
Assert.notNull(this.asyncWebRequest, "Not initialized");
this.writer = new LifecyclePrintWriter(getResponse().getWriter(), this.asyncWebRequest);
}
}
catch (IOException ex) {
handleIOException(ex, "Failed to get PrintWriter");
}
finally {
releaseLock(level);
}
return this.writer;
}
@Override
public void flushBuffer() throws IOException {
obtainLockAndCheckState();
int level = obtainLockAndCheckState();
try {
getResponse().flushBuffer();
}
@@ -269,32 +289,40 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
handleIOException(ex, "ServletResponse failed to flushBuffer");
}
finally {
releaseLock();
releaseLock(level);
}
}
private void obtainLockAndCheckState() throws AsyncRequestNotUsableException {
/**
* Return 0 if checks passed and lock is not needed, 1 if checks passed
* and lock is held, or raise AsyncRequestNotUsableException.
*/
private int obtainLockAndCheckState() throws AsyncRequestNotUsableException {
Assert.notNull(this.asyncWebRequest, "Not initialized");
if (this.asyncWebRequest.state != State.NEW) {
this.asyncWebRequest.stateLock.lock();
if (this.asyncWebRequest.state != State.ASYNC) {
this.asyncWebRequest.stateLock.unlock();
throw new AsyncRequestNotUsableException("Response not usable after " +
(this.asyncWebRequest.state == State.COMPLETED ?
"async request completion" : "onError notification") + ".");
}
if (this.asyncWebRequest.state == State.NEW) {
return 0;
}
this.asyncWebRequest.stateLock.lock();
if (this.asyncWebRequest.state == State.ASYNC) {
return 1;
}
this.asyncWebRequest.stateLock.unlock();
throw new AsyncRequestNotUsableException("Response not usable after " +
(this.asyncWebRequest.state == State.COMPLETED ?
"async request completion" : "response errors") + ".");
}
void handleIOException(IOException ex, String msg) throws AsyncRequestNotUsableException {
Assert.notNull(this.asyncWebRequest, "Not initialized");
this.asyncWebRequest.transitionToErrorState();
throw new AsyncRequestNotUsableException(msg, ex);
this.asyncWebRequest.state = State.ERROR;
throw new AsyncRequestNotUsableException(msg + ": " + ex.getMessage(), ex);
}
void releaseLock() {
void releaseLock(int level) {
Assert.notNull(this.asyncWebRequest, "Not initialized");
if (this.asyncWebRequest.state != State.NEW) {
if (level > 0) {
this.asyncWebRequest.stateLock.unlock();
}
}
@@ -304,6 +332,7 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
/**
* Wraps a ServletOutputStream to prevent use after Servlet container onError
* notifications, and after async request completion.
* @since 5.3.33
*/
private static final class LifecycleServletOutputStream extends ServletOutputStream {
@@ -328,7 +357,7 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
@Override
public void write(int b) throws IOException {
this.response.obtainLockAndCheckState();
int level = this.response.obtainLockAndCheckState();
try {
this.delegate.write(b);
}
@@ -336,12 +365,12 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
this.response.handleIOException(ex, "ServletOutputStream failed to write");
}
finally {
this.response.releaseLock();
this.response.releaseLock(level);
}
}
public void write(byte[] buf, int offset, int len) throws IOException {
this.response.obtainLockAndCheckState();
int level = this.response.obtainLockAndCheckState();
try {
this.delegate.write(buf, offset, len);
}
@@ -349,13 +378,13 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
this.response.handleIOException(ex, "ServletOutputStream failed to write");
}
finally {
this.response.releaseLock();
this.response.releaseLock(level);
}
}
@Override
public void flush() throws IOException {
this.response.obtainLockAndCheckState();
int level = this.response.obtainLockAndCheckState();
try {
this.delegate.flush();
}
@@ -363,13 +392,13 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
this.response.handleIOException(ex, "ServletOutputStream failed to flush");
}
finally {
this.response.releaseLock();
this.response.releaseLock(level);
}
}
@Override
public void close() throws IOException {
this.response.obtainLockAndCheckState();
int level = this.response.obtainLockAndCheckState();
try {
this.delegate.close();
}
@@ -377,7 +406,7 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
this.response.handleIOException(ex, "ServletOutputStream failed to close");
}
finally {
this.response.releaseLock();
this.response.releaseLock(level);
}
}
@@ -387,6 +416,7 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
/**
* Wraps a PrintWriter to prevent use after Servlet container onError
* notifications, and after async request completion.
* @since 5.3.33
*/
private static final class LifecyclePrintWriter extends PrintWriter {
@@ -402,24 +432,26 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
@Override
public void flush() {
if (tryObtainLockAndCheckState()) {
int level = tryObtainLockAndCheckState();
if (level > -1) {
try {
this.delegate.flush();
}
finally {
releaseLock();
releaseLock(level);
}
}
}
@Override
public void close() {
if (tryObtainLockAndCheckState()) {
int level = tryObtainLockAndCheckState();
if (level > -1) {
try {
this.delegate.close();
}
finally {
releaseLock();
releaseLock(level);
}
}
}
@@ -431,24 +463,26 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
@Override
public void write(int c) {
if (tryObtainLockAndCheckState()) {
int level = tryObtainLockAndCheckState();
if (level > -1) {
try {
this.delegate.write(c);
}
finally {
releaseLock();
releaseLock(level);
}
}
}
@Override
public void write(char[] buf, int off, int len) {
if (tryObtainLockAndCheckState()) {
int level = tryObtainLockAndCheckState();
if (level > -1) {
try {
this.delegate.write(buf, off, len);
}
finally {
releaseLock();
releaseLock(level);
}
}
}
@@ -460,12 +494,13 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
@Override
public void write(String s, int off, int len) {
if (tryObtainLockAndCheckState()) {
int level = tryObtainLockAndCheckState();
if (level > -1) {
try {
this.delegate.write(s, off, len);
}
finally {
releaseLock();
releaseLock(level);
}
}
}
@@ -475,33 +510,28 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
this.delegate.write(s);
}
private boolean tryObtainLockAndCheckState() {
if (state() == State.NEW) {
return true;
/**
* Return 0 if checks passed and lock is not needed, 1 if checks passed
* and lock is held, and -1 if checks did not pass.
*/
private int tryObtainLockAndCheckState() {
if (this.asyncWebRequest.state == State.NEW) {
return 0;
}
if (stateLock().tryLock()) {
if (state() == State.ASYNC) {
return true;
}
stateLock().unlock();
this.asyncWebRequest.stateLock.lock();
if (this.asyncWebRequest.state == State.ASYNC) {
return 1;
}
return false;
this.asyncWebRequest.stateLock.unlock();
return -1;
}
private void releaseLock() {
if (state() != State.NEW) {
stateLock().unlock();
private void releaseLock(int level) {
if (level > 0) {
this.asyncWebRequest.stateLock.unlock();
}
}
private State state() {
return this.asyncWebRequest.state;
}
private Lock stateLock() {
return this.asyncWebRequest.stateLock;
}
// Plain delegates
@Override
@@ -639,28 +669,28 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
/**
* Represents a state for {@link StandardServletAsyncWebRequest} to be in.
* <p><pre>
* NEW
* |
* v
* ASYNC----> +
* | |
* v |
* ERROR |
* | |
* v |
* COMPLETED <--+
* +------ NEW
* | |
* | v
* | ASYNC ----> +
* | | |
* | v |
* +----> ERROR |
* | |
* v |
* COMPLETED <---+
* </pre>
* @since 5.3.33
*/
private enum State {
/** New request (thas may not do async handling). */
/** New request (may not start async handling). */
NEW,
/** Async handling has started. */
ASYNC,
/** onError notification received, or ServletOutputStream failed. */
/** ServletOutputStream failed, or onError notification received. */
ERROR,
/** onComplete notification received. */