diff --git a/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go b/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go index 087cd0a7e..4c9a652f0 100644 --- a/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go +++ b/gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go @@ -168,8 +168,94 @@ func (sd *streamDecompressor) Close() { } } +// 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 61d9a95f2..cc2bee53e 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" @@ -280,3 +281,147 @@ func TestStreamDecompressor_RoundTrip(t *testing.T) { require.NoError(t, err) assert.Equal(t, original, final) } + +// ============================================================================= +// 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 6e109c595..eea227ba6 100644 --- a/gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go +++ b/gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go @@ -118,6 +118,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. @@ -127,6 +131,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 diff --git a/gateway/gateway-runtime/policy-engine/internal/kernel/translator.go b/gateway/gateway-runtime/policy-engine/internal/kernel/translator.go index e6dd5d397..4ee78cf69 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") }