From a94809b6aab7e7da93cdff359dfce0743177447f Mon Sep 17 00:00:00 2001
From: Bjørn Christian Seime <bjorncs@yahooinc.com>
Date: Mon, 26 Jun 2023 14:37:48 +0200
Subject: Add configurable limit for max request content size

---
 .../http/server/jetty/ServletRequestReader.java    | 37 ++++++++++++++++++++--
 .../jdisc.http.jdisc.http.connector.def            |  3 ++
 2 files changed, 38 insertions(+), 2 deletions(-)

(limited to 'container-core/src/main')

diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletRequestReader.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletRequestReader.java
index 2f2c48e0b48..75ef655c60c 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletRequestReader.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/ServletRequestReader.java
@@ -1,16 +1,19 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jdisc.http.server.jetty;
 
+import com.yahoo.jdisc.Response;
 import com.yahoo.jdisc.handler.CompletionHandler;
 import com.yahoo.jdisc.handler.ContentChannel;
 import jakarta.servlet.ReadListener;
 import jakarta.servlet.ServletInputStream;
 import jakarta.servlet.http.HttpServletRequest;
+import org.eclipse.jetty.server.Request;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -33,6 +36,7 @@ import java.util.logging.Logger;
  */
 class ServletRequestReader {
 
+
     private enum State {
         NOT_STARTED, READING, ALL_DATA_READ, REQUEST_CONTENT_CLOSED
     }
@@ -96,12 +100,15 @@ class ServletRequestReader {
     private final CompletableFuture<Void> finishedFuture = new CompletableFuture<>();
 
     ServletRequestReader(
-            HttpServletRequest req,
+            Request req,
             ContentChannel requestContentChannel,
             Janitor janitor,
             RequestMetricReporter metricReporter) {
         this.req = Objects.requireNonNull(req);
-        this.requestContentChannel = Objects.requireNonNull(requestContentChannel);
+        long maxContentSize = RequestUtils.getConnector(req).connectorConfig().maxContentSize();
+        this.requestContentChannel = maxContentSize >= 0
+                ? new ByteLimitedContentChannel(Objects.requireNonNull(requestContentChannel), maxContentSize)
+                : Objects.requireNonNull(requestContentChannel);
         this.janitor = Objects.requireNonNull(janitor);
         this.metricReporter = Objects.requireNonNull(metricReporter);
     }
@@ -259,4 +266,30 @@ class ServletRequestReader {
         }
     }
 
+    private static class ByteLimitedContentChannel implements ContentChannel {
+        private final long maxContentSize;
+        private final AtomicLong bytesWritten = new AtomicLong();
+        private final ContentChannel delegate;
+
+        ByteLimitedContentChannel(ContentChannel delegate, long maxContentSize) {
+            this.delegate = delegate;
+            this.maxContentSize = maxContentSize;
+        }
+
+        @Override
+        public void write(ByteBuffer buf, CompletionHandler handler) {
+            long written = bytesWritten.addAndGet(buf.remaining());
+            if (written > maxContentSize) {
+                handler.failed(new RequestException(
+                        Response.Status.REQUEST_TOO_LONG,
+                        "Request content length %d exceeds limit of %d bytes".formatted(written, maxContentSize)));
+                return;
+            }
+            delegate.write(buf, handler);
+        }
+
+        @Override public void close(CompletionHandler h) { delegate.close(h); }
+        @Override public void onError(Throwable t) { delegate.onError(t); }
+    }
+
 }
diff --git a/container-core/src/main/resources/configdefinitions/jdisc.http.jdisc.http.connector.def b/container-core/src/main/resources/configdefinitions/jdisc.http.jdisc.http.connector.def
index bdcc3f9e40a..3c01012fd9e 100644
--- a/container-core/src/main/resources/configdefinitions/jdisc.http.jdisc.http.connector.def
+++ b/container-core/src/main/resources/configdefinitions/jdisc.http.jdisc.http.connector.def
@@ -22,6 +22,9 @@ responseHeaderSize                  int      default=65536
 # The accept queue size (also known as accept backlog).
 acceptQueueSize                     int      default=0
 
+# Max content size allowed for requests. Set to -1 to disable.
+maxContentSize                      long     default=-1
+
 # Whether the server socket reuses addresses.
 reuseAddress                        bool     default=true
 
-- 
cgit v1.2.3