From 35c9e3eb2831cdf9e543dd3d0546f2d8ebda6770 Mon Sep 17 00:00:00 2001 From: Andrii Bezzub Date: Tue, 23 Jun 2026 15:58:33 +0200 Subject: [PATCH 1/2] fix(proxy): capture streamed responses with a non-blocking tee MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The response-body log sampler (captureBody) did a blocking io.ReadFull of up to MaxBodySize (8 KB) inside the forward path, before the response reached the client. For any incrementally-produced text response — Server-Sent Events, application/x-ndjson, chunked JSON — this withheld the status line and every chunk (including upstream keepalive pings sent during a long time-to-first-token) until 8 KB accumulated, or until the whole stream ended for sub-8 KB responses. A streaming client received nothing, tripped its first-byte timeout, and retried in a loop. Direct (un-proxied) clients were unaffected because events flowed incrementally. Add capturingBody, which tees up to MaxBodySize into a buffer as the downstream consumer reads, so capture never blocks the forward path and covers every text content type (the same isTextContentType gate as before, now non-blocking). Both response paths use it. captureBody's blocking read is retained for request bodies, which are already fully available. On the intercept path the copy is owned by httputil.ReverseProxy, so the canonical log line for a text response is emitted from the body's Close — carrying the streamed sample and full duration. ReverseProxy closes the body unconditionally, so it fires exactly once; protocol upgrades (101) and non-text responses are logged synchronously instead. The plain-HTTP path captures during its own copy and logs synchronously after it. This does not change the llm-gateway response-policy path, which still reads the full response (io.ReadAll) to evaluate it before forwarding. Tests: capturingBody never reads ahead (regression guard using a non-SSE ndjson stream) and captures a bounded sample while forwarding the full body; intercept tests that assert on the canonical log line wait for the deferred entry. --- proxy/intercept_test.go | 51 +++++++++------ proxy/proxy.go | 135 +++++++++++++++++++++++++++++++--------- proxy/proxy_test.go | 88 ++++++++++++++++++++++++++ 3 files changed, 227 insertions(+), 47 deletions(-) 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{ From b72f19f55e45263a502aabec24176fcacc4e7723 Mon Sep 17 00:00:00 2001 From: Andrii Bezzub Date: Tue, 23 Jun 2026 15:58:33 +0200 Subject: [PATCH 2/2] docs(changelog): add v0.14.1 entry for the streamed-response capture fix (#37) --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) 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