Early Rejection of Large Requests Based on Content-Length Header

core/src/main/java/com/linecorp/armeria/common/ContentTooLargeException.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -50,6 +50,7 @@ public static ContentTooLargeExceptionBuilder builder() {
     private final long maxContentLength;
     private final long contentLength;
     private final long transferred;
+    private final boolean earlyRejection;
 
     private ContentTooLargeException() {
         this(false);
@@ -62,16 +63,18 @@ private ContentTooLargeException(boolean neverSample) {
         maxContentLength = -1;
         transferred = -1;
         contentLength = -1;
+        earlyRejection = false;
     }
 
     ContentTooLargeException(long maxContentLength, long contentLength, long transferred,
-                             @Nullable Throwable cause) {
-        super(toString(maxContentLength, contentLength, transferred), cause);
+                             boolean earlyRejection, @Nullable Throwable cause) {
+        super(toString(maxContentLength, contentLength, transferred, earlyRejection), cause);
 
         neverSample = false;
         this.transferred = transferred;
         this.contentLength = contentLength;
         this.maxContentLength = maxContentLength;
+        this.earlyRejection = earlyRejection;
     }
 
     /**
@@ -96,6 +99,13 @@ public long maxContentLength() {
         return maxContentLength;
     }
 
+    /**
+     * Returns whether the exception is raised when reading content-length header.
+     */
+    public boolean earlyRejection() {
+        return earlyRejection;
+    }
+
     @Override
     public Throwable fillInStackTrace() {
         if (!neverSample && Flags.verboseExceptionSampler().isSampled(getClass())) {
@@ -105,7 +115,8 @@ public Throwable fillInStackTrace() {
     }
 
     @Nullable
-    private static String toString(long maxContentLength, long contentLength, long transferred) {
+    private static String toString(long maxContentLength, long contentLength, long transferred,
+                                   boolean earlyRejection) {
         try (TemporaryThreadLocals ttl = TemporaryThreadLocals.acquire()) {
             final StringBuilder buf = ttl.stringBuilder();
             if (maxContentLength >= 0) {
@@ -117,6 +128,9 @@ private static String toString(long maxContentLength, long contentLength, long t
             if (transferred >= 0) {
                 buf.append(", transferred: ").append(transferred);
             }
+            if (earlyRejection) {
+                buf.append(", earlyRejection: ").append("true");
+            }
             return buf.length() != 0 ? buf.substring(2) : null;
         }
     }

core/src/main/java/com/linecorp/armeria/common/ContentTooLargeExceptionBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -30,6 +30,8 @@ public final class ContentTooLargeExceptionBuilder {
     private long maxContentLength = -1;
     private long contentLength = -1;
     private long transferred = -1;
+    private boolean earlyRejection;
+
     @Nullable
     private Throwable cause;
 
@@ -83,13 +85,24 @@ public ContentTooLargeExceptionBuilder cause(Throwable cause) {
         return this;
     }
 
+    /**
+     * Sets the exception as early rejection.
+     */
+    @UnstableApi
+    public ContentTooLargeExceptionBuilder earlyRejection(boolean isEarlyRejection) {
+        this.earlyRejection = isEarlyRejection;
+        return this;
+    }
+
     /**
      * Returns a new instance of {@link ContentTooLargeException}.
      */
     public ContentTooLargeException build() {
-        if (maxContentLength < 0 && contentLength < 0 && transferred < 0 && cause == null) {
+        if (maxContentLength < 0 && contentLength < 0 &&
+                transferred < 0 && !earlyRejection && cause == null) {
             return ContentTooLargeException.get();
         }
-        return new ContentTooLargeException(maxContentLength, contentLength, transferred, cause);
+        return new ContentTooLargeException(maxContentLength, contentLength,
+                transferred, earlyRejection, cause);
     }
 }

core/src/main/java/com/linecorp/armeria/server/DecodedHttpRequest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -143,6 +143,16 @@ default CompletableFuture<Void> whenAggregated() {
      */
     long requestStartTimeMicros();
 
+    /**
+     * Returns the maximum allowed length of the content of the request.
+     */
+    long maxRequestLength();
+
+    /**
+     * Returns the transferred bytes of the request.
+     */
+    long transferredBytes();
+
     /**
      * Returns whether the request is an HTTP/1.1 webSocket request.
      */

core/src/main/java/com/linecorp/armeria/server/DecodedHttpRequestWriter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -19,10 +19,5 @@
 import com.linecorp.armeria.common.HttpRequestWriter;
 
 interface DecodedHttpRequestWriter extends DecodedHttpRequest, HttpRequestWriter {
-
-    long maxRequestLength();
-
-    long transferredBytes();
-
     void increaseTransferredBytes(long delta);
 }

core/src/main/java/com/linecorp/armeria/server/EmptyContentDecodedHttpRequest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -260,4 +260,14 @@ public long requestStartTimeNanos() {
     public long requestStartTimeMicros() {
         return requestStartTimeMicros;
     }
+
+    @Override
+    public long maxRequestLength() {
+        return 0;
+    }
+
+    @Override
+    public long transferredBytes() {
+        return 0;
+    }
 }

core/src/main/java/com/linecorp/armeria/server/Http1RequestDecoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -208,8 +208,8 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
                     // Validate the 'content-length' header.
                     final String contentLengthStr = headers.get(HttpHeaderNames.CONTENT_LENGTH);
                     final boolean contentEmpty;
+                    long contentLength = 0;
                     if (contentLengthStr != null) {
-                        long contentLength;
                         try {
                             contentLength = Long.parseLong(contentLengthStr);
                         } catch (NumberFormatException ignored) {
@@ -280,6 +280,10 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
                     final boolean endOfStream = contentEmpty && !transferEncodingChunked;
                     this.req = req = DecodedHttpRequest.of(endOfStream, eventLoop, id, 1, headers,
                                                            keepAlive, inboundTrafficController, routingCtx);
+                    final long maxRequestLength = req.maxRequestLength();
+                    if (maxRequestLength > 0 && contentLength > maxRequestLength) {
+                        abortLargeRequest(ctx, req, id, endOfStream, keepAliveHandler, true);
+                    }
                     cfg.serverMetrics().increasePendingHttp1Requests();
                     ctx.fireChannelRead(req);
                 } else {
@@ -313,33 +317,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
                     final long maxContentLength = decodedReq.maxRequestLength();
                     final long transferredLength = decodedReq.transferredBytes();
                     if (maxContentLength > 0 && transferredLength > maxContentLength) {
-                        final ContentTooLargeException cause =
-                                ContentTooLargeException.builder()
-                                                        .maxContentLength(maxContentLength)
-                                                        .contentLength(req.headers())
-                                                        .transferred(transferredLength)
-                                                        .build();
-                        discarding = true;
-                        req = null;
-                        final boolean shouldReset;
-                        if (encoder instanceof ServerHttp1ObjectEncoder) {
-                            if (encoder.isResponseHeadersSent(id, 1)) {
-                                ctx.channel().close();
-                            } else {
-                                keepAliveHandler.disconnectWhenFinished();
-                            }
-                            shouldReset = false;
-                        } else {
-                            // Upgraded to HTTP/2. Reset only if the remote peer is still open.
-                            shouldReset = !endOfStream;
-                        }
-
-                        // Wrap the cause with the returned status to let LoggingService correctly log the
-                        // status.
-                        final HttpStatusException httpStatusException =
-                                HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
-                        decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
-                        decodedReq.abortResponse(httpStatusException, true);
+                        abortLargeRequest(ctx, decodedReq, id, endOfStream, keepAliveHandler, false);
                         return;
                     }
 
@@ -377,6 +355,38 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
         }
     }
 
+    private void abortLargeRequest(ChannelHandlerContext ctx, DecodedHttpRequest decodedReq, int id,
+                                   boolean endOfStream, KeepAliveHandler keepAliveHandler, boolean isEarlyRejection) {
+        final ContentTooLargeException cause =
+                ContentTooLargeException.builder()
+                        .maxContentLength(decodedReq.maxRequestLength())
+                        .transferred(decodedReq.transferredBytes())
+                        .contentLength(decodedReq.headers())
+                        .earlyRejection(isEarlyRejection)
+                        .build();
+        discarding = true;
+        req = null;
+        final boolean shouldReset;
+        if (encoder instanceof ServerHttp1ObjectEncoder) {
+            if (encoder.isResponseHeadersSent(id, 1)) {
+                ctx.channel().close();
+            } else {
+                keepAliveHandler.disconnectWhenFinished();
+            }
+            shouldReset = false;
+        } else {
+            // Upgraded to HTTP/2. Reset only if the remote peer is still open.
+            shouldReset = !endOfStream;
+        }
+
+        // Wrap the cause with the returned status to let LoggingService correctly log the
+        // status.
+        final HttpStatusException httpStatusException =
+                HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
+        decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
+        decodedReq.abortResponse(httpStatusException, true);
+    }
+
     private void removeFromPipelineIfUpgraded(ChannelHandlerContext ctx, boolean endOfStream) {
         if (endOfStream && encoder instanceof ServerHttp2ObjectEncoder) {
             // An HTTP/1 connection has been upgraded to HTTP/2.

core/src/main/java/com/linecorp/armeria/server/Http2RequestDecoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 LINE Corporation
+ * Copyright 2024 LINE Corporation
  *
  * LINE Corporation licenses this file to you under the Apache License,
  * version 2.0 (the "License"); you may not use this file except in compliance
@@ -165,8 +165,8 @@ public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers
 
             // Validate the 'content-length' header if exists.
             final String contentLengthStr = headers.get(HttpHeaderNames.CONTENT_LENGTH);
+            long contentLength = 0;
             if (contentLengthStr != null) {
-                long contentLength;
                 try {
                     contentLength = Long.parseLong(contentLengthStr);
                 } catch (NumberFormatException ignored) {
@@ -206,6 +206,10 @@ public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers
             final EventLoop eventLoop = ctx.channel().eventLoop();
             req = DecodedHttpRequest.of(endOfStream, eventLoop, id, streamId, headers, true,
                                         inboundTrafficController, routingCtx);
+            final long maxRequestLength = req.maxRequestLength();
+            if (maxRequestLength > 0 && contentLength > maxRequestLength) {
+                abortLargeRequest(req, endOfStream, true);
+            }
             requests.put(streamId, req);
             cfg.serverMetrics().increasePendingHttp2Requests();
             ctx.fireChannelRead(req);
@@ -321,20 +325,7 @@ public int onDataRead(
         final long maxContentLength = decodedReq.maxRequestLength();
         final long transferredLength = decodedReq.transferredBytes();
         if (maxContentLength > 0 && transferredLength > maxContentLength) {
-            assert encoder != null;
-            final ContentTooLargeException cause =
-                    ContentTooLargeException.builder()
-                                            .maxContentLength(maxContentLength)
-                                            .contentLength(decodedReq.headers())
-                                            .transferred(transferredLength)
-                                            .build();
-
-            final boolean shouldReset = !endOfStream;
-
-            final HttpStatusException httpStatusException =
-                    HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
-            decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
-            decodedReq.abortResponse(httpStatusException, true);
+            abortLargeRequest(decodedReq, endOfStream, false);
         } else if (decodedReq.isOpen()) {
             try {
                 // The decodedReq will be automatically closed if endOfStream is true.
@@ -350,6 +341,25 @@ public int onDataRead(
         return dataLength + padding;
     }
 
+    private void abortLargeRequest(DecodedHttpRequest decodedReq, boolean endOfStream,
+                                   boolean isEarlyRejection) {
+        assert encoder != null;
+        final ContentTooLargeException cause =
+                ContentTooLargeException.builder()
+                        .maxContentLength(decodedReq.maxRequestLength())
+                        .contentLength(decodedReq.headers())
+                        .transferred(decodedReq.transferredBytes())
+                        .earlyRejection(isEarlyRejection)
+                        .build();
+
+        final boolean shouldReset = !endOfStream;
+
+        final HttpStatusException httpStatusException =
+                HttpStatusException.of(HttpStatus.REQUEST_ENTITY_TOO_LARGE, cause);
+        decodedReq.setShouldResetOnlyIfRemoteIsOpen(shouldReset);
+        decodedReq.abortResponse(httpStatusException, true);
+    }
+
     private void writeInvalidRequestPathResponse(int streamId, @Nullable RequestHeaders headers) {
         writeErrorResponse(streamId, headers, HttpStatus.BAD_REQUEST,
                            "Invalid request path", null);