router: bound inference request body size and HTTP listener connections

charts/kthena/charts/networking/README.md

@@ -149,6 +149,34 @@ kthenaRouter:
 | `kthenaRouter.terminationGracePeriodSeconds` | int    | `330`   | Pod termination grace period for the router                             |
 | `kthenaRouter.drainTimeout`                  | string | `"5m"`  | Time allowed for the router to drain in-flight requests before shutdown |
 
+### Request and Connection Limits
+
+The router listeners bound request size and connection lifetime so that a single
+client cannot exhaust router memory or connections. Requests below the limits
+behave exactly as before, and no write deadline is applied, so long-running
+streaming inference responses are never truncated.
+
+```yaml
+kthenaRouter:
+  requestLimits:
+    maxRequestBodyBytes: 33554432
+    readHeaderTimeout: 10s
+    idleTimeout: 120s
+    maxHeaderBytes: 1048576
+```
+
+| Parameter                                     | Type   | Default    | Description                                                                                                             |
+| --------------------------------------------- | ------ | ---------- | ------------------------------------------------------------------------------------------------------------------------- |
+| `kthenaRouter.requestLimits.maxRequestBodyBytes` | int    | `33554432` | Largest inference request body accepted, in bytes (32Mi). A larger request is rejected with HTTP 413 before it is buffered. Set to `0` to disable the limit. |
+| `kthenaRouter.requestLimits.readHeaderTimeout`   | string | `"10s"`    | Maximum time a client may take to send the complete request headers                                                     |
+| `kthenaRouter.requestLimits.idleTimeout`         | string | `"120s"`   | Maximum time an idle keep-alive connection is kept open between requests                                                |
+| `kthenaRouter.requestLimits.maxHeaderBytes`      | int    | `1048576`  | Largest request header block accepted, in bytes (1Mi). A larger header block is rejected with HTTP 431. `net/http` allows a few KiB of slack above this value. |
+
+Each value maps to an environment variable on the router container
+(`MAX_REQUEST_BODY_BYTES`, `READ_HEADER_TIMEOUT`, `IDLE_TIMEOUT`,
+`MAX_HEADER_BYTES`). An unparsable or non-positive timeout or header size falls
+back to its default, so a bad value never leaves a listener unbounded.
+
 ## Installation
 
 ### Basic Installation

charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml

@@ -127,6 +127,18 @@ spec:
 
             - name: DRAIN_TIMEOUT
               value: {{ .Values.kthenaRouter.drainTimeout | quote }}
+
+            # Request and connection limits
+            {{- with .Values.kthenaRouter.requestLimits }}
+            - name: MAX_REQUEST_BODY_BYTES
+              value: {{ .maxRequestBodyBytes | quote }}
+            - name: READ_HEADER_TIMEOUT
+              value: {{ .readHeaderTimeout | quote }}
+            - name: IDLE_TIMEOUT
+              value: {{ .idleTimeout | quote }}
+            - name: MAX_HEADER_BYTES
+              value: {{ .maxHeaderBytes | quote }}
+            {{- end }}
           resources: {{- toYaml .Values.kthenaRouter.resource | nindent 12 }}
           livenessProbe:
             httpGet:

charts/kthena/charts/networking/values.yaml

@@ -39,6 +39,23 @@ kthenaRouter:
   # -- Drain timeout for kthena-router graceful shutdown.
   # -- This should be less than terminationGracePeriodSeconds.
   drainTimeout: 5m
+  # requestLimits bounds the request size and connection lifetime accepted by the
+  # router listeners. No write deadline is applied, so streaming responses are
+  # unaffected.
+  requestLimits:
+    # maxRequestBodyBytes is the largest inference request body accepted,
+    # in bytes (default: 33554432, i.e. 32Mi). Larger requests get HTTP 413.
+    # Set to 0 to disable the limit.
+    maxRequestBodyBytes: 33554432
+    # readHeaderTimeout is how long a client may take to send the complete
+    # request headers (default: 10s)
+    readHeaderTimeout: 10s
+    # idleTimeout is how long an idle keep-alive connection is kept open
+    # (default: 120s)
+    idleTimeout: 120s
+    # maxHeaderBytes is the largest request header block accepted, in bytes
+    # (default: 1048576, i.e. 1Mi). Larger headers get HTTP 431.
+    maxHeaderBytes: 1048576
   # fairness configuration for request scheduling
   fairness:
     # enabled controls whether fairness scheduling is active

charts/kthena/values.yaml

@@ -122,6 +122,22 @@ networking:
     # -- Drain timeout for kthena-router graceful shutdown.
     # -- This should be less than terminationGracePeriodSeconds.
     drainTimeout: 5m
+    # requestLimits bounds the request size and connection lifetime accepted by
+    # the router listeners. No write deadline is applied, so streaming inference
+    # responses are unaffected.
+    requestLimits:
+      # -- Largest inference request body accepted, in bytes.<br/>
+      # A larger request is rejected with HTTP 413 before it is buffered.<br/>
+      # Set to `0` to disable the limit.
+      maxRequestBodyBytes: 33554432
+      # -- Maximum time a client may take to send the complete request headers.<br/>
+      # Bounds slow-header clients holding connections open.
+      readHeaderTimeout: 10s
+      # -- Maximum time an idle keep-alive connection is kept open between requests.
+      idleTimeout: 120s
+      # -- Largest request header block accepted, in bytes.<br/>
+      # A larger header block is rejected with HTTP 431.
+      maxHeaderBytes: 1048576
 
 global:
   # -- Certificate Management Mode.<br/>

cmd/kthena-router/app/router.go

@@ -92,6 +92,7 @@ func (s *Server) startRouter(ctx context.Context, router *router.Router, store d
 			readyCheck:       s.HasSynced,
 			activeRequests:   router.ActiveRequestCount,
 			drainTimeout:     s.drainTimeout,
+			limits:           s.limits,
 			startLog:         fmt.Sprintf("Starting default server on port %s", s.Port),
 			shutdownStartLog: "Shutting down default HTTP server ...",
 			shutdownDoneLog:  "Default HTTP server exited",
@@ -205,6 +206,7 @@ type listenerConfig struct {
 	readyCheck     func() bool
 	activeRequests func() int64
 	drainTimeout   time.Duration
+	limits         serverLimits
 	// Gateway mode (non-nil => use gateway branch).
 	gateway          *listenerGatewayConfig
 	startLog         string
@@ -214,6 +216,19 @@ type listenerConfig struct {
 	logListenErr     func(err error)
 }
 
+// newHTTPServer builds a listener HTTP server with the connection-level bounds
+// applied. WriteTimeout is intentionally left unset so that long-running
+// streaming inference responses are not truncated.
+func newHTTPServer(addr string, handler http.Handler, limits serverLimits) *http.Server {
+	return &http.Server{
+		Addr:              addr,
+		Handler:           handler,
+		ReadHeaderTimeout: limits.readHeaderTimeout,
+		IdleTimeout:       limits.idleTimeout,
+		MaxHeaderBytes:    limits.maxHeaderBytes,
+	}
+}
+
 // startListener: build Gin, listen, graceful shutdown on ctx.
 func startListener(ctx context.Context, cfg listenerConfig) *http.Server {
 	engine := gin.New()
@@ -277,10 +292,7 @@ func startListener(ctx context.Context, cfg listenerConfig) *http.Server {
 	} else {
 		klog.Fatal("startListener: invalid listenerConfig (need gateway or defaultRouter+readyCheck)")
 	}
-	srv := &http.Server{
-		Addr:    cfg.addr,
-		Handler: engine.Handler(),
-	}
+	srv := newHTTPServer(cfg.addr, engine.Handler(), cfg.limits)
 
 	go func() {
 		klog.Info(cfg.startLog)
@@ -525,6 +537,7 @@ func (lm *ListenerManager) addListenerToPort(port int32, config ListenerConfig,
 			gateway:          &listenerGatewayConfig{lm: lm, port: port},
 			activeRequests:   lm.router.ActiveRequestCount,
 			drainTimeout:     lm.server.drainTimeout,
+			limits:           lm.server.limits,
 			startLog:         fmt.Sprintf("Starting Gateway listener server on port %d", port),
 			shutdownStartLog: fmt.Sprintf("Shutting down Gateway listener server on port %d ...", port),
 			shutdownDoneLog:  "",

cmd/kthena-router/app/router_listener_test.go

@@ -17,7 +17,13 @@ limitations under the License.
 package app
 
 import (
+	"io"
+	"net"
+	"net/http"
+	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	gatewayv1 "sigs.k8s.io/gateway-api/apis/v1"
@@ -188,3 +194,159 @@ func TestMatchedListenerIsStableAfterGatewayUpdate(t *testing.T) {
 		t.Fatalf("matched listener changed to %q after update", matched.GatewayKey)
 	}
 }
+
+// serveWithLimits starts a listener-equivalent HTTP server on a random local
+// port and returns its address.
+func serveWithLimits(t *testing.T, limits serverLimits, handler http.Handler) string {
+	t.Helper()
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen failed: %v", err)
+	}
+	srv := newHTTPServer(ln.Addr().String(), handler, limits)
+	go func() {
+		_ = srv.Serve(ln)
+	}()
+	t.Cleanup(func() {
+		_ = srv.Close()
+	})
+	return ln.Addr().String()
+}
+
+func TestNewHTTPServerAppliesLimits(t *testing.T) {
+	limits := serverLimits{
+		readHeaderTimeout: 3 * time.Second,
+		idleTimeout:       45 * time.Second,
+		maxHeaderBytes:    8192,
+	}
+
+	srv := newHTTPServer(":8080", http.NotFoundHandler(), limits)
+
+	if srv.ReadHeaderTimeout != limits.readHeaderTimeout {
+		t.Errorf("ReadHeaderTimeout = %v, want %v", srv.ReadHeaderTimeout, limits.readHeaderTimeout)
+	}
+	if srv.IdleTimeout != limits.idleTimeout {
+		t.Errorf("IdleTimeout = %v, want %v", srv.IdleTimeout, limits.idleTimeout)
+	}
+	if srv.MaxHeaderBytes != limits.maxHeaderBytes {
+		t.Errorf("MaxHeaderBytes = %d, want %d", srv.MaxHeaderBytes, limits.maxHeaderBytes)
+	}
+	// A global write deadline would truncate streaming inference responses.
+	if srv.WriteTimeout != 0 {
+		t.Errorf("WriteTimeout = %v, want 0 so that streaming responses are not truncated", srv.WriteTimeout)
+	}
+	// ReadTimeout would bound the upload of a large inference request body;
+	// the body size limit handles that instead.
+	if srv.ReadTimeout != 0 {
+		t.Errorf("ReadTimeout = %v, want 0", srv.ReadTimeout)
+	}
+}
+
+func TestNewHTTPServerClosesSlowHeaderConnections(t *testing.T) {
+	const readHeaderTimeout = 200 * time.Millisecond
+
+	addr := serveWithLimits(t, serverLimits{
+		readHeaderTimeout: readHeaderTimeout,
+		idleTimeout:       time.Minute,
+		maxHeaderBytes:    defaultMaxHeaderBytes,
+	}, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Error("handler must not be reached for an incomplete request")
+	}))
+
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatalf("dial failed: %v", err)
+	}
+	defer conn.Close()
+
+	// Send the request line and one header, then stall without the blank line
+	// that terminates the header block.
+	if _, err := conn.Write([]byte("POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n")); err != nil {
+		t.Fatalf("write failed: %v", err)
+	}
+
+	start := time.Now()
+	if err := conn.SetReadDeadline(time.Now().Add(10 * readHeaderTimeout)); err != nil {
+		t.Fatalf("set read deadline failed: %v", err)
+	}
+	if _, err := conn.Read(make([]byte, 1)); err == nil {
+		t.Fatal("expected the server to close the connection, got a response")
+	} else if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+		t.Fatalf("connection still open after %v, want it closed after %v", time.Since(start), readHeaderTimeout)
+	}
+	if elapsed := time.Since(start); elapsed < readHeaderTimeout {
+		t.Errorf("connection closed after %v, want at least %v", elapsed, readHeaderTimeout)
+	}
+}
+
+func TestNewHTTPServerRejectsOversizedHeaders(t *testing.T) {
+	const maxHeaderBytes = 1024
+
+	var handlerCalled atomic.Bool
+	addr := serveWithLimits(t, serverLimits{
+		readHeaderTimeout: 10 * time.Second,
+		idleTimeout:       time.Minute,
+		maxHeaderBytes:    maxHeaderBytes,
+	}, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		handlerCalled.Store(true)
+	}))
+
+	req, err := http.NewRequest(http.MethodGet, "http://"+addr+"/healthz", nil)
+	if err != nil {
+		t.Fatalf("new request failed: %v", err)
+	}
+	// net/http allows a few KiB of slack above MaxHeaderBytes, so overshoot it
+	// by a wide margin.
+	req.Header.Set("X-Oversized", strings.Repeat("a", 16*maxHeaderBytes))
+
+	// The server may reply 431 or close the connection while the client is
+	// still writing; either way the request must not reach the handler.
+	resp, err := http.DefaultClient.Do(req)
+	if err == nil {
+		defer resp.Body.Close()
+		if resp.StatusCode != http.StatusRequestHeaderFieldsTooLarge {
+			t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusRequestHeaderFieldsTooLarge)
+		}
+	}
+	if handlerCalled.Load() {
+		t.Error("handler was invoked for a request with oversized headers")
+	}
+}
+
+func TestNewHTTPServerAllowsLongStreamingResponses(t *testing.T) {
+	const (
+		chunks        = 5
+		chunkInterval = 100 * time.Millisecond
+	)
+
+	// Both bounds are shorter than the response duration; neither may truncate it.
+	addr := serveWithLimits(t, serverLimits{
+		readHeaderTimeout: 100 * time.Millisecond,
+		idleTimeout:       100 * time.Millisecond,
+		maxHeaderBytes:    defaultMaxHeaderBytes,
+	}, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		for i := 0; i < chunks; i++ {
+			if _, err := io.WriteString(w, "data: chunk\n\n"); err != nil {
+				return
+			}
+			w.(http.Flusher).Flush()
+			time.Sleep(chunkInterval)
+		}
+	}))
+
+	resp, err := http.Get("http://" + addr + "/v1/chat/completions")
+	if err != nil {
+		t.Fatalf("request failed: %v", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("reading the streamed body failed: %v", err)
+	}
+	if got := strings.Count(string(body), "data: chunk"); got != chunks {
+		t.Errorf("received %d chunks, want %d", got, chunks)
+	}
+}