diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 2f268e7..ce6507c 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -302,17 +302,7 @@ func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, d // ServeHTTP implements http.Handler. func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - // Auth endpoints — no JWT required (they issue tokens) - if r.URL.Path == "/auth/token" { - h.handleAuthToken(w, r) - return - } - if r.URL.Path == "/auth/refresh" { - h.handleAuthRefresh(w, r) - return - } - if r.URL.Path == "/auth/exchange" { - h.handleAuthExchange(w, r) + if h.handleAuthEndpoint(w, r) { return } @@ -322,68 +312,9 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - // Try to find a matching route (including body-based dispatch for /v1/chat/completions). - // Note: /workflows endpoint is deprecated. Use X-Service: workflow + X-Resource headers instead. - route, err := h.RouteRequest(r) - - // Check if this is a model validation error (from body-based dispatch) - if validationErr, ok := err.(*modelValidationError); ok { - // This is a client error, not a routing error - var status int - var problemType string - var title string - var detail string - - switch validationErr.Kind { - case "invalid_json": - status = http.StatusBadRequest - problemType = "https://api.example.com/problems/invalid-request-body" - title = "Invalid Request Body" - detail = validationErr.Message - case "missing_model", "empty_model", "null_model": - status = http.StatusBadRequest - problemType = "https://api.example.com/problems/missing-model" - title = "Missing Model" - detail = "The 'model' field is required and must be a non-empty string" - case "unknown_model": - status = http.StatusBadRequest - problemType = "https://api.example.com/problems/unknown-model" - title = "Unknown Model" - detail = fmt.Sprintf("Model %q is not available. See valid_models for available options.", validationErr.Model) - default: - status = http.StatusBadRequest - problemType = "https://api.example.com/problems/invalid-request" - title = "Invalid Request" - detail = validationErr.Message - } - - // Get list of valid models (only for model-related errors) - var validModels []string - if validationErr.Kind == "unknown_model" || validationErr.Kind == "missing_model" || validationErr.Kind == "empty_model" || validationErr.Kind == "null_model" { - validModels = h.getValidModels() - } - - writeProblemDetail(w, status, problemType, title, detail, validModels) - logging.Errorf("client error", validationErr, map[string]string{ - "path": r.URL.Path, - "method": r.Method, - "reason": validationErr.Kind, - }) - return - } - - if err != nil || route == nil { - // Route not found or error determining route - w.WriteHeader(http.StatusNotFound) - fmt.Fprintf(w, "not found") - if err != nil { - logging.Errorf("routing failed", err, map[string]string{ - "path": r.URL.Path, - "method": r.Method, - }) - } else { - logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil) - } + // Resolve route; handles validation errors and 404. + route, ok := h.routeOrError(w, r) + if !ok { return } @@ -391,154 +322,30 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { // Must happen before any routing — even unauthenticated paths. identity.StripIncoming(r) - // JWT Authentication for /v1/* endpoints - if h.jwtValidator != nil && strings.HasPrefix(r.URL.Path, "/v1/") { - authHeader := r.Header.Get("Authorization") - if authHeader == "" { - writeProblemDetail(w, http.StatusUnauthorized, - "https://api.example.com/problems/unauthorized", - "Unauthorized", - "Authorization header required", - nil) - logging.Errorf("auth failed", fmt.Errorf("missing auth header"), map[string]string{ - "path": r.URL.Path, - }) - return - } - - claims, err := h.jwtValidator.ValidateBearerToken(authHeader) - if err != nil { - writeProblemDetail(w, http.StatusForbidden, - "https://api.example.com/problems/forbidden", - "Forbidden", - "JWT validation failed", - nil) - logging.Errorf("auth failed", err, map[string]string{ - "path": r.URL.Path, - }) - return - } - - // Inject identity headers for downstream services - identity.Inject(r, claims) - - // Check required capability if configured - if h.config.Auth.RequiredCapability != "" { - if !h.jwtValidator.CheckPermissions(claims, h.config.Auth.RequiredCapability, "*") { - writeProblemDetail(w, http.StatusForbidden, - "https://api.example.com/problems/insufficient-permissions", - "Insufficient Permissions", - fmt.Sprintf("Required capability: %s", h.config.Auth.RequiredCapability), - nil) - logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{ - "path": r.URL.Path, - "required": h.config.Auth.RequiredCapability, - }) - return - } - } + // JWT Authentication for /v1/* endpoints. + if !h.authenticateV1(w, r) { + return } - // Note: Body size checking already happened in RouteRequest (body was read for model dispatch). - // For other paths, we still need to enforce the cap. - // For /v1/chat/completions, the body was already read and validated. - - // Enforce request body size cap for non-chat routes - if r.URL.Path != "/v1/chat/completions" { - if route.Upstream.MaxBodySize > 0 && r.ContentLength > route.Upstream.MaxBodySize { - w.WriteHeader(http.StatusRequestEntityTooLarge) - fmt.Fprintf(w, "request body too large") - logging.Errorf("request rejected", fmt.Errorf("body_too_large"), map[string]string{ - "reason": "body_too_large", - "route": route.Name, - "upstream": route.Upstream.Address, - "content_length": fmt.Sprintf("%d", r.ContentLength), - "max_body_size": fmt.Sprintf("%d", route.Upstream.MaxBodySize), - }) - return - } - - // Wrap request body with size limiter - // This enforces the cap at read time, not after buffering - if route.Upstream.MaxBodySize > 0 && r.Body != nil { - r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize)) - } + // Enforce request body size cap for non-chat routes. + if !h.enforceBodyLimit(w, r, route) { + return } - // CAG: warmup deprioritisation + per-backend semaphore (issues #45/#46) - var selectedBackend *Backend + // CAG: warmup deprioritisation + per-backend semaphore (issues #45/#46). cacheKey := r.Header.Get("X-Poimen-Cache-Key") - if sched := h.schedulers[route.ModelName]; sched != nil { - var ok bool - selectedBackend, route, ok = h.acquireScheduledBackend(r.Context(), w, r, sched, route, cacheKey) - if !ok { - return - } - if selectedBackend != nil { - defer selectedBackend.Release() - } + sched := h.schedulers[route.ModelName] + proxy, backend, ok := h.setupCAGProxy(r.Context(), w, r, route, sched, cacheKey) + if !ok { + return + } + if backend != nil { + defer backend.Release() } - // Create the reverse proxy - proxy := httputil.NewSingleHostReverseProxy(&url.URL{ - Scheme: "http", - Host: route.Upstream.Address, - }) - - // Set the director to apply path rewriting - proxy.Director = route.Director - - // Use the connection-pooled transport wrapped with tracing - proxy.Transport = tracing.NewTransport(route.Transport) - - // Set error handler: log upstream errors and exclude failed backend from routing. - proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) { - if selectedBackend != nil { - if sched := h.schedulers[route.ModelName]; sched != nil { - sched.MarkFailed(selectedBackend.Address) - } - } - w.WriteHeader(http.StatusBadGateway) - fmt.Fprintf(w, "upstream error") - logging.Errorf("upstream error", err, map[string]string{ - "upstream": route.Upstream.Address, - "path": r.URL.Path, - }) - } - - // CAG: capture vllm-num-cached-tokens from response headers (issue #46) - if cacheKey != "" && r.URL.Path == "/v1/chat/completions" { - proxy.ModifyResponse = func(resp *http.Response) error { - raw := resp.Header.Get("vllm-num-cached-tokens") - if raw == "" { - return nil - } - n, err := strconv.ParseInt(raw, 10, 64) - if err != nil || n <= 0 { - return nil - } - observability.Global().IncrPrefixCacheHitTokens(n) - logging.Event(map[string]interface{}{ - "event": "llm_request", - "cache_key": cacheKey, - "cached_tokens": n, - }) - return nil - } - } - - // Note: We apply connection timeout via Transport dialer, but NOT read timeout as a context deadline. - // Read timeout should apply to inactivity (socket read timeout), not total request duration. - // A streaming response that's continuously sending should not be cut off. - // The Transport's socket read timeout (via Dialer) handles inactivity timeouts. - - // For streaming responses (SSE, chunked), disable buffering to ensure events - // reach clients immediately. Issue #33: X-Accel-Buffering:no tells nginx/Ingress - // to stream instead of buffer. ResponseController.Flush() in upstream handler - // pairs with this to deliver unbuffered chunks. + // For streaming responses (SSE, chunked), disable buffering. Issue #33. w.Header().Set("X-Accel-Buffering", "no") - // Serve the request through the proxy proxy.ServeHTTP(w, r) } @@ -606,6 +413,216 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr return backend, route, true } +// handleAuthEndpoint routes /auth/* requests and returns true if handled. +func (h *Handler) handleAuthEndpoint(w http.ResponseWriter, r *http.Request) bool { + switch r.URL.Path { + case "/auth/token": + h.handleAuthToken(w, r) + return true + case "/auth/refresh": + h.handleAuthRefresh(w, r) + return true + case "/auth/exchange": + h.handleAuthExchange(w, r) + return true + } + return false +} + +// routeOrError resolves a route for r, writing error responses on failure. +// Returns (route, true) on success or (nil, false) when a response has been written. +func (h *Handler) routeOrError(w http.ResponseWriter, r *http.Request) (*Route, bool) { + route, err := h.RouteRequest(r) + if validationErr, ok := err.(*modelValidationError); ok { + h.writeModelValidationError(w, validationErr, h.getValidModels()) + logging.Errorf("client error", validationErr, map[string]string{ + "path": r.URL.Path, + "method": r.Method, + "reason": validationErr.Kind, + }) + return nil, false + } + if err != nil || route == nil { + w.WriteHeader(http.StatusNotFound) + fmt.Fprintf(w, "not found") + if err != nil { + logging.Errorf("routing failed", err, map[string]string{ + "path": r.URL.Path, + "method": r.Method, + }) + } else { + logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil) + } + return nil, false + } + return route, true +} + +// writeModelValidationError writes a problem-detail response for model validation errors. +func (h *Handler) writeModelValidationError(w http.ResponseWriter, validationErr *modelValidationError, validModels []string) { + var status int + var problemType, title, detail string + var models []string + + switch validationErr.Kind { + case "invalid_json": + status = http.StatusBadRequest + problemType = "https://api.example.com/problems/invalid-request-body" + title = "Invalid Request Body" + detail = validationErr.Message + case "missing_model", "empty_model", "null_model": + status = http.StatusBadRequest + problemType = "https://api.example.com/problems/missing-model" + title = "Missing Model" + detail = "The 'model' field is required and must be a non-empty string" + models = validModels + case "unknown_model": + status = http.StatusBadRequest + problemType = "https://api.example.com/problems/unknown-model" + title = "Unknown Model" + detail = fmt.Sprintf("Model %q is not available. See valid_models for available options.", validationErr.Model) + models = validModels + default: + status = http.StatusBadRequest + problemType = "https://api.example.com/problems/invalid-request" + title = "Invalid Request" + detail = validationErr.Message + } + + writeProblemDetail(w, status, problemType, title, detail, models) +} + +// authenticateV1 validates JWT for /v1/* paths. +// Returns false when auth fails (response already written). +func (h *Handler) authenticateV1(w http.ResponseWriter, r *http.Request) bool { + if h.jwtValidator == nil || !strings.HasPrefix(r.URL.Path, "/v1/") { + return true + } + authHeader := r.Header.Get("Authorization") + if authHeader == "" { + writeProblemDetail(w, http.StatusUnauthorized, + "https://api.example.com/problems/unauthorized", + "Unauthorized", + "Authorization header required", + nil) + logging.Errorf("auth failed", fmt.Errorf("missing auth header"), map[string]string{ + "path": r.URL.Path, + }) + return false + } + claims, err := h.jwtValidator.ValidateBearerToken(authHeader) + if err != nil { + writeProblemDetail(w, http.StatusForbidden, + "https://api.example.com/problems/forbidden", + "Forbidden", + "JWT validation failed", + nil) + logging.Errorf("auth failed", err, map[string]string{ + "path": r.URL.Path, + }) + return false + } + identity.Inject(r, claims) + if h.config.Auth.RequiredCapability != "" { + if !h.jwtValidator.CheckPermissions(claims, h.config.Auth.RequiredCapability, "*") { + writeProblemDetail(w, http.StatusForbidden, + "https://api.example.com/problems/insufficient-permissions", + "Insufficient Permissions", + fmt.Sprintf("Required capability: %s", h.config.Auth.RequiredCapability), + nil) + logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{ + "path": r.URL.Path, + "required": h.config.Auth.RequiredCapability, + }) + return false + } + } + return true +} + +// enforceBodyLimit caps request body size for non-chat routes. +// Returns false when the request is rejected (response already written). +func (h *Handler) enforceBodyLimit(w http.ResponseWriter, r *http.Request, route *Route) bool { + if r.URL.Path == "/v1/chat/completions" { + return true + } + if route.Upstream.MaxBodySize > 0 && r.ContentLength > route.Upstream.MaxBodySize { + w.WriteHeader(http.StatusRequestEntityTooLarge) + fmt.Fprintf(w, "request body too large") + logging.Errorf("request rejected", fmt.Errorf("body_too_large"), map[string]string{ + "reason": "body_too_large", + "route": route.Name, + "upstream": route.Upstream.Address, + "content_length": fmt.Sprintf("%d", r.ContentLength), + "max_body_size": fmt.Sprintf("%d", route.Upstream.MaxBodySize), + }) + return false + } + if route.Upstream.MaxBodySize > 0 && r.Body != nil { + r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize)) + } + return true +} + +// setupCAGProxy acquires a scheduled backend (if any), builds a reverse proxy, +// wires ErrorHandler and ModifyResponse, and returns (proxy, backend, true). +// Returns (nil, nil, false) when the backend slot cannot be acquired (response written). +func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *http.Request, route *Route, sched *Scheduler, cacheKey string) (*httputil.ReverseProxy, *Backend, bool) { + var selectedBackend *Backend + if sched != nil { + var ok bool + selectedBackend, route, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey) + if !ok { + return nil, nil, false + } + } + + proxy := httputil.NewSingleHostReverseProxy(&url.URL{ + Scheme: "http", + Host: route.Upstream.Address, + }) + proxy.Director = route.Director + proxy.Transport = tracing.NewTransport(route.Transport) + + capturedBackend := selectedBackend + capturedRoute := route + proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) { + if capturedBackend != nil { + if s := h.schedulers[capturedRoute.ModelName]; s != nil { + s.MarkFailed(capturedBackend.Address) + } + } + w.WriteHeader(http.StatusBadGateway) + fmt.Fprintf(w, "upstream error") + logging.Errorf("upstream error", err, map[string]string{ + "upstream": capturedRoute.Upstream.Address, + "path": r.URL.Path, + }) + } + + if cacheKey != "" && r.URL.Path == "/v1/chat/completions" { + proxy.ModifyResponse = func(resp *http.Response) error { + raw := resp.Header.Get("vllm-num-cached-tokens") + if raw == "" { + return nil + } + n, err := strconv.ParseInt(raw, 10, 64) + if err != nil || n <= 0 { + return nil + } + observability.Global().IncrPrefixCacheHitTokens(n) + logging.Event(map[string]interface{}{ + "event": "llm_request", + "cache_key": cacheKey, + "cached_tokens": n, + }) + return nil + } + } + + return proxy, selectedBackend, true +} + // getValidModels returns a sorted list of all configured model names. func (h *Handler) getValidModels() []string { var models []string