update SizeLimitHandler with fixes needed for appengine tests (#11570)

* update SizeLimitHandler with fixes needed for appengine tests --------- Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
jetty · Apr 24, 2024 · 0c14484 · 0c14484
1 parent c10ec98
commit 0c14484
Showing 1 changed file with 11 additions and 3 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
@@ -109,6 +109,7 @@ private class SizeLimitResponseWrapper extends Response.Wrapper
     {
         private final HttpFields.Mutable _httpFields;
         private long _written = 0;
+        private HttpException.RuntimeException _failure;
 
         public SizeLimitResponseWrapper(Request request, Response wrapped)
         {
@@ -119,7 +120,7 @@ public SizeLimitResponseWrapper(Request request, Response wrapped)
                 @Override
                 public HttpField onAddField(HttpField field)
                 {
-                    if (field.getHeader().is(HttpHeader.CONTENT_LENGTH.asString()))
+                    if (field.getHeader() == HttpHeader.CONTENT_LENGTH)
                     {
                         long contentLength = field.getLongValue();
                         if (_responseLimit >= 0 && contentLength > _responseLimit)
@@ -139,12 +140,19 @@ public HttpFields.Mutable getHeaders()
         @Override
         public void write(boolean last, ByteBuffer content, Callback callback)
         {
+            if (_failure != null)
+            {
+                callback.failed(_failure);
+                return;
+            }
+
             if (content != null && content.remaining() > 0)
             {
                 if (_responseLimit >= 0 && (_written + content.remaining())  > _responseLimit)
                 {
-                    String message = "Response body is too large: %d>%d".formatted(_written + content.remaining(), _responseLimit);
-                    callback.failed(new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500, message));
+                    _failure = new HttpException.RuntimeException(HttpStatus.INTERNAL_SERVER_ERROR_500,
+                        "Response body is too large: %d>%d".formatted(_written + content.remaining(), _responseLimit));
+                    callback.failed(_failure);
                     return;
                 }
                 _written += content.remaining();