From e6ab8f128cb7264271f77336c1b1f7fc15c99ec2 Mon Sep 17 00:00:00 2001 From: tlnarayana005 Date: Tue, 7 Jul 2026 23:39:30 +0530 Subject: [PATCH] fix(gateway): harden AI provider HTTP error handling and context propagation - Cap error response body reads to 4KB in both OpenRouter and Ollama providers. Previously, io.ReadAll(resp.Body) on non-2xx responses could read unlimited data from a malicious/misconfigured upstream, causing OOM. Now uses io.LimitReader to cap at 4096 bytes, matching the existing pattern in verifyPayment(). - Handle context.Canceled alongside context.DeadlineExceeded in AI provider error paths. During graceful shutdown, Go cancels contexts (not deadline-exceeds them). Without this fix, shutdown cancellations were misclassified and wrapped in connection-error messages instead of being returned as clean context errors. - Apply the same context.Canceled fix to handleSummarize() in main.go so the gateway returns 504 (not 502) during shutdown-induced cancellations, giving clients actionable retry signals. - Add unit tests covering both the body-limit cap and context.Canceled behavior for both providers. --- gateway/internal/ai/error_handling_test.go | 116 +++++++++++++++++++++ gateway/internal/ai/ollama.go | 12 ++- gateway/internal/ai/openrouter.go | 13 ++- gateway/main.go | 4 +- 4 files changed, 138 insertions(+), 7 deletions(-) create mode 100644 gateway/internal/ai/error_handling_test.go diff --git a/gateway/internal/ai/error_handling_test.go b/gateway/internal/ai/error_handling_test.go new file mode 100644 index 00000000..00036deb --- /dev/null +++ b/gateway/internal/ai/error_handling_test.go @@ -0,0 +1,116 @@ +package ai + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// TestOpenRouterErrorBodyLimit verifies that the OpenRouter provider caps error +// response body reads at maxErrorResponseBytes to prevent OOM from a malicious +// upstream. +func TestOpenRouterErrorBodyLimit(t *testing.T) { + // Serve a body larger than the limit (8KB > 4KB limit) + oversizedBody := strings.Repeat("X", 8192) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(oversizedBody)) + })) + defer srv.Close() + + p := &OpenRouterProvider{ + apiKey: "test-key", + model: "test-model", + url: srv.URL, + } + + _, err := p.Generate(context.Background(), "test prompt") + if err == nil { + t.Fatal("expected error from 500 response, got nil") + } + + // The error message should NOT contain the full 8KB body + errMsg := err.Error() + if len(errMsg) > int(maxErrorResponseBytes)+200 { + t.Errorf("error message too long (%d bytes), expected capped at ~%d", len(errMsg), maxErrorResponseBytes) + } +} + +// TestOllamaErrorBodyLimit verifies that the Ollama provider caps error +// response body reads at maxOllamaErrorResponseBytes. +func TestOllamaErrorBodyLimit(t *testing.T) { + oversizedBody := strings.Repeat("Y", 8192) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(oversizedBody)) + })) + defer srv.Close() + + p := &OllamaProvider{ + url: srv.URL, + model: "test-model", + } + + _, err := p.Generate(context.Background(), "test prompt") + if err == nil { + t.Fatal("expected error from 500 response, got nil") + } + + errMsg := err.Error() + if len(errMsg) > int(maxOllamaErrorResponseBytes)+200 { + t.Errorf("error message too long (%d bytes), expected capped at ~%d", len(errMsg), maxOllamaErrorResponseBytes) + } +} + +// TestOpenRouterContextCanceled verifies that context.Canceled is returned +// (not wrapped in a connection error) when the context is canceled. +func TestOpenRouterContextCanceled(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Never respond — let the context cancel + <-r.Context().Done() + })) + defer srv.Close() + + p := &OpenRouterProvider{ + apiKey: "test-key", + model: "test-model", + url: srv.URL, + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + _, err := p.Generate(ctx, "test prompt") + if err == nil { + t.Fatal("expected error from canceled context, got nil") + } + if err != context.Canceled { + t.Errorf("expected context.Canceled, got: %v", err) + } +} + +// TestOllamaContextCanceled verifies the same for the Ollama provider. +func TestOllamaContextCanceled(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + })) + defer srv.Close() + + p := &OllamaProvider{ + url: srv.URL, + model: "test-model", + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := p.Generate(ctx, "test prompt") + if err == nil { + t.Fatal("expected error from canceled context, got nil") + } + if err != context.Canceled { + t.Errorf("expected context.Canceled, got: %v", err) + } +} diff --git a/gateway/internal/ai/ollama.go b/gateway/internal/ai/ollama.go index 57847cba..9b70d16e 100644 --- a/gateway/internal/ai/ollama.go +++ b/gateway/internal/ai/ollama.go @@ -11,6 +11,11 @@ import ( "os" ) +// maxErrorResponseBytes caps the amount of upstream error response body we +// read on non-200 responses. Prevents OOM if a misbehaving Ollama instance +// streams an unbounded error payload. +const maxOllamaErrorResponseBytes int64 = 4096 + // OllamaProvider implements the Provider interface for Ollama API type OllamaProvider struct { url string @@ -53,15 +58,16 @@ func (p *OllamaProvider) Generate(ctx context.Context, text string) (string, err resp, err := http.DefaultClient.Do(req) if err != nil { - if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded { - return "", context.DeadlineExceeded + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || + ctx.Err() == context.DeadlineExceeded || ctx.Err() == context.Canceled { + return "", ctx.Err() } return "", fmt.Errorf("failed to connect to Ollama: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) + body, _ := io.ReadAll(io.LimitReader(resp.Body, maxOllamaErrorResponseBytes)) return "", fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, string(body)) } diff --git a/gateway/internal/ai/openrouter.go b/gateway/internal/ai/openrouter.go index 36bc48ab..b7a27607 100644 --- a/gateway/internal/ai/openrouter.go +++ b/gateway/internal/ai/openrouter.go @@ -12,6 +12,12 @@ import ( "os" ) +// maxErrorResponseBytes caps the amount of upstream error response body we are +// willing to read. Without this limit, a misbehaving or malicious upstream +// could stream an unbounded response on an error status, exhausting gateway +// memory. +const maxErrorResponseBytes int64 = 4096 + // OpenRouterProvider implements the Provider interface for OpenRouter API type OpenRouterProvider struct { apiKey string @@ -59,8 +65,9 @@ func (p *OpenRouterProvider) Generate(ctx context.Context, text string) (string, resp, err := http.DefaultClient.Do(req) if err != nil { - if errors.Is(err, context.DeadlineExceeded) || ctx.Err() == context.DeadlineExceeded { - return "", context.DeadlineExceeded + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || + ctx.Err() == context.DeadlineExceeded || ctx.Err() == context.Canceled { + return "", ctx.Err() } return "", err } @@ -68,7 +75,7 @@ func (p *OpenRouterProvider) Generate(ctx context.Context, text string) (string, // Check status code before decoding if resp.StatusCode < 200 || resp.StatusCode >= 300 { - body, _ := io.ReadAll(resp.Body) + body, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorResponseBytes)) return "", fmt.Errorf("openrouter returned status %d: %s", resp.StatusCode, string(body)) } diff --git a/gateway/main.go b/gateway/main.go index 953c94d1..8d68e524 100644 --- a/gateway/main.go +++ b/gateway/main.go @@ -523,7 +523,9 @@ func handleSummarize(c *gin.Context) { // 3. Call AI Service summary, err := aiProvider.Generate(c.Request.Context(), req.Text) if err != nil { - if errors.Is(err, context.DeadlineExceeded) || errors.Is(c.Request.Context().Err(), context.DeadlineExceeded) { + ctxErr := c.Request.Context().Err() + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || + ctxErr == context.DeadlineExceeded || ctxErr == context.Canceled { respondError(c, 504, "upstream_timeout", err) } else { respondError(c, 502, "upstream_unavailable", err)