diff --git a/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go b/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go index 087cd0a7ef..7151e2c852 100644 --- a/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go +++ b/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go @@ -23,7 +23,7 @@ import ( "compress/gzip" "fmt" "io" - "runtime" + "sync" "github.com/andybalholm/brotli" ) @@ -47,129 +47,274 @@ func decompressBody(body []byte, encoding string) ([]byte, error) { } } -// streamDecompressor provides true per-chunk streaming decompression using an io.Pipe -// and a persistent decoder goroutine. Incoming compressed chunks are written to the -// pipe; the goroutine owns the stateful gzip/brotli reader and pushes decompressed -// bytes to an output channel as complete blocks become available. +// streamDecompressor provides true per-chunk streaming decompression. It adapts +// Envoy's push-style chunk delivery to Go's pull-style gzip/brotli reader: a single +// decoder goroutine owns a stateful reader that pulls compressed bytes from an +// internal buffer (fed by FeedChunk) and writes decompressed bytes to an unbounded +// output buffer. +// +// The synchronization is a single sync.Cond guarding shared state. Two invariants make +// it deadlock-free and race-free, unlike the previous io.Pipe + bounded-channel design: +// +// 1. The decoder goroutine NEVER blocks on output — decompressed bytes go into a +// bytes.Buffer that FeedChunk drains on every call. (The old design used a bounded +// channel; when it filled, the goroutine blocked mid-decode, stopped consuming +// input, and the whole stream wedged. That was the mid-stream stall.) A single +// decode burst is capped at maxStreamAccumulatorSize so a decompression bomb fails +// terminally instead of exhausting memory. +// 2. FeedChunk returns exactly when the decoder has consumed all input fed so far and +// is blocked waiting for more (feederBlocked), or has finished/errored (done). At +// that point every byte decodable from the fed input is already in the output +// buffer — so output is complete per chunk, with no racy "maybe next time" drain. type streamDecompressor struct { - pipeWriter *io.PipeWriter - outChan chan []byte - errChan chan error + mu sync.Mutex + cond *sync.Cond + + inbuf []byte // compressed bytes fed but not yet consumed by the decoder + out bytes.Buffer // decompressed bytes not yet returned to the caller (unbounded) + closed bool // endOfStream fed: feeder returns io.EOF once inbuf drains + feederBlocked bool // decoder is parked waiting for more input (inbuf empty) + done bool // decoder goroutine has exited + decodeErr error // terminal decode error (nil on clean io.EOF) + passthrough bool // unknown encoding: FeedChunk returns bytes unchanged } // newStreamDecompressor starts the background decoder goroutine and returns a // streamDecompressor ready to accept chunks via FeedChunk. func newStreamDecompressor(encoding string) *streamDecompressor { - pr, pw := io.Pipe() - outChan := make(chan []byte, 64) - errChan := make(chan error, 1) - - go func() { - defer close(outChan) - var r io.Reader - switch encoding { - case "gzip": - gr, err := gzip.NewReader(pr) - if err != nil { - select { - case errChan <- fmt.Errorf("gzip.NewReader: %w", err): - default: - } - _ = pr.CloseWithError(err) - return - } - defer gr.Close() - r = gr - case "br": - r = brotli.NewReader(pr) - default: - r = pr + sd := &streamDecompressor{} + sd.cond = sync.NewCond(&sd.mu) + + if encoding != "gzip" && encoding != "br" { + sd.passthrough = true + return sd + } + + go sd.decodeLoop(encoding) + return sd +} + +// feederReader is the io.Reader handed to the gzip/brotli decoder. Each Read hands the +// decoder whatever compressed bytes are currently buffered; when the buffer is empty it +// parks (recording feederBlocked=true so FeedChunk knows all fed input is consumed) until +// FeedChunk supplies more or signals end-of-stream. +type feederReader struct{ sd *streamDecompressor } + +func (f feederReader) Read(p []byte) (int, error) { + sd := f.sd + sd.mu.Lock() + defer sd.mu.Unlock() + for len(sd.inbuf) == 0 { + if sd.closed { + return 0, io.EOF } + sd.feederBlocked = true + sd.cond.Broadcast() // wake any FeedChunk waiting for "all input consumed" + sd.cond.Wait() + } + sd.feederBlocked = false + n := copy(p, sd.inbuf) + sd.inbuf = sd.inbuf[n:] + return n, nil +} - buf := make([]byte, 32*1024) - for { - n, err := r.Read(buf) - if n > 0 { - out := make([]byte, n) - copy(out, buf[:n]) - outChan <- out - } - if err == io.EOF { - return - } - if err != nil { - select { - case errChan <- err: - default: - } +// decodeLoop owns the stateful decoder for the life of the stream and copies its output +// into the unbounded buffer. It exits on io.EOF (clean) or any decode error. +func (sd *streamDecompressor) decodeLoop(encoding string) { + var r io.Reader + switch encoding { + case "gzip": + gr, err := gzip.NewReader(feederReader{sd: sd}) + if err != nil { + sd.finish(fmt.Errorf("gzip.NewReader: %w", err)) + return + } + defer gr.Close() + r = gr + case "br": + r = brotli.NewReader(feederReader{sd: sd}) + } + + buf := make([]byte, 32*1024) + for { + n, err := r.Read(buf) + if n > 0 { + sd.mu.Lock() + // Bound the decompressed output: a small compressed chunk can expand to an + // arbitrarily large decode burst (decompression bomb). Since FeedChunk drains + // sd.out on each call, this caps a single burst — matching the force-flush + // guard the uncompressed streaming path applies at maxStreamAccumulatorSize. + // On overflow we stop before copying the offending block and fail terminally. + if sd.out.Len()+n > maxStreamAccumulatorSize { + sd.mu.Unlock() + sd.finish(fmt.Errorf("decompressed stream exceeds maximum allowed size (%d bytes)", maxStreamAccumulatorSize)) return } + sd.out.Write(buf[:n]) // bytes.Buffer copies, so reusing buf is safe + sd.mu.Unlock() + } + if err == io.EOF { + sd.finish(nil) + return + } + if err != nil { + sd.finish(err) + return } - }() + } +} - return &streamDecompressor{pipeWriter: pw, outChan: outChan, errChan: errChan} +// finish records the decoder's terminal state and wakes any waiting FeedChunk. +func (sd *streamDecompressor) finish(err error) { + sd.mu.Lock() + sd.decodeErr = err + sd.done = true + sd.cond.Broadcast() + sd.mu.Unlock() } -// FeedChunk writes compressed bytes into the decoder and returns whatever decompressed -// bytes are immediately available. +// FeedChunk feeds compressed bytes to the decoder and returns every decompressed byte +// derivable from the input fed so far. Output may legitimately be empty for an +// intermediate chunk when the decoder needs more input to complete a block. // -// For intermediate chunks the output may be empty — the decoder needs more input before -// a full DEFLATE/brotli block can be decoded. Callers must tolerate empty output. -// -// On endOfStream=true the pipe writer is closed and FeedChunk blocks until all remaining -// decompressed bytes have been flushed by the goroutine. +// It blocks only until the decoder has consumed all fed input and parked for more (or +// finished) — bounded by the CPU cost of decoding this chunk, never on downstream +// delivery — so it cannot wedge the ext_proc stream. func (sd *streamDecompressor) FeedChunk(chunk []byte, endOfStream bool) ([]byte, error) { - if len(chunk) > 0 { - if _, err := sd.pipeWriter.Write(chunk); err != nil { - return nil, fmt.Errorf("stream decompressor write: %w", err) - } + if sd.passthrough { + return chunk, nil } + sd.mu.Lock() + defer sd.mu.Unlock() + + if len(chunk) > 0 { + sd.inbuf = append(sd.inbuf, chunk...) + } if endOfStream { - _ = sd.pipeWriter.Close() - var result []byte - for data := range sd.outChan { - result = append(result, data...) - } - select { - case err := <-sd.errChan: - if err != nil { - return result, err - } - default: - } - return result, nil + sd.closed = true } + sd.feederBlocked = false + sd.cond.Broadcast() // wake the feeder to consume the new input / observe close - // pw.Write blocks until the goroutine's pr.Read has consumed all our bytes. - // Yield so the goroutine can finish its r.Read call and push decoded output to - // outChan before we do the non-blocking drain below. - runtime.Gosched() + // Wait until the decoder has drained all fed input and parked (feederBlocked with an + // empty inbuf), or has exited. Either way, all output decodable from the fed input is + // now in sd.out. + for !sd.done && !(sd.feederBlocked && len(sd.inbuf) == 0) { + sd.cond.Wait() + } - var result []byte - for { - select { - case data, ok := <-sd.outChan: - if !ok { - return result, nil - } - result = append(result, data...) - default: - return result, nil - } + result := make([]byte, sd.out.Len()) + copy(result, sd.out.Bytes()) + sd.out.Reset() + + if sd.decodeErr != nil && sd.decodeErr != io.EOF { + return result, sd.decodeErr } + return result, nil } -// Close releases the decompressor's pipe and drains the goroutine. Call this on -// error paths where endOfStream will never arrive. +// Close releases the decoder goroutine on error paths where endOfStream will never +// arrive. It signals end-of-input and returns without joining, so it never hangs; the +// goroutine observes the close, drains, and exits on its own. Safe to call multiple times. func (sd *streamDecompressor) Close() { - _ = sd.pipeWriter.CloseWithError(io.ErrClosedPipe) - for range sd.outChan { + if sd.passthrough { + return + } + sd.mu.Lock() + sd.closed = true + sd.cond.Broadcast() + sd.mu.Unlock() +} + +// streamWriteFlushCloser is the common interface implemented by *gzip.Writer and +// *brotli.Writer: incremental writes, a mid-stream flush that keeps the stream open, +// and a final close that emits the trailer. +type streamWriteFlushCloser interface { + io.WriteCloser + Flush() error +} + +// streamCompressor provides stateful, per-chunk streaming compression that produces +// a SINGLE continuous compressed stream across all chunks. +// +// This is the compression counterpart of streamDecompressor and must be used instead +// of recompressBody for streaming bodies. recompressBody creates a fresh writer and +// Close()es it on every call, so each chunk becomes an independent, self-contained +// gzip/brotli member. A concatenation of independent members is not what the +// Content-Encoding header promises (one stream for the whole body), and downstream +// HTTP decoders (e.g. the Anthropic/Claude Code client) decode only the first member, +// see its end-of-stream trailer, and treat the response as finished — dropping every +// subsequent chunk. streamCompressor instead keeps one writer alive for the life of +// the stream, flushing (Z_SYNC_FLUSH) after each chunk and closing only at EndOfStream. +type streamCompressor struct { + buf *bytes.Buffer + writer streamWriteFlushCloser + encoding string + closed bool +} + +// newStreamCompressor returns a streamCompressor for the given Content-Encoding. +// For unknown encodings the writer is nil and FeedChunk passes bytes through unchanged. +func newStreamCompressor(encoding string) *streamCompressor { + buf := &bytes.Buffer{} + switch encoding { + case "gzip": + return &streamCompressor{buf: buf, writer: gzip.NewWriter(buf), encoding: encoding} + case "br": + return &streamCompressor{buf: buf, writer: brotli.NewWriter(buf), encoding: encoding} + default: + return &streamCompressor{buf: buf, encoding: encoding} + } +} + +// FeedChunk compresses a chunk of the decompressed body and returns the compressed +// bytes produced so far. On intermediate chunks the writer is flushed (keeping the +// stream open); on endOfStream the writer is closed, emitting the final block and +// trailer. The returned bytes belong to the caller (a fresh copy). +func (sc *streamCompressor) FeedChunk(chunk []byte, endOfStream bool) ([]byte, error) { + // Passthrough for unknown encodings. + if sc.writer == nil { + return chunk, nil + } + if sc.closed { + return nil, fmt.Errorf("stream compressor: write after close") + } + if len(chunk) > 0 { + if _, err := sc.writer.Write(chunk); err != nil { + return nil, fmt.Errorf("stream compressor write: %w", err) + } + } + if endOfStream { + if err := sc.writer.Close(); err != nil { + return nil, fmt.Errorf("stream compressor close: %w", err) + } + sc.closed = true + } else { + if err := sc.writer.Flush(); err != nil { + return nil, fmt.Errorf("stream compressor flush: %w", err) + } + } + out := make([]byte, sc.buf.Len()) + copy(out, sc.buf.Bytes()) + sc.buf.Reset() + return out, nil +} + +// Close releases the compressor's writer on error paths where endOfStream will never +// arrive. Safe to call multiple times. +func (sc *streamCompressor) Close() { + if sc.writer != nil && !sc.closed { + _ = sc.writer.Close() + sc.closed = true } } // recompressBody re-compresses body bytes using the original Content-Encoding. // Used to restore compression after policies have processed the decompressed body. +// NOTE: This produces a complete standalone stream and is only correct for +// non-streaming (fully buffered) bodies. For streaming bodies use streamCompressor, +// which keeps a single stream open across chunks. // Supported encodings: "gzip", "br" (Brotli). Unknown encodings are returned as-is. func recompressBody(body []byte, encoding string) ([]byte, error) { switch encoding { diff --git a/gateway/gateway-runtime/policy-engine/internal/kernel/decompression_test.go b/gateway/gateway-runtime/policy-engine/internal/kernel/decompression_test.go index 61d9a95f2b..390963e82c 100644 --- a/gateway/gateway-runtime/policy-engine/internal/kernel/decompression_test.go +++ b/gateway/gateway-runtime/policy-engine/internal/kernel/decompression_test.go @@ -21,6 +21,7 @@ package kernel import ( "bytes" "compress/gzip" + "io" "testing" "time" @@ -258,6 +259,53 @@ func TestStreamDecompressor_Close_DoesNotHang(t *testing.T) { } } +// TestStreamDecompressor_Close_ReleasesParkedDecoder is the regression test for the +// goroutine leak that occurred when a policy terminated a compressed response stream +// mid-flight: the decoder had consumed a non-final chunk and parked waiting for input +// that never arrived (EndOfStream never reached it). Close() must release that parked +// goroutine — otherwise every early-terminated compressed stream leaks one goroutine. +func TestStreamDecompressor_Close_ReleasesParkedDecoder(t *testing.T) { + // Build a gzip stream flushed after the first event so a leading, non-terminal + // slice is independently decodable — mimicking a mid-stream chunk. + var buf bytes.Buffer + zw := gzip.NewWriter(&buf) + _, err := zw.Write([]byte("event: message_start\n\n")) + require.NoError(t, err) + require.NoError(t, zw.Flush()) + midStreamBytes := append([]byte{}, buf.Bytes()...) + + sd := newStreamDecompressor("gzip") + + // Feed the mid-stream chunk WITHOUT EndOfStream. FeedChunk returns once the + // decoder has consumed it and parked waiting for more input. + _, err = sd.FeedChunk(midStreamBytes, false) + require.NoError(t, err) + + // Sanity-check the decoder is parked (not done) — this is the leak state. + sd.mu.Lock() + require.True(t, sd.feederBlocked, "decoder should be parked waiting for more input") + require.False(t, sd.done, "decoder should not have exited before Close") + sd.mu.Unlock() + + // The stream is terminated early; Close() must release the parked goroutine. + sd.Close() + + deadline := time.After(2 * time.Second) + for { + sd.mu.Lock() + done := sd.done + sd.mu.Unlock() + if done { + return // goroutine exited — no leak + } + select { + case <-deadline: + t.Fatal("decoder goroutine was not released after Close() — leak on mid-stream termination") + case <-time.After(5 * time.Millisecond): + } + } +} + // TestStreamDecompressor_RoundTrip verifies the full cycle the streaming path // performs: incoming compressed bytes → decompress → policy modifies → recompress. // The final recompressed bytes must decompress back to the (modified) original. @@ -280,3 +328,278 @@ func TestStreamDecompressor_RoundTrip(t *testing.T) { require.NoError(t, err) assert.Equal(t, original, final) } + +// TestStreamDecompressor_ManyChunks_NoStall is the regression test for the +// mid-stream stall that broke large Claude Code tasks. The previous io.Pipe + +// bounded-channel (cap 64) design deadlocked once the decoder produced more +// undrained output blocks than the channel could hold while a chunk was still +// being consumed: the decoder goroutine blocked on the full channel, stopped +// reading the pipe, so the pipe Write never returned, so FeedChunk never drained +// the channel. It surfaced only on long/large responses (many chunks) — short +// ones never filled the channel, which is why the two-chunk tests missed it. +// +// This feeds a long stream as many separate sync-flushed chunks and asserts every +// byte comes back, within a deadline so a regression fails as a timeout, not a hang. +func TestStreamDecompressor_ManyChunks_NoStall(t *testing.T) { + // Build a gzip stream flushed after each of many logical chunks, mirroring how + // an upstream (Anthropic) emits a continuous gzip stream with per-event flushes. + const numChunks = 500 + var raw bytes.Buffer + var compressed bytes.Buffer + zw := gzip.NewWriter(&compressed) + flushOffsets := make([]int, 0, numChunks) + for i := 0; i < numChunks; i++ { + // ~250 bytes per SSE event, including ~200 bytes of filler text. + filler := string(bytes.Repeat([]byte("x"), 200)) + line := []byte(`data: {"type":"content_block_delta","index":0,` + + `"delta":{"type":"text_delta","text":"token-` + filler + `"}}` + "\n\n") + raw.Write(line) + _, err := zw.Write(line) + require.NoError(t, err) + require.NoError(t, zw.Flush()) // sync-flush → a decodable boundary, like SSE + flushOffsets = append(flushOffsets, compressed.Len()) + } + require.NoError(t, zw.Close()) + + sd := newStreamDecompressor("gzip") + compBytes := compressed.Bytes() + // Close() appended the gzip trailer after the last recorded flush offset — extend + // the final segment to the true end so the EOS chunk carries the trailer. + flushOffsets[len(flushOffsets)-1] = len(compBytes) + + // Feed the compressed stream in slices aligned to the flush boundaries, each as a + // separate FeedChunk call — the shape that accumulated backlog in the old design. + done := make(chan []byte, 1) + errCh := make(chan error, 1) + go func() { + var out []byte + prev := 0 + for i, off := range flushOffsets { + eos := i == len(flushOffsets)-1 + part, err := sd.FeedChunk(compBytes[prev:off], eos) + if err != nil { + errCh <- err + return + } + out = append(out, part...) + prev = off + } + done <- out + }() + + select { + case err := <-errCh: + t.Fatalf("FeedChunk returned error: %v", err) + case out := <-done: + assert.Equal(t, raw.Bytes(), out, "every decompressed byte must be returned across all chunks") + case <-time.After(15 * time.Second): + sd.Close() + t.Fatal("streamDecompressor stalled on a long many-chunk stream (regression: mid-stream deadlock)") + } +} + +// TestStreamDecompressor_HighRatioChunk_NoStall feeds a single chunk that +// decompresses to far more than the old 64-block channel could buffer, exercising +// the second face of the same bug: one FeedChunk whose decoded output alone would +// have overrun the bounded channel mid-consume. With the unbounded output buffer +// this must complete promptly and return all bytes. +func TestStreamDecompressor_HighRatioChunk_NoStall(t *testing.T) { + // 8 MiB of highly compressible data → tiny compressed input, ~256 blocks of + // 32 KiB out — well past the old cap of 64. + original := bytes.Repeat([]byte("A"), 8<<20) + compressed := gzipCompress(original) + + sd := newStreamDecompressor("gzip") + + done := make(chan []byte, 1) + errCh := make(chan error, 1) + go func() { + out, err := sd.FeedChunk(compressed, true) + if err != nil { + errCh <- err + return + } + done <- out + }() + + select { + case err := <-errCh: + t.Fatalf("FeedChunk returned error: %v", err) + case out := <-done: + assert.Equal(t, len(original), len(out)) + assert.True(t, bytes.Equal(original, out)) + case <-time.After(15 * time.Second): + sd.Close() + t.Fatal("streamDecompressor stalled on a high-ratio single chunk (regression: bounded-channel overrun)") + } +} + +// TestStreamDecompressor_OutputCap_FailsTerminally verifies the decompression-bomb +// guard: a tiny compressed input that expands past maxStreamAccumulatorSize must fail +// terminally rather than growing the output buffer without bound. +func TestStreamDecompressor_OutputCap_FailsTerminally(t *testing.T) { + // Compresses to a few KB but decompresses to just over the 10 MB cap. + original := bytes.Repeat([]byte("A"), maxStreamAccumulatorSize+1<<20) + compressed := gzipCompress(original) + + sd := newStreamDecompressor("gzip") + + done := make(chan error, 1) + go func() { + _, err := sd.FeedChunk(compressed, true) + done <- err + }() + + select { + case err := <-done: + require.Error(t, err, "expected a terminal error when decompressed output exceeds the cap") + assert.Contains(t, err.Error(), "maximum allowed size") + case <-time.After(15 * time.Second): + sd.Close() + t.Fatal("streamDecompressor hung instead of failing terminally on an oversized decode burst") + } +} + +// ============================================================================= +// streamCompressor Tests +// ============================================================================= + +// gzipHeaderCount counts the number of gzip member headers (magic bytes 1f 8b 08) +// in a byte slice. A correct single continuous stream has exactly one. +func gzipHeaderCount(data []byte) int { + return bytes.Count(data, []byte{0x1f, 0x8b, 0x08}) +} + +// singleMemberGunzip decodes only the FIRST gzip member, mimicking downstream HTTP +// decoders (such as the Claude Code client) that stop at the first member's trailer. +func singleMemberGunzip(t *testing.T, data []byte) []byte { + t.Helper() + zr, err := gzip.NewReader(bytes.NewReader(data)) + require.NoError(t, err) + zr.Multistream(false) // do NOT transparently continue into later members + out, err := io.ReadAll(zr) + require.NoError(t, err) + return out +} + +// TestStreamCompressor_Gzip_MultipleChunks is the regression test for the analytics +// streaming bug: multiple chunks must be re-compressed into ONE continuous gzip stream. +func TestStreamCompressor_Gzip_MultipleChunks(t *testing.T) { + chunks := [][]byte{ + []byte("event: message_start\ndata: {\"type\":\"message_start\"}\n\n"), + []byte("event: content_block_delta\ndata: {\"delta\":\"Hi\"}\n\n"), + []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + } + var full []byte + for _, c := range chunks { + full = append(full, c...) + } + + sc := newStreamCompressor("gzip") + var compressed []byte + for i, c := range chunks { + endOfStream := i == len(chunks)-1 + out, err := sc.FeedChunk(c, endOfStream) + require.NoError(t, err) + compressed = append(compressed, out...) + } + + // The whole body must be a SINGLE gzip member — this is the property the bug violated. + assert.Equal(t, 1, gzipHeaderCount(compressed), + "streaming compression must emit exactly one gzip member, not one per chunk") + + // A single-member decoder (like the downstream client) must recover ALL chunks, + // not just the first. This is what failed for Claude Code before the fix. + assert.Equal(t, full, singleMemberGunzip(t, compressed)) + + // And a normal (multistream) decode must also yield the full body. + final, err := decompressBody(compressed, "gzip") + require.NoError(t, err) + assert.Equal(t, full, final) +} + +// TestRecompressBody_PerChunk_ProducesMultipleMembers documents the OLD, buggy +// behaviour that streamCompressor replaces: re-compressing each chunk independently +// produces N gzip members, and a single-member decoder recovers only the first chunk. +func TestRecompressBody_PerChunk_ProducesMultipleMembers(t *testing.T) { + chunks := [][]byte{ + []byte("event: message_start\n\n"), + []byte("event: content_block_delta\n\n"), + []byte("event: message_stop\n\n"), + } + + var buggy []byte + for _, c := range chunks { + out, err := recompressBody(c, "gzip") + require.NoError(t, err) + buggy = append(buggy, out...) + } + + // Per-chunk recompress yields one member per chunk... + assert.Equal(t, len(chunks), gzipHeaderCount(buggy)) + // ...and a single-member decoder sees only the first chunk — the dropped-stream bug. + assert.Equal(t, chunks[0], singleMemberGunzip(t, buggy)) +} + +// TestStreamCompressor_Brotli_MultipleChunks verifies the brotli path round-trips +// across multiple chunks into a single continuous stream. +func TestStreamCompressor_Brotli_MultipleChunks(t *testing.T) { + chunks := [][]byte{ + []byte("first-chunk-payload"), + []byte("second-chunk-payload"), + []byte("third-chunk-payload"), + } + var full []byte + for _, c := range chunks { + full = append(full, c...) + } + + sc := newStreamCompressor("br") + var compressed []byte + for i, c := range chunks { + out, err := sc.FeedChunk(c, i == len(chunks)-1) + require.NoError(t, err) + compressed = append(compressed, out...) + } + + final, err := decompressBody(compressed, "br") + require.NoError(t, err) + assert.Equal(t, full, final) +} + +// TestStreamCompressor_Gzip_EmptyFinalChunk mirrors the real Envoy flow where the +// EndOfStream chunk carries zero bytes: the trailer must still be emitted so the +// stream is well-formed. +func TestStreamCompressor_Gzip_EmptyFinalChunk(t *testing.T) { + sc := newStreamCompressor("gzip") + + out1, err := sc.FeedChunk([]byte("payload-one"), false) + require.NoError(t, err) + out2, err := sc.FeedChunk([]byte("payload-two"), false) + require.NoError(t, err) + // Final chunk with no data, only EndOfStream — as Envoy delivers it. + out3, err := sc.FeedChunk(nil, true) + require.NoError(t, err) + + compressed := append(append(append([]byte{}, out1...), out2...), out3...) + assert.Equal(t, 1, gzipHeaderCount(compressed)) + assert.Equal(t, []byte("payload-onepayload-two"), singleMemberGunzip(t, compressed)) +} + +// TestStreamCompressor_UnknownEncoding_Passthrough verifies unknown encodings pass +// bytes through unchanged. +func TestStreamCompressor_UnknownEncoding_Passthrough(t *testing.T) { + sc := newStreamCompressor("identity") + out, err := sc.FeedChunk([]byte("raw-bytes"), true) + require.NoError(t, err) + assert.Equal(t, []byte("raw-bytes"), out) +} + +// TestStreamCompressor_WriteAfterClose returns an error rather than corrupting output. +func TestStreamCompressor_WriteAfterClose(t *testing.T) { + sc := newStreamCompressor("gzip") + _, err := sc.FeedChunk([]byte("data"), true) + require.NoError(t, err) + _, err = sc.FeedChunk([]byte("more"), false) + require.Error(t, err) +} diff --git a/gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go b/gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go index 3d11d38f81..bb562c2b1d 100644 --- a/gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go +++ b/gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go @@ -124,6 +124,10 @@ type PolicyExecutionContext struct { // requestStreamDecomp performs per-chunk decompression for compressed streaming // request bodies. Nil when the request is not Content-Encoded. requestStreamDecomp *streamDecompressor + // requestStreamComp performs per-chunk re-compression for compressed streaming + // request bodies, keeping a single continuous stream open across chunks. Nil until + // the first compressed chunk is forwarded. + requestStreamComp *streamCompressor // isStreamingResponse is set to true during response headers processing when // streaming indicators are detected AND the policy chain supports streaming. @@ -133,6 +137,10 @@ type PolicyExecutionContext struct { // responseStreamDecomp performs per-chunk decompression for compressed streaming // response bodies. Nil when the response is not Content-Encoded. responseStreamDecomp *streamDecompressor + // responseStreamComp performs per-chunk re-compression for compressed streaming + // response bodies, keeping a single continuous stream open across chunks. Nil until + // the first compressed chunk is forwarded. + responseStreamComp *streamCompressor // streamTerminated is set when a policy returns TerminateStream=true. Any // subsequent upstream chunks that Envoy delivers after we have already sent // EndOfStream downstream are silently suppressed — the downstream connection @@ -853,6 +861,17 @@ func (ec *PolicyExecutionContext) processStreamingResponseBody( if execResult.StreamTerminated { ec.streamTerminated = true } + + // Release the decompressor goroutine on the terminal chunk. On a normal + // EndOfStream the decoder already exited on its own after the feeder returned + // io.EOF; but when a policy terminates the stream early, EndOfStream never + // reaches the decoder, so it stays parked waiting for input that will never + // arrive — a goroutine leak. Close() releases it and is idempotent/safe on an + // already-finished decoder. + if ec.responseStreamDecomp != nil && (chunk.EndOfStream || execResult.StreamTerminated) { + ec.responseStreamDecomp.Close() + ec.responseStreamDecomp = nil + } return TranslateStreamingResponseChunkAction(execResult, chunk, ec) } diff --git a/gateway/gateway-runtime/policy-engine/internal/kernel/translator.go b/gateway/gateway-runtime/policy-engine/internal/kernel/translator.go index e6dd5d3976..4ee78cf69a 100644 --- a/gateway/gateway-runtime/policy-engine/internal/kernel/translator.go +++ b/gateway/gateway-runtime/policy-engine/internal/kernel/translator.go @@ -1514,17 +1514,31 @@ func TranslateStreamingRequestChunkAction(result *executor.StreamingRequestExecu // Re-compress the output if the original request was Content-Encoded. // The upstream receives the Content-Encoding header as-is, so the body // bytes must match the encoding the upstream expects. + // + // Use a stateful streamCompressor (kept for the life of the stream), NOT + // recompressBody: per-chunk recompressBody emits an independent closed member per + // chunk, so the upstream decoder would stop after the first one. See + // TranslateStreamingResponseChunkAction for the full rationale. if execCtx.requestContentEncoding != "" { - recompressed, err := recompressBody(outputBody, execCtx.requestContentEncoding) + if execCtx.requestStreamComp == nil { + execCtx.requestStreamComp = newStreamCompressor(execCtx.requestContentEncoding) + } + recompressed, err := execCtx.requestStreamComp.FeedChunk(outputBody, originalChunk.EndOfStream) if err != nil { - slog.Warn("[streaming] failed to re-compress request body; sending uncompressed — Content-Encoding mismatch", + // The Content-Encoding header has already been forwarded to the upstream in + // streaming mode and cannot be retracted. Sending the body uncompressed under + // a committed Content-Encoding would yield a stream the upstream cannot decode. + // Fail the ext_proc stream so Envoy aborts the request rather than delivering + // corrupt bytes. The forwarded Content-Encoding state is left intact. + slog.Error("[streaming] failed to re-compress request body; terminating stream to avoid Content-Encoding mismatch", "encoding", execCtx.requestContentEncoding, "error", err, ) - execCtx.requestContentEncoding = "" - } else { - outputBody = recompressed + execCtx.requestStreamComp.Close() + return nil, fmt.Errorf("failed to re-compress streaming request body (encoding %q): %w", + execCtx.requestContentEncoding, err) } + outputBody = recompressed } analyticsData := make(map[string]any) @@ -1585,21 +1599,40 @@ func TranslateStreamingResponseChunkAction(result *executor.StreamingResponseExe outputBody = originalChunk.Chunk } + // If a policy terminated the stream early (e.g. guardrail intervention), force + // EndOfStream so Envoy closes the connection cleanly after delivering the final chunk. + endOfStream := originalChunk.EndOfStream || result.StreamTerminated + // Re-compress the output if the original response was Content-Encoded. // Response headers (including Content-Encoding) are already committed downstream // in streaming mode and cannot be changed — the body must match the encoding // the client expects. + // + // Use a stateful streamCompressor (created once and kept for the life of the + // stream), NOT recompressBody: recompressBody would emit an independent, fully + // closed compressed member per chunk, and downstream clients decode only the first + // member and treat the stream as finished — dropping every subsequent chunk. if execCtx.responseContentEncoding != "" { - recompressed, err := recompressBody(outputBody, execCtx.responseContentEncoding) + if execCtx.responseStreamComp == nil { + execCtx.responseStreamComp = newStreamCompressor(execCtx.responseContentEncoding) + } + recompressed, err := execCtx.responseStreamComp.FeedChunk(outputBody, endOfStream) if err != nil { - slog.Warn("[streaming] failed to re-compress response body; sending uncompressed — Content-Encoding mismatch", + // The Content-Encoding header has already been committed to the downstream + // client in streaming mode and cannot be changed. Sending the body + // uncompressed under a committed Content-Encoding would yield a stream the + // client cannot decode. Fail the ext_proc stream so Envoy aborts the response + // rather than delivering corrupt bytes. The committed Content-Encoding state + // is left intact. + slog.Error("[streaming] failed to re-compress response body; terminating stream to avoid Content-Encoding mismatch", "encoding", execCtx.responseContentEncoding, "error", err, ) - execCtx.responseContentEncoding = "" - } else { - outputBody = recompressed + execCtx.responseStreamComp.Close() + return nil, fmt.Errorf("failed to re-compress streaming response body (encoding %q): %w", + execCtx.responseContentEncoding, err) } + outputBody = recompressed } analyticsData := make(map[string]any) @@ -1634,9 +1667,6 @@ func TranslateStreamingResponseChunkAction(result *executor.StreamingResponseExe mergeDynamicMetadata(execCtx.dynamicMetadata, dm) } - // If a policy terminated the stream early (e.g. guardrail intervention), force - // EndOfStream so Envoy closes the connection cleanly after delivering the final chunk. - endOfStream := originalChunk.EndOfStream || result.StreamTerminated if result.StreamTerminated { slog.Info("[streaming] stream terminated by policy; forcing EndOfStream on final chunk") }