From d5bff28c5cd4075b93e1b0a92b715ad737e440d6 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 00:26:28 +0100 Subject: [PATCH] =?UTF-8?q?fix:=20harden=20scaffold=20after=20review=20?= =?UTF-8?q?=E2=80=94=20manifest=20drift=20check,=20vet-safe=20logging,=20p?= =?UTF-8?q?anic-path=20header=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- errors.go | 49 ++++++++++++++++++++--- llmgateway.go | 7 ++-- llmgateway_test.go | 96 ++++++++++++++++++++++++++++++++++++++++++++-- logger.go | 4 +- manifest_test.go | 32 ++++++++++++++-- 5 files changed, 170 insertions(+), 18 deletions(-) diff --git a/errors.go b/errors.go index f4878f7..c432693 100644 --- a/errors.go +++ b/errors.go @@ -4,23 +4,62 @@ import ( "encoding/json" "net/http" "runtime/debug" + "strconv" ) +// statusTrackingWriter wraps an http.ResponseWriter to record whether a +// response has already been committed (header written, or body bytes +// written under the net/http implicit-200 rule). recoverPanic uses this to +// avoid writing a second, conflicting response after a handler panics +// mid-stream. Flush delegates to the underlying http.Flusher when present — +// later tasks' SSE streaming depends on Flush passing through. +type statusTrackingWriter struct { + http.ResponseWriter + wroteHeader bool +} + +// WriteHeader records that the response was committed, then delegates. +func (w *statusTrackingWriter) WriteHeader(status int) { + w.wroteHeader = true + w.ResponseWriter.WriteHeader(status) +} + +// Write records an implicit header commit if none happened yet, then +// delegates. +func (w *statusTrackingWriter) Write(b []byte) (int, error) { + w.wroteHeader = true + return w.ResponseWriter.Write(b) +} + +// Flush delegates to the underlying ResponseWriter's Flush when it +// implements http.Flusher, and is a no-op otherwise. +func (w *statusTrackingWriter) Flush() { + if f, ok := w.ResponseWriter.(http.Flusher); ok { + f.Flush() + } +} + // writeOAIError writes an OpenAI-compatible error envelope: -// {"error":{"message":...,"type":...,"code":...}}. +// {"error":{"message":...,"type":...,"code":...}}. code is a JSON string — +// OpenAI-compatible SDKs expect string-or-null, not an integer. func writeOAIError(w http.ResponseWriter, status int, errType, msg string) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) _ = json.NewEncoder(w).Encode(map[string]any{ - "error": map[string]any{"message": msg, "type": errType, "code": status}, + "error": map[string]any{"message": msg, "type": errType, "code": strconv.Itoa(status)}, }) } -// recoverPanic recovers a panic from a handler, logs it, and writes a 500 -// error envelope instead of letting the panic escape ServeHTTP. -func recoverPanic(w http.ResponseWriter, g *Gateway) { +// recoverPanic recovers a panic from a handler and logs it. It writes a 500 +// error envelope only if the response has not already been committed — a +// handler that panics after writing headers or body bytes (mid-stream) must +// not get a second, conflicting response appended. +func recoverPanic(w *statusTrackingWriter, g *Gateway) { if rec := recover(); rec != nil { g.errorf("panic recovered: %v\n%s", rec, debug.Stack()) + if w.wroteHeader { + return + } writeOAIError(w, http.StatusInternalServerError, "server_error", "internal error") } } diff --git a/llmgateway.go b/llmgateway.go index 809ac9b..96915dd 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -125,10 +125,11 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str // ServeHTTP is the internal router. It grows in later tasks. func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { - defer recoverPanic(w, g) + sw := &statusTrackingWriter{ResponseWriter: w} + defer recoverPanic(sw, g) if g.cfg.PassthroughUnknown { - g.next.ServeHTTP(w, r) + g.next.ServeHTTP(sw, r) return } - writeOAIError(w, http.StatusNotFound, "invalid_request_error", "unknown route") + writeOAIError(sw, http.StatusNotFound, "invalid_request_error", "unknown route") } diff --git a/llmgateway_test.go b/llmgateway_test.go index 577311d..a481a63 100644 --- a/llmgateway_test.go +++ b/llmgateway_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "os" + "strconv" "strings" "testing" ) @@ -83,7 +84,7 @@ func TestWriteOAIError_WritesEnvelope(t *testing.T) { Error struct { Message string `json:"message"` Type string `json:"type"` - Code int `json:"code"` + Code string `json:"code"` } `json:"error"` } if err := json.NewDecoder(rec.Body).Decode(&body); err != nil { @@ -95,8 +96,9 @@ func TestWriteOAIError_WritesEnvelope(t *testing.T) { if body.Error.Type != "invalid_request_error" { t.Fatalf("want type %q, got %q", "invalid_request_error", body.Error.Type) } - if body.Error.Code != http.StatusBadRequest { - t.Fatalf("want code %d, got %d", http.StatusBadRequest, body.Error.Code) + wantCode := strconv.Itoa(http.StatusBadRequest) + if body.Error.Code != wantCode { + t.Fatalf("want code %q (JSON string), got %q", wantCode, body.Error.Code) } } @@ -119,8 +121,9 @@ func TestServeHTTP_PanicRecovery_Returns500Envelope(t *testing.T) { defer func() { os.Stderr = origStderr }() rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} func() { - defer recoverPanic(rec, gw) + defer recoverPanic(sw, gw) panic("boom") }() @@ -199,3 +202,88 @@ func TestGateway_Errorf_WritesErrorPrefix(t *testing.T) { t.Fatalf("want stderr to contain %q, got %q", want, buf.String()) } } + +func TestStatusTrackingWriter_WriteHeaderSetsFlag(t *testing.T) { + rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} + if sw.wroteHeader { + t.Fatal("want wroteHeader false before any write") + } + sw.WriteHeader(http.StatusAccepted) + if !sw.wroteHeader { + t.Fatal("want wroteHeader true after WriteHeader") + } + if rec.Code != http.StatusAccepted { + t.Fatalf("want delegated status %d, got %d", http.StatusAccepted, rec.Code) + } +} + +func TestStatusTrackingWriter_WriteSetsFlag(t *testing.T) { + rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} + if _, err := sw.Write([]byte("hi")); err != nil { + t.Fatalf("Write: %v", err) + } + if !sw.wroteHeader { + t.Fatal("want wroteHeader true after an implicit-header Write") + } + if rec.Body.String() != "hi" { + t.Fatalf("want delegated body %q, got %q", "hi", rec.Body.String()) + } +} + +func TestStatusTrackingWriter_FlushDelegates(t *testing.T) { + rec := httptest.NewRecorder() + sw := &statusTrackingWriter{ResponseWriter: rec} + sw.Flush() + if !rec.Flushed { + t.Fatal("want Flush to delegate to the underlying http.Flusher") + } +} + +// TestServeHTTP_PanicAfterHeadersCommitted_DoesNotOverwriteResponse is the +// regression test for the recoverPanic header-guard: a handler that panics +// after it has already written a status code and body bytes (e.g. +// mid-stream) must not get a second, conflicting error envelope appended. +func TestServeHTTP_PanicAfterHeadersCommitted_DoesNotOverwriteResponse(t *testing.T) { + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + if _, err := w.Write([]byte("partial")); err != nil { + t.Fatalf("next handler Write: %v", err) + } + panic("boom mid-stream") + }) + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.PassthroughUnknown = true + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + origStderr := os.Stderr + pr, pw, _ := os.Pipe() + os.Stderr = pw + defer func() { os.Stderr = origStderr }() + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/stream", nil) + h.ServeHTTP(rec, req) + + _ = pw.Close() // closing the pipe write end to unblock the read; error not actionable in a test + os.Stderr = origStderr + var logBuf bytes.Buffer + if _, err := io.Copy(&logBuf, pr); err != nil { + t.Fatalf("io.Copy: %v", err) + } + + if rec.Code != http.StatusOK { + t.Fatalf("want status left at %d (already committed by next), got %d", http.StatusOK, rec.Code) + } + if got := rec.Body.String(); got != "partial" { + t.Fatalf("want body left as %q (no error envelope appended), got %q", "partial", got) + } + if !strings.Contains(logBuf.String(), "llmgw[llmgw] ERROR ") { + t.Fatalf("want panic still logged even though no envelope was written, got %q", logBuf.String()) + } +} diff --git a/logger.go b/logger.go index c60cad6..1cd8905 100644 --- a/logger.go +++ b/logger.go @@ -8,11 +8,11 @@ import ( // logf writes an info-level log line to stderr. No timestamps — Traefik // adds its own. Never log key material. func (g *Gateway) logf(format string, args ...any) { - fmt.Fprintf(os.Stderr, "llmgw[%s] INFO "+format+"\n", append([]any{g.name}, args...)...) + fmt.Fprintf(os.Stderr, "llmgw[%s] INFO %s\n", g.name, fmt.Sprintf(format, args...)) } // errorf writes an error-level log line to stderr. No timestamps — Traefik // adds its own. Never log key material. func (g *Gateway) errorf(format string, args ...any) { - fmt.Fprintf(os.Stderr, "llmgw[%s] ERROR "+format+"\n", append([]any{g.name}, args...)...) + fmt.Fprintf(os.Stderr, "llmgw[%s] ERROR %s\n", g.name, fmt.Sprintf(format, args...)) } diff --git a/manifest_test.go b/manifest_test.go index 2dc3dc4..1ac5bbf 100644 --- a/manifest_test.go +++ b/manifest_test.go @@ -1,6 +1,7 @@ package traefikllmgateway import ( + "bytes" "context" "encoding/json" "net/http" @@ -10,10 +11,13 @@ import ( "testing" ) -// yamlToJSONLite converts the small, known-shape testData block of -// .traefik.yml to JSON via a scripted extraction — we cannot import a YAML -// lib. Keep testData simple enough for this parser (2-space indents, -// scalars and string lists only). +// testDataConfig decodes a hand-maintained JSON mirror of the .traefik.yml +// testData block — we cannot import a YAML library (stdlib only). The +// mirror is drift-checked against the raw manifest by the assertions below: +// every top-level key and scalar value the mirror encodes must also appear +// in .traefik.yml's raw bytes, so an edit to either side that removes or +// renames a testData field fails this test instead of silently going +// stale. func testDataConfig(t *testing.T) *Config { t.Helper() raw, err := os.ReadFile(".traefik.yml") @@ -23,6 +27,26 @@ func testDataConfig(t *testing.T) *Config { if !regexp.MustCompile(`(?m)^testData:`).Match(raw) { t.Fatal(".traefik.yml missing testData") } + + // Drift check: every top-level key and scalar value referenced by the + // mirror blob below must appear in the raw manifest bytes. + mirrorTokens := []string{ + "providers", "groups", "users", // top-level testData keys + "openai", "type: openai", "test-key", "gpt-test", // providers.openai + "default", // groups.default and users.inline[0].group + "tester", // users.inline[0].name + "test-user-key", // users.inline[0].apiKey + } + var missing []string + for _, tok := range mirrorTokens { + if !bytes.Contains(raw, []byte(tok)) { + missing = append(missing, tok) + } + } + if len(missing) > 0 { + t.Fatalf("manifest drift: mirror blob references %v not found in .traefik.yml — update the mirror or the manifest so they match", missing) + } + cfg := CreateConfig() // mirror of testData, kept in sync by this test living next to the manifest blob := `{"providers":{"openai":{"type":"openai","apiKey":"test-key","models":["gpt-test"]}},