From 68db7d51704ecff983458c5e2df11d9cb6c91783 Mon Sep 17 00:00:00 2001 From: Gite vivek kumar <71180467+vivek-gite@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:13:54 +0530 Subject: [PATCH] feat(router): add request and connection limits for kthenaRouter This update introduces configurable request limits for the kthenaRouter, including maximum request body size, read header timeout, idle timeout, and maximum header bytes. These limits help prevent excessive resource usage by clients and ensure more stable server performance. The new parameters are documented in the Helm chart values and integrated into the router's deployment configuration. - Added requestLimits section in values.yaml - Updated README.md to include new parameters - Implemented limits in the router's HTTP server configuration - Added tests to validate request body size enforcement and header limits Signed-off-by: [Gite Vivek Kumar] [vivekkumargite@outlook.com] Signed-off-by: Gite Vivek Kumar --- charts/kthena/charts/networking/README.md | 28 +++ .../kthena-router/component/deployment.yaml | 12 ++ charts/kthena/charts/networking/values.yaml | 17 ++ charts/kthena/values.yaml | 16 ++ cmd/kthena-router/app/router.go | 21 ++- cmd/kthena-router/app/router_listener_test.go | 162 ++++++++++++++++++ cmd/kthena-router/app/server.go | 58 ++++++- cmd/kthena-router/app/server_test.go | 65 +++++++ .../docs/reference/helm-chart-values.md | 4 + pkg/kthena-router/router/router.go | 38 ++++ pkg/kthena-router/router/router_test.go | 123 +++++++++++++ 11 files changed, 537 insertions(+), 7 deletions(-) diff --git a/charts/kthena/charts/networking/README.md b/charts/kthena/charts/networking/README.md index 30aa3cd63f..40169fd2a2 100644 --- a/charts/kthena/charts/networking/README.md +++ b/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 diff --git a/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml b/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml index d5d87a972f..c14591b49c 100644 --- a/charts/kthena/charts/networking/templates/kthena-router/component/deployment.yaml +++ b/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: diff --git a/charts/kthena/charts/networking/values.yaml b/charts/kthena/charts/networking/values.yaml index 42779cf9cc..67c12eea69 100644 --- a/charts/kthena/charts/networking/values.yaml +++ b/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 diff --git a/charts/kthena/values.yaml b/charts/kthena/values.yaml index a4fa23f1ea..8bb73923e6 100644 --- a/charts/kthena/values.yaml +++ b/charts/kthena/values.yaml @@ -118,6 +118,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.
+ # A larger request is rejected with HTTP 413 before it is buffered.
+ # Set to `0` to disable the limit. + maxRequestBodyBytes: 33554432 + # -- Maximum time a client may take to send the complete request headers.
+ # 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.
+ # A larger header block is rejected with HTTP 431. + maxHeaderBytes: 1048576 global: # -- Certificate Management Mode.
diff --git a/cmd/kthena-router/app/router.go b/cmd/kthena-router/app/router.go index 143050d0f3..57f7bf5b0c 100644 --- a/cmd/kthena-router/app/router.go +++ b/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: "", diff --git a/cmd/kthena-router/app/router_listener_test.go b/cmd/kthena-router/app/router_listener_test.go index 926951623d..9f7fa9b6e1 100644 --- a/cmd/kthena-router/app/router_listener_test.go +++ b/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) + } +} diff --git a/cmd/kthena-router/app/server.go b/cmd/kthena-router/app/server.go index e44931b8d7..b5fccf1a3d 100644 --- a/cmd/kthena-router/app/server.go +++ b/cmd/kthena-router/app/server.go @@ -18,7 +18,9 @@ package app import ( "context" + "net/http" "os" + "strconv" "time" "k8s.io/client-go/tools/cache" @@ -30,6 +32,28 @@ import ( const defaultDrainTimeout = 5 * time.Minute +const ( + // defaultReadHeaderTimeout bounds how long a client may take to send the + // complete request headers, so a slow-header client cannot hold a + // connection and its goroutine indefinitely. + defaultReadHeaderTimeout = 10 * time.Second + // defaultIdleTimeout bounds how long an idle keep-alive connection is kept + // open between requests. + defaultIdleTimeout = 120 * time.Second + // defaultMaxHeaderBytes matches net/http's own default, so the limit only + // becomes stricter when an operator asks for it. + defaultMaxHeaderBytes = http.DefaultMaxHeaderBytes +) + +// serverLimits are the connection-level bounds applied to every router HTTP +// listener. WriteTimeout is deliberately absent: streaming inference responses +// are long-lived and a global write deadline would truncate them. +type serverLimits struct { + readHeaderTimeout time.Duration + idleTimeout time.Duration + maxHeaderBytes int +} + type Server struct { store datastore.Store controllers Controller @@ -45,6 +69,8 @@ type Server struct { KubeAPIBurst int // drainTimeout is HTTP server shutdown grace; not datastore state. drainTimeout time.Duration + // limits are the connection-level bounds shared by every HTTP listener. + limits serverLimits } func NewServer(port string, enableTLS bool, cert, key string, enableGatewayAPI bool, enableGatewayAPIInferenceExtension bool, debugPort int, kubeAPIQPS float32, kubeAPIBurst int) *Server { @@ -60,17 +86,43 @@ func NewServer(port string, enableTLS bool, cert, key string, enableGatewayAPI b KubeAPIQPS: kubeAPIQPS, KubeAPIBurst: kubeAPIBurst, drainTimeout: parseDrainTimeout(), + limits: parseServerLimits(), } } func parseDrainTimeout() time.Duration { - if v := os.Getenv("DRAIN_TIMEOUT"); v != "" { + return parsePositiveDurationEnv("DRAIN_TIMEOUT", defaultDrainTimeout) +} + +// parseServerLimits reads the listener bounds from READ_HEADER_TIMEOUT, +// IDLE_TIMEOUT and MAX_HEADER_BYTES. Invalid or non-positive values fall back +// to the defaults, so a bad value can never leave a listener unbounded. +func parseServerLimits() serverLimits { + return serverLimits{ + readHeaderTimeout: parsePositiveDurationEnv("READ_HEADER_TIMEOUT", defaultReadHeaderTimeout), + idleTimeout: parsePositiveDurationEnv("IDLE_TIMEOUT", defaultIdleTimeout), + maxHeaderBytes: parseMaxHeaderBytes(), + } +} + +func parsePositiveDurationEnv(key string, fallback time.Duration) time.Duration { + if v := os.Getenv(key); v != "" { if d, err := time.ParseDuration(v); err == nil && d > 0 { return d } - klog.Warningf("Invalid DRAIN_TIMEOUT %q, using default %v", v, defaultDrainTimeout) + klog.Warningf("Invalid %s %q, using default %v", key, v, fallback) + } + return fallback +} + +func parseMaxHeaderBytes() int { + if v := os.Getenv("MAX_HEADER_BYTES"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + klog.Warningf("Invalid MAX_HEADER_BYTES %q, using default %v", v, defaultMaxHeaderBytes) } - return defaultDrainTimeout + return defaultMaxHeaderBytes } func (s *Server) Run(ctx context.Context) { diff --git a/cmd/kthena-router/app/server_test.go b/cmd/kthena-router/app/server_test.go index 54a9865318..49e6f734cb 100644 --- a/cmd/kthena-router/app/server_test.go +++ b/cmd/kthena-router/app/server_test.go @@ -18,6 +18,7 @@ package app import ( "testing" + "time" "github.com/stretchr/testify/assert" ) @@ -40,3 +41,67 @@ func TestNewServerDebugPortDefault(t *testing.T) { }) } } + +// TestParseServerLimits tests that listener bounds are read from the +// environment and that invalid values never leave a listener unbounded. +func TestParseServerLimits(t *testing.T) { + testCases := []struct { + name string + readHeaderTimeout string + idleTimeout string + maxHeaderBytes string + want serverLimits + }{ + { + name: "defaults when unset", + want: serverLimits{ + readHeaderTimeout: defaultReadHeaderTimeout, + idleTimeout: defaultIdleTimeout, + maxHeaderBytes: defaultMaxHeaderBytes, + }, + }, + { + name: "overrides are honoured", + readHeaderTimeout: "3s", + idleTimeout: "45s", + maxHeaderBytes: "8192", + want: serverLimits{ + readHeaderTimeout: 3 * time.Second, + idleTimeout: 45 * time.Second, + maxHeaderBytes: 8192, + }, + }, + { + name: "unparsable values fall back to defaults", + readHeaderTimeout: "ten-seconds", + idleTimeout: "120", + maxHeaderBytes: "1MiB", + want: serverLimits{ + readHeaderTimeout: defaultReadHeaderTimeout, + idleTimeout: defaultIdleTimeout, + maxHeaderBytes: defaultMaxHeaderBytes, + }, + }, + { + name: "non-positive values fall back to defaults", + readHeaderTimeout: "0s", + idleTimeout: "-1s", + maxHeaderBytes: "0", + want: serverLimits{ + readHeaderTimeout: defaultReadHeaderTimeout, + idleTimeout: defaultIdleTimeout, + maxHeaderBytes: defaultMaxHeaderBytes, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("READ_HEADER_TIMEOUT", tc.readHeaderTimeout) + t.Setenv("IDLE_TIMEOUT", tc.idleTimeout) + t.Setenv("MAX_HEADER_BYTES", tc.maxHeaderBytes) + + assert.Equal(t, tc.want, parseServerLimits()) + }) + } +} diff --git a/docs/kthena/docs/reference/helm-chart-values.md b/docs/kthena/docs/reference/helm-chart-values.md index 482de115b7..5b758a04e5 100644 --- a/docs/kthena/docs/reference/helm-chart-values.md +++ b/docs/kthena/docs/reference/helm-chart-values.md @@ -32,6 +32,10 @@ A Helm chart for deploying Kthena | networking.kthenaRouter.image.repository | string | `"ghcr.io/volcano-sh/kthena-router"` | Image repository for Kthena Router. | | networking.kthenaRouter.image.tag | string | `"latest"` | Image tag for Kthena Router. | | networking.kthenaRouter.port | int | `8080` | Container port for Kthena Router. | +| networking.kthenaRouter.requestLimits.idleTimeout | string | `"120s"` | Maximum time an idle keep-alive connection is kept open between requests. | +| networking.kthenaRouter.requestLimits.maxHeaderBytes | int | `1048576` | Largest request header block accepted, in bytes.
A larger header block is rejected with HTTP 431. | +| networking.kthenaRouter.requestLimits.maxRequestBodyBytes | int | `33554432` | Largest inference request body accepted, in bytes.
A larger request is rejected with HTTP 413 before it is buffered.
Set to `0` to disable the limit. | +| networking.kthenaRouter.requestLimits.readHeaderTimeout | string | `"10s"` | Maximum time a client may take to send the complete request headers.
Bounds slow-header clients holding connections open. | | networking.kthenaRouter.sessionBoost.enabled | bool | `false` | Enable session-boost scheduling. Mutually exclusive with fairness. | | networking.kthenaRouter.sessionBoost.gracePeriod | string | `"0s"` | Wait time after a request completes for a same-session follow-up.
Disabled by default (`0s`). | | networking.kthenaRouter.sessionBoost.header | string | `"X-Session-ID"` | HTTP header used to identify conversation sessions. | diff --git a/pkg/kthena-router/router/router.go b/pkg/kthena-router/router/router.go index 4e0375b3ad..d90a846ab7 100644 --- a/pkg/kthena-router/router/router.go +++ b/pkg/kthena-router/router/router.go @@ -91,6 +91,32 @@ func getEnvBool(key string, fallback bool) bool { var EnableFairnessScheduling = getEnvBool("ENABLE_FAIRNESS_SCHEDULING", false) var EnableSessionBoost = getEnvBool("ENABLE_SESSION_BOOST", false) +// defaultMaxRequestBodyBytes bounds the inference request body the router +// buffers. It is generous enough for multimodal and tool-call payloads while +// keeping a single client from forcing an unbounded allocation. +const defaultMaxRequestBodyBytes = 32 << 20 // 32 MiB + +// MaxRequestBodyBytes is the largest inference request body the router accepts, +// configured by MAX_REQUEST_BODY_BYTES. A larger request is rejected with +// HTTP 413 before the payload is fully buffered. A non-positive value disables +// the limit. +var MaxRequestBodyBytes = parseMaxRequestBodyBytes() + +// parseMaxRequestBodyBytes reads the request body limit in bytes from the +// MAX_REQUEST_BODY_BYTES environment variable. Setting it to a non-positive +// value (e.g. "0") explicitly disables the limit. An invalid value falls back +// to defaultMaxRequestBodyBytes. +func parseMaxRequestBodyBytes() int64 { + if s, ok := os.LookupEnv("MAX_REQUEST_BODY_BYTES"); ok { + if n, err := strconv.ParseInt(s, 10, 64); err == nil { + // A non-positive value explicitly disables the limit. + return n + } + klog.Warningf("Invalid MAX_REQUEST_BODY_BYTES %q, using default %v", s, defaultMaxRequestBodyBytes) + } + return defaultMaxRequestBodyBytes +} + type Router struct { scheduler scheduler.Scheduler authenticator *auth.JWTAuthenticator @@ -585,8 +611,20 @@ func (r *Router) doLoadbalance(c *gin.Context, modelRequest ModelRequest) error } func ParseModelRequest(c *gin.Context) (ModelRequest, error) { + // Cap the body before reading it so an oversized payload is rejected + // instead of being buffered in full. + if MaxRequestBodyBytes > 0 { + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, MaxRequestBodyBytes) + } bodyBytes, err := io.ReadAll(c.Request.Body) if err != nil { + var maxBytesErr *http.MaxBytesError + if errors.As(err, &maxBytesErr) { + // Report the limit only; the client does not need router internals. + msg := fmt.Sprintf("request body exceeds the %d byte limit", maxBytesErr.Limit) + c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, msg) + return nil, errors.New(msg) + } c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) return nil, err } diff --git a/pkg/kthena-router/router/router_test.go b/pkg/kthena-router/router/router_test.go index 2e3c889d8b..4f7c659261 100644 --- a/pkg/kthena-router/router/router_test.go +++ b/pkg/kthena-router/router/router_test.go @@ -1327,6 +1327,129 @@ func TestParseModelRequestValidatesModelName(t *testing.T) { } } +// countingReader counts how many bytes were actually pulled from the request +// body, so a test can prove an oversized payload was never fully buffered. +type countingReader struct { + r io.Reader + n int64 +} + +func (c *countingReader) Read(p []byte) (int, error) { + n, err := c.r.Read(p) + c.n += int64(n) + return n, err +} + +// modelRequestBody returns a valid inference request body of exactly size +// bytes, padding the prompt to reach the requested length. +func modelRequestBody(t *testing.T, size int) string { + t.Helper() + const prefix = `{"model":"test-model","prompt":"` + const suffix = `"}` + if size < len(prefix)+len(suffix) { + t.Fatalf("size %d is too small for a valid request body", size) + } + return prefix + strings.Repeat("a", size-len(prefix)-len(suffix)) + suffix +} + +func TestParseModelRequestEnforcesMaxRequestBodyBytes(t *testing.T) { + const limit int64 = 512 + + tests := []struct { + name string + limit int64 + bodySize int + wantStatus int + }{ + { + name: "below limit", + limit: limit, + bodySize: int(limit) - 1, + }, + { + name: "exactly at limit", + limit: limit, + bodySize: int(limit), + }, + { + name: "one byte above limit", + limit: limit, + bodySize: int(limit) + 1, + wantStatus: http.StatusRequestEntityTooLarge, + }, + { + name: "far above limit", + limit: limit, + bodySize: int(limit) * 8, + wantStatus: http.StatusRequestEntityTooLarge, + }, + { + name: "limit disabled", + limit: 0, + bodySize: int(limit) * 8, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + prevLimit := MaxRequestBodyBytes + MaxRequestBodyBytes = tt.limit + defer func() { MaxRequestBodyBytes = prevLimit }() + + counter := &countingReader{r: strings.NewReader(modelRequestBody(t, tt.bodySize))} + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request, _ = http.NewRequest(http.MethodPost, "/v1/chat/completions", io.NopCloser(counter)) + + got, err := ParseModelRequest(c) + if tt.wantStatus != 0 { + assert.Error(t, err) + assert.Nil(t, got) + assert.Equal(t, tt.wantStatus, w.Code) + assert.Contains(t, w.Body.String(), "request body exceeds") + // The payload must be rejected before it is fully buffered. + assert.LessOrEqual(t, counter.n, tt.limit+1) + return + } + assert.NoError(t, err) + assert.Equal(t, "test-model", got["model"]) + assert.Equal(t, int64(tt.bodySize), counter.n) + }) + } +} + +func TestParseMaxRequestBodyBytes(t *testing.T) { + tests := []struct { + name string + env string + set bool + want int64 + }{ + {name: "unset uses default", want: defaultMaxRequestBodyBytes}, + {name: "explicit value", env: "1048576", set: true, want: 1048576}, + {name: "zero disables the limit", env: "0", set: true, want: 0}, + {name: "negative disables the limit", env: "-1", set: true, want: -1}, + {name: "invalid value uses default", env: "1MiB", set: true, want: defaultMaxRequestBodyBytes}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if tt.set { + t.Setenv("MAX_REQUEST_BODY_BYTES", tt.env) + } else { + prev, existed := os.LookupEnv("MAX_REQUEST_BODY_BYTES") + os.Unsetenv("MAX_REQUEST_BODY_BYTES") + defer func() { + if existed { + os.Setenv("MAX_REQUEST_BODY_BYTES", prev) + } + }() + } + assert.Equal(t, tt.want, parseMaxRequestBodyBytes()) + }) + } +} + func TestAccessLogConfigurationFromEnv(t *testing.T) { // Save original environment variables originalEnabled := os.Getenv("ACCESS_LOG_ENABLED")