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 56a3683300..d8dcf22aaf 100644 --- a/charts/kthena/values.yaml +++ b/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.
+ # 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 39dc53cc56..24c69e1156 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.replicas | int | `1` | Number of Kthena Router instances to run. | | 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`). | diff --git a/pkg/kthena-router/router/router.go b/pkg/kthena-router/router/router.go index 50b216d77c..125e6b2163 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 @@ -591,8 +617,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 4a9a4399ec..22a13deb14 100644 --- a/pkg/kthena-router/router/router_test.go +++ b/pkg/kthena-router/router/router_test.go @@ -1388,6 +1388,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")