diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d7827d..2fb3944 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ Gatekeeper is a standalone credential-injecting TLS-intercepting proxy. It trans Gatekeeper is pre-1.0. The configuration schema and credential source interface may change between minor versions. +## v0.14.1 — 2026-06-23 + +### Fixed + +- **Streamed responses are no longer buffered during capture** — the response-body log sampler did a blocking read of up to `MaxBodySize` (8 KB) before forwarding, so any incrementally-produced text response (Server-Sent Events, `application/x-ndjson`, chunked JSON, …) had its status line and every chunk — including the keepalive pings an upstream sends during a long time-to-first-token — withheld from the client until 8 KB accumulated, or until the stream ended for responses under 8 KB. The client received nothing, tripped its first-byte timeout, and retried in a loop; most visibly on large or cache-cold LLM streaming requests (e.g. `/v1/messages`), which worked when sent directly (un-proxied). The sampler now captures text responses lazily via a non-blocking tee instead of a blocking read-ahead, so a slow or streamed response (SSE, ndjson, chunked JSON) is forwarded immediately rather than withheld. The canonical log line for a text response is now written when the body completes rather than when its headers arrive; non-text and non-streaming responses are unchanged. This does **not** apply to hosts with an `llm-gateway` Keep engine: response-policy evaluation reads the full body before forwarding, so streams on those hosts are still buffered end-to-end ([#37](https://github.com/majorcontext/gatekeeper/pull/37)) + ## v0.14.0 — 2026-06-23 ### Added diff --git a/proxy/intercept_test.go b/proxy/intercept_test.go index 6b7f5ec..4017b17 100644 --- a/proxy/intercept_test.go +++ b/proxy/intercept_test.go @@ -15,10 +15,31 @@ import ( "sync" "sync/atomic" "testing" + "time" keeplib "github.com/majorcontext/keep" ) +// captureLog installs a logger that records logged requests and returns a +// function that waits for the next one. The canonical log line for a text +// response is written when its body is closed (capture streams lazily), so +// tests must wait for it rather than read a shared variable synchronously. +func captureLog(t *testing.T, p *Proxy) func() RequestLogData { + t.Helper() + ch := make(chan RequestLogData, 8) + p.SetLogger(func(d RequestLogData) { ch <- d }) + return func() RequestLogData { + t.Helper() + select { + case d := <-ch: + return d + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for canonical log line") + return RequestLogData{} + } + } +} + // interceptTestSetup creates a proxy with TLS interception enabled and an HTTPS // backend server. The proxy is configured to trust the backend's TLS cert and // the returned client trusts the proxy's interception CA. @@ -110,17 +131,16 @@ func TestIntercept_CredentialInjectionCanonicalLog(t *testing.T) { setup.Proxy.SetCredentialWithGrant(setup.BackendHost, "Authorization", "Bearer granted-token", "my-grant") - var logged RequestLogData - setup.Proxy.SetLogger(func(data RequestLogData) { - logged = data - }) + waitLog := captureLog(t, setup.Proxy) resp, err := setup.Client.Get(setup.Backend.URL + "/resource") if err != nil { t.Fatalf("request: %v", err) } + io.ReadAll(resp.Body) resp.Body.Close() + logged := waitLog() if !logged.AuthInjected { t.Error("expected AuthInjected=true") } @@ -249,17 +269,16 @@ func TestIntercept_CanonicalLogFields(t *testing.T) { setup.Proxy.SetCredentialWithGrant(setup.BackendHost, "Authorization", "Bearer tok", "test-grant") - var logged RequestLogData - setup.Proxy.SetLogger(func(data RequestLogData) { - logged = data - }) + waitLog := captureLog(t, setup.Proxy) resp, err := setup.Client.Get(setup.Backend.URL + "/some/path") if err != nil { t.Fatalf("request: %v", err) } + io.ReadAll(resp.Body) resp.Body.Close() + logged := waitLog() if logged.Method != "GET" { t.Errorf("Method = %q, want GET", logged.Method) } @@ -625,10 +644,7 @@ func TestIntercept_CaptureHeaders_StrippedBeforeForwarding(t *testing.T) { setup.Proxy.SetCaptureHeaders([]string{"X-Workspace-Slug", "X-Request-Source"}) - var logged RequestLogData - setup.Proxy.SetLogger(func(data RequestLogData) { - logged = data - }) + waitLog := captureLog(t, setup.Proxy) req, _ := http.NewRequest("GET", setup.Backend.URL+"/test", nil) req.Header.Set("X-Workspace-Slug", "sneaky-plum") @@ -641,6 +657,7 @@ func TestIntercept_CaptureHeaders_StrippedBeforeForwarding(t *testing.T) { } defer resp.Body.Close() io.ReadAll(resp.Body) + logged := waitLog() // Verify capture headers were stripped before forwarding. if receivedHeaders.Get("X-Workspace-Slug") != "" { @@ -769,10 +786,7 @@ func TestIntercept_UserID_ContextResolver(t *testing.T) { return nil, false }) - var logged RequestLogData - setup.Proxy.SetLogger(func(data RequestLogData) { - logged = data - }) + waitLog := captureLog(t, setup.Proxy) // Rebuild the client with proxy auth credentials (user:token). proxyURL := mustParseURL(setup.ProxyServer.URL) @@ -793,6 +807,7 @@ func TestIntercept_UserID_ContextResolver(t *testing.T) { } defer resp.Body.Close() io.ReadAll(resp.Body) + logged := waitLog() if logged.UserID != "alice" { t.Errorf("UserID = %q, want %q (CONNECT path)", logged.UserID, "alice") @@ -951,8 +966,7 @@ func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) { }, true }) - var logged RequestLogData - p.SetLogger(func(data RequestLogData) { logged = data }) + waitLog := captureLog(t, p) proxyServer := httptest.NewServer(p) t.Cleanup(proxyServer.Close) @@ -979,6 +993,7 @@ func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) { } io.ReadAll(resp.Body) resp.Body.Close() + logged := waitLog() if resp.StatusCode != http.StatusBadRequest { t.Errorf("status = %d, want 400", resp.StatusCode) diff --git a/proxy/proxy.go b/proxy/proxy.go index 593adb0..2a2c7eb 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -372,6 +372,12 @@ type readCloserWrapper struct { // the captured data and a new ReadCloser that streams the full content. // For small bodies (<=MaxBodySize), the body is fully buffered. // For large bodies, only MaxBodySize is buffered; the rest streams through. +// +// captureBody does a blocking read-ahead, so it must only be used on bodies that +// are already fully available — i.e. request bodies. For response bodies, which +// may be produced incrementally (Server-Sent Events, ndjson, chunked JSON), use +// capturingBody: a read-ahead there would withhold the response from the client +// until MaxBodySize accumulated, starving streamed responses. func captureBody(body io.ReadCloser, contentType string) ([]byte, io.ReadCloser) { if body == nil { return nil, nil @@ -407,6 +413,54 @@ func captureBody(body io.ReadCloser, contentType string) ([]byte, io.ReadCloser) return captured, &readCloserWrapper{Reader: fullBody, Closer: body} } +// capturingBody wraps a response body, copying up to limit bytes into an +// in-memory buffer as the body is read, then invoking onClose (if non-nil) +// exactly once when the body is closed. Capture happens lazily as the downstream +// consumer reads, so — unlike captureBody's read-ahead — it never blocks the +// forwarding of a streamed response. Used for response-body logging on the +// streaming proxy paths. +// +// Not safe for concurrent Read/Close; the proxy paths read the body to +// completion and then close it, which is sequential. +type capturingBody struct { + rc io.ReadCloser + buf bytes.Buffer + limit int + onClose func(captured []byte) + closed bool +} + +func newCapturingBody(rc io.ReadCloser, limit int, onClose func(captured []byte)) *capturingBody { + return &capturingBody{rc: rc, limit: limit, onClose: onClose} +} + +func (c *capturingBody) Read(p []byte) (int, error) { + n, err := c.rc.Read(p) + if n > 0 { + if room := c.limit - c.buf.Len(); room > 0 { + if room > n { + room = n + } + c.buf.Write(p[:room]) + } + } + return n, err +} + +// Captured returns the bytes buffered so far (up to limit). +func (c *capturingBody) Captured() []byte { return c.buf.Bytes() } + +func (c *capturingBody) Close() error { + err := c.rc.Close() + if !c.closed { + c.closed = true + if c.onClose != nil { + c.onClose(c.buf.Bytes()) + } + } + return err +} + // FilterHeaders creates a copy of headers with sensitive values filtered. // injectedHeaders is a set of lower-cased header names whose values should be // redacted (credential headers the proxy injected). @@ -1832,52 +1886,56 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // Forward request resp, err := httpTransport.RoundTrip(outReq) - duration := time.Since(start) - - // Capture response body and headers - var respBody []byte - var respHeaders http.Header - statusCode := http.StatusBadGateway - var responseSize int64 = -1 - if resp != nil { - respHeaders = resp.Header.Clone() - respBody, resp.Body = captureBody(resp.Body, resp.Header.Get("Content-Type")) - statusCode = resp.StatusCode - responseSize = resp.ContentLength - } - p.logRequest(r, RequestLogData{ + // Fields shared by the error and success log lines. + logData := RequestLogData{ Method: r.Method, URL: r.URL.String(), Host: host, Path: r.URL.Path, RequestType: "http", - StatusCode: statusCode, - Duration: duration, - Err: err, RequestHeaders: originalReqHeaders, - ResponseHeaders: respHeaders, RequestBody: reqBody, - ResponseBody: respBody, RequestSize: r.ContentLength, - ResponseSize: responseSize, + ResponseSize: -1, InjectedHeaders: credResult.InjectedHeaders, Grants: credResult.Grants, - }) + } if err != nil { + logData.StatusCode = http.StatusBadGateway + logData.Duration = time.Since(start) + logData.Err = err + p.logRequest(r, logData) http.Error(w, "moat proxy: upstream request failed", http.StatusBadGateway) return } defer resp.Body.Close() + logData.StatusCode = resp.StatusCode + logData.ResponseHeaders = resp.Header.Clone() + logData.ResponseSize = resp.ContentLength + for key, values := range resp.Header { for _, value := range values { w.Header().Add(key, value) } } w.WriteHeader(resp.StatusCode) - _, _ = io.Copy(w, resp.Body) + + // Stream to the client, capturing a bounded sample for logging without a + // blocking read-ahead (which would starve a streamed response). The log line + // is written after the body completes. Non-text bodies carry no useful + // sample, so they stream straight through. + if isTextContentType(resp.Header.Get("Content-Type")) { + cb := newCapturingBody(resp.Body, MaxBodySize, nil) + _, _ = io.Copy(w, cb) + logData.ResponseBody = cb.Captured() + } else { + _, _ = io.Copy(w, resp.Body) + } + logData.Duration = time.Since(start) + p.logRequest(r, logData) } func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) { @@ -2334,9 +2392,6 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req if logURL == "" { logURL = req.URL.String() } - var respBody []byte - respBody, resp.Body = captureBody(resp.Body, resp.Header.Get("Content-Type")) - // Use pre-injection headers so credential values don't appear in logs. preHeaders, _ := req.Context().Value(interceptPreInjHeadersKey{}).(http.Header) if preHeaders == nil { @@ -2344,7 +2399,7 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req } reqBody, _ := req.Context().Value(interceptReqBodyKey{}).([]byte) - p.logRequest(r, RequestLogData{ + logData := RequestLogData{ RequestID: req.Header.Get("X-Request-Id"), Method: req.Method, URL: logURL, @@ -2352,11 +2407,9 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req Path: req.URL.Path, RequestType: "connect", StatusCode: resp.StatusCode, - Duration: time.Since(reqStartFromContext(req.Context())), RequestHeaders: preHeaders, ResponseHeaders: resp.Header.Clone(), RequestBody: reqBody, - ResponseBody: respBody, RequestSize: req.ContentLength, ResponseSize: resp.ContentLength, AuthInjected: len(credResult.InjectedHeaders) > 0, @@ -2364,7 +2417,31 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req Grants: credResult.Grants, Denied: llmDenied, DenyReason: llmDenyReason, - }) + } + reqStart := reqStartFromContext(req.Context()) + + // For text responses (SSE, ndjson, chunked JSON, …) defer the canonical + // log line until the body is read and closed, capturing a bounded sample + // as it streams. A blocking read-ahead here would withhold the status + // line and every chunk from the client until the sample filled, starving + // streamed responses and tripping the client's first-byte timeout. + // + // The deferred line is emitted exactly once: httputil.ReverseProxy + // closes the response body unconditionally (success and copy-error + // paths), so onClose always runs. Responses with no useful streamed + // sample are logged synchronously instead: non-text bodies, and protocol + // upgrades (101), whose body is the hijacked connection rather than a + // readable stream. + if isTextContentType(resp.Header.Get("Content-Type")) && resp.StatusCode != http.StatusSwitchingProtocols { + resp.Body = newCapturingBody(resp.Body, MaxBodySize, func(captured []byte) { + logData.ResponseBody = captured + logData.Duration = time.Since(reqStart) + p.logRequest(r, logData) + }) + } else { + logData.Duration = time.Since(reqStart) + p.logRequest(r, logData) + } return nil }, diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index 6ddb973..f625893 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -19,6 +19,7 @@ import ( "sync" "sync/atomic" "testing" + "time" ) func TestProxy_ForwardsRequests(t *testing.T) { @@ -1408,6 +1409,93 @@ func TestCaptureBody_NilBody(t *testing.T) { } } +// blockingBody delivers one initial chunk, then blocks on Read until released. +// It mimics a streaming upstream that has emitted an early chunk (status line, +// first record, keepalive ping) but not yet MaxBodySize bytes of body — e.g. +// during a long time-to-first-token. +type blockingBody struct { + first []byte + sent bool + release chan struct{} +} + +func (b *blockingBody) Read(p []byte) (int, error) { + if !b.sent { + b.sent = true + return copy(p, b.first), nil + } + <-b.release + return 0, io.EOF +} + +func (b *blockingBody) Close() error { return nil } + +// TestCapturingBody_StreamsAndCaptures verifies the full body passes through +// while a bounded sample is captured, and onClose fires exactly once with it. +func TestCapturingBody_StreamsAndCaptures(t *testing.T) { + full := strings.Repeat("a", MaxBodySize+500) + var got []byte + var calls int + cb := newCapturingBody(io.NopCloser(strings.NewReader(full)), MaxBodySize, func(c []byte) { + calls++ + got = c + }) + + streamed, err := io.ReadAll(cb) + if err != nil { + t.Fatalf("read: %v", err) + } + if string(streamed) != full { + t.Errorf("streamed %d bytes, want %d (body must pass through in full)", len(streamed), len(full)) + } + + if err := cb.Close(); err != nil { + t.Fatalf("close: %v", err) + } + if err := cb.Close(); err != nil { // second close must not re-fire onClose + t.Fatalf("second close: %v", err) + } + if calls != 1 { + t.Errorf("onClose called %d times, want 1", calls) + } + if len(got) != MaxBodySize { + t.Errorf("captured %d bytes, want %d (bounded sample)", len(got), MaxBodySize) + } +} + +// TestCapturingBody_NeverBlocksOnSlowStream is the regression guard for the +// first-byte-timeout bug: a streamed body that delivers one record then blocks +// must not trigger a read-ahead. capturingBody reads only what the consumer +// requests, so the first record is available immediately — regardless of content +// type. This uses a non-SSE stream (application/x-ndjson) to show the fix is not +// keyed on a media-type allowlist. +func TestCapturingBody_NeverBlocksOnSlowStream(t *testing.T) { + first := []byte(`{"chunk":1}` + "\n") + body := &blockingBody{first: first, release: make(chan struct{})} + defer close(body.release) + cb := newCapturingBody(body, MaxBodySize, nil) + + done := make(chan []byte, 1) + go func() { + buf := make([]byte, 64) + n, _ := cb.Read(buf) + done <- buf[:n] + }() + + select { + case got := <-done: + if !bytes.Equal(got, first) { + t.Errorf("first read = %q, want %q", got, first) + } + case <-time.After(2 * time.Second): + t.Fatal("capturingBody.Read blocked instead of returning the available record") + } + + if !bytes.Equal(cb.Captured(), first) { + t.Errorf("captured = %q, want %q", cb.Captured(), first) + } +} + // TestFilterHeaders_RedactsInjectedCredential verifies credential redaction. func TestFilterHeaders_RedactsInjectedCredential(t *testing.T) { headers := http.Header{