From 9de992524da1a6a66b001c935ff23c032c88f8f6 Mon Sep 17 00:00:00 2001 From: Lachlan Roberts Date: Wed, 13 Dec 2023 11:41:35 +1100 Subject: [PATCH] Do not use HttpStream.Wrapper in SizeLimitHandler Signed-off-by: Lachlan Roberts --- .../jetty/server/SizeLimitHandler.java | 121 ++++++++++-------- .../server/handler/SizeLimitHandlerTest.java | 25 ++++ 2 files changed, 91 insertions(+), 55 deletions(-) diff --git a/jetty-core/jetty-server/src/main/java/org/eclipse/jetty/server/SizeLimitHandler.java b/jetty-core/jetty-server/src/main/java/org/eclipse/jetty/server/SizeLimitHandler.java index 862f1459015d..6e7c7749d143 100644 --- a/jetty-core/jetty-server/src/main/java/org/eclipse/jetty/server/SizeLimitHandler.java +++ b/jetty-core/jetty-server/src/main/java/org/eclipse/jetty/server/SizeLimitHandler.java @@ -21,7 +21,6 @@ import org.eclipse.jetty.http.HttpFields; import org.eclipse.jetty.http.HttpHeader; import org.eclipse.jetty.http.HttpStatus; -import org.eclipse.jetty.http.MetaData; import org.eclipse.jetty.io.Content; import org.eclipse.jetty.util.Callback; @@ -40,8 +39,6 @@ public class SizeLimitHandler extends Handler.Wrapper { private final long _requestLimit; private final long _responseLimit; - private long _read = 0; - private long _written = 0; /** * @param requestLimit The request body size limit in bytes or -1 for no limit @@ -68,76 +65,90 @@ public boolean handle(Request request, Response response, Callback callback) thr } } - HttpFields.Mutable.Wrapper httpFields = new HttpFields.Mutable.Wrapper(response.getHeaders()) + SizeLimitRequestWrapper wrappedRequest = new SizeLimitRequestWrapper(request); + SizeLimitResponseWrapper wrappedResponse = new SizeLimitResponseWrapper(wrappedRequest, response); + return super.handle(wrappedRequest, wrappedResponse, callback); + } + + private class SizeLimitRequestWrapper extends Request.Wrapper + { + private long _read = 0; + + public SizeLimitRequestWrapper(Request wrapped) { - @Override - public HttpField onAddField(HttpField field) + super(wrapped); + } + + @Override + public Content.Chunk read() + { + Content.Chunk chunk = super.read(); + if (chunk == null) + return null; + if (chunk.getFailure() != null) + return chunk; + + // Check request content limit. + ByteBuffer content = chunk.getByteBuffer(); + if (content != null && content.remaining() > 0) { - if (field.getHeader().is(HttpHeader.CONTENT_LENGTH.asString())) + _read += content.remaining(); + if (_requestLimit >= 0 && _read > _requestLimit) { - long contentLength = field.getLongValue(); - if (_responseLimit >= 0 && contentLength > _responseLimit) - throw new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, "Response body is too large: " + contentLength + ">" + _responseLimit); + BadMessageException e = new BadMessageException(HttpStatus.PAYLOAD_TOO_LARGE_413, "Request body is too large: " + _read + ">" + _requestLimit); + getWrapped().fail(e); + return null; } - return super.onAddField(field); } - }; - response = new Response.Wrapper(request, response) - { - @Override - public HttpFields.Mutable getHeaders() - { - return httpFields; - } - }; + return chunk; + } + } - request.addHttpStreamWrapper(httpStream -> new HttpStream.Wrapper(httpStream) - { - @Override - public Content.Chunk read() - { - Content.Chunk chunk = super.read(); - if (chunk == null) - return null; - if (chunk.getFailure() != null) - return chunk; + private class SizeLimitResponseWrapper extends Response.Wrapper + { + private final HttpFields.Mutable _httpFields; + private long _written = 0; - // Check request content limit. - ByteBuffer content = chunk.getByteBuffer(); - if (content != null && content.remaining() > 0) + public SizeLimitResponseWrapper(Request request, Response wrapped) { + super(request, wrapped); + + _httpFields = new HttpFields.Mutable.Wrapper(wrapped.getHeaders()) + { + @Override + public HttpField onAddField(HttpField field) { - _read += content.remaining(); - if (_requestLimit >= 0 && _read > _requestLimit) + if (field.getHeader().is(HttpHeader.CONTENT_LENGTH.asString())) { - BadMessageException e = new BadMessageException(HttpStatus.PAYLOAD_TOO_LARGE_413, "Request body is too large: " + _read + ">" + _requestLimit); - request.fail(e); - return null; + long contentLength = field.getLongValue(); + if (_responseLimit >= 0 && contentLength > _responseLimit) + throw new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, "Response body is too large: " + contentLength + ">" + _responseLimit); } + return super.onAddField(field); } + }; + } - return chunk; - } + @Override + public HttpFields.Mutable getHeaders() { + return _httpFields; + } - @Override - public void send(MetaData.Request request, MetaData.Response response, boolean last, ByteBuffer content, Callback callback) + @Override + public void write(boolean last, ByteBuffer content, Callback callback) + { + if (content != null && content.remaining() > 0) { - // Check response content limit. - if (content != null && content.remaining() > 0) + if (_responseLimit >= 0 && (_written + content.remaining()) > _responseLimit) { - if (_responseLimit >= 0 && (_written + content.remaining()) > _responseLimit) - { - callback.failed(new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, "Response body is too large: " + - _written + content.remaining() + ">" + _responseLimit)); - return; - } - _written += content.remaining(); + callback.failed(new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, "Response body is too large: " + + _written + content.remaining() + ">" + _responseLimit)); + return; } - - super.send(request, response, last, content, callback); + _written += content.remaining(); } - }); - return super.handle(request, response, callback); + super.write(last, content, callback); + } } } diff --git a/jetty-core/jetty-server/src/test/java/org/eclipse/jetty/server/handler/SizeLimitHandlerTest.java b/jetty-core/jetty-server/src/test/java/org/eclipse/jetty/server/handler/SizeLimitHandlerTest.java index e226a949e715..d54af550ed2b 100644 --- a/jetty-core/jetty-server/src/test/java/org/eclipse/jetty/server/handler/SizeLimitHandlerTest.java +++ b/jetty-core/jetty-server/src/test/java/org/eclipse/jetty/server/handler/SizeLimitHandlerTest.java @@ -267,4 +267,29 @@ public boolean handle(Request request, Response response, Callback callback) thr assertThat(response.getContent(), containsString(">8192")); } } + + @Test + public void testMultipleRequests() throws Exception + { + String message = "x".repeat(1024); + _contextHandler.setHandler(new Handler.Abstract() + { + @Override + public boolean handle(Request request, Response response, Callback callback) throws Exception + { + response.write(true, BufferUtil.toBuffer(message), callback); + return true; + } + }); + + _server.start(); + + for (int i = 0; i < 1000; i++) + { + HttpTester.Response response = HttpTester.parseResponse( + _local.getResponse("GET /ctx/hello HTTP/1.0\r\n\r\n")); + assertThat(response.getStatus(), equalTo(200)); + assertThat(response.getContent(), equalTo(message)); + } + } }