refactor(proxy): decompose ServeHTTP into 6 helpers (comp 38→9, CRAP<30)
Extract handleAuthEndpoint, routeOrError, writeModelValidationError, authenticateV1, enforceBodyLimit, setupCAGProxy. No behavior change. All tests pass with -race.
This commit is contained in:
Admin Bot
+229
-212
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user