Fix compressed streaming responses at the kernel

gateway/gateway-runtime/policy-engine/go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/go-viper/mapstructure/v2 v2.5.0
 	github.com/google/cel-go v0.26.1
 	github.com/google/uuid v1.6.0
+	github.com/klauspost/compress v1.18.6
 	github.com/knadh/koanf/parsers/toml/v2 v2.2.0
 	github.com/knadh/koanf/providers/confmap v1.0.0
 	github.com/knadh/koanf/providers/file v1.2.1
@@ -40,7 +41,6 @@ require (
 	github.com/go-logr/logr v1.4.3 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect
-	github.com/klauspost/compress v1.18.6 // indirect
 	github.com/knadh/koanf/maps v0.1.2 // indirect
 	github.com/mitchellh/copystructure v1.2.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.2 // indirect

gateway/gateway-runtime/policy-engine/internal/constants/constants.go

@@ -105,6 +105,7 @@ const (
 	TerminalReasonNoPolicyChain        = "no_policy_chain"        // route resolved but no chain registered
 	TerminalReasonUnknownMessageType   = "unknown_message_type"   // unrecognised ext_proc message
 	TerminalReasonProcessingFailed     = "processing_failed"      // a phase returned a fatal (stream-ending) error with no ImmediateResponse to classify
+	TerminalReasonUnsupportedEncoding  = "unsupported_encoding"   // Content-Encoding the kernel cannot round-trip, on a body the policy chain requires
 
 	// Analytics metadata and property keys shared across packages.
 	GuardrailHitMetadataKey  = "isGuardrailHit"

gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go

@@ -20,35 +20,109 @@ package kernel
 
 import (
 	"bytes"
+	"compress/flate"
 	"compress/gzip"
+	"compress/zlib"
 	"errors"
 	"fmt"
 	"io"
 	"sync"
 
 	"github.com/andybalholm/brotli"
+	"github.com/klauspost/compress/zstd"
 )
 
+// Content-coding tokens the kernel can round-trip. These are the values stored
+// in requestContentEncoding/responseContentEncoding — already lowercased, and
+// for "deflate" already resolved to one of the two wire variants below.
+const (
+	encodingGzip = "gzip"
+	encodingBr   = "br"
+	encodingZstd = "zstd"
+	// encodingDeflate is "deflate" in its RFC 9110 / RFC 1950 form: DEFLATE data
+	// inside a zlib wrapper.
+	encodingDeflate = "deflate"
+	// encodingDeflateRaw is an INTERNAL token, never a wire value. Some servers
+	// and clients send "Content-Encoding: deflate" carrying bare RFC 1951 DEFLATE
+	// with no zlib wrapper. Both are decodable, but the two are not interchangeable
+	// on output — re-encoding raw input as zlib-wrapped (or vice versa) hands the
+	// peer a body its decoder rejects. Recording which variant arrived lets the
+	// kernel emit the same one back; the Content-Encoding header itself is never
+	// rewritten and stays "deflate" either way.
+	encodingDeflateRaw = "deflate-raw"
+	// encodingIdentity is the no-op coding: present in the header but meaning the
+	// body is not encoded at all.
+	encodingIdentity = "identity"
+)
+
+// zstdDecoderConcurrency/zstdEncoderConcurrency pin the zstd codec to a single
+// goroutine per stream. The library defaults to GOMAXPROCS workers per
+// encoder/decoder, which on a proxy handling many concurrent bodies multiplies
+// into thousands of goroutines for no throughput gain at these body sizes.
+const (
+	zstdDecoderConcurrency = 1
+	zstdEncoderConcurrency = 1
+)
+
+// resolveDeflateVariant inspects the first bytes of a "deflate" body and reports
+// the concrete variant token to record for it.
+//
+// A zlib stream (RFC 1950) starts with a 2-byte header: the low nibble of the
+// first byte is the compression method (8 == DEFLATE) and the big-endian pair is
+// a multiple of 31. Bare DEFLATE data effectively never satisfies both, so this
+// check distinguishes the two reliably. Too few bytes to tell yet is treated as
+// the RFC-conformant zlib form.
+func resolveDeflateVariant(body []byte) string {
+	if len(body) < 2 {
+		return encodingDeflate
+	}
+	if body[0]&0x0f == 0x08 && (uint16(body[0])<<8|uint16(body[1]))%31 == 0 {
+		return encodingDeflate
+	}
+	return encodingDeflateRaw
+}
+
 // ErrDecompressedTooLarge is returned when decompressed output exceeds the
 // configured ceiling — the signature of a decompression bomb.
 var ErrDecompressedTooLarge = errors.New("decompressed body exceeds maximum allowed size")
 
 // decompressBody decompresses body bytes based on the Content-Encoding value.
-// Supported encodings: "gzip", "br" (Brotli). Unknown encodings are returned as-is.
+// Supported encodings: gzip, br, zstd, and both deflate variants. Unknown
+// encodings are returned as-is — callers must not reach this with one, since
+// isRecompressibleEncoding gates every call site (an unsupported encoding is
+// rejected outright rather than handed to policies as opaque bytes).
 // Output is capped at maxBytes (<= 0 means unbounded); exceeding it returns
 // ErrDecompressedTooLarge, never a truncated body.
 func decompressBody(body []byte, encoding string, maxBytes int64) ([]byte, error) {
 	switch encoding {
-	case "gzip":
+	case encodingGzip:
 		r, err := gzip.NewReader(bytes.NewReader(body))
 		if err != nil {
 			return nil, fmt.Errorf("gzip reader: %w", err)
 		}
 		defer r.Close()
 		return readLimited(r, maxBytes)
-	case "br":
+	case encodingBr:
 		r := brotli.NewReader(bytes.NewReader(body))
 		return readLimited(r, maxBytes)
+	case encodingZstd:
+		r, err := zstd.NewReader(bytes.NewReader(body), zstd.WithDecoderConcurrency(zstdDecoderConcurrency))
+		if err != nil {
+			return nil, fmt.Errorf("zstd reader: %w", err)
+		}
+		defer r.Close()
+		return readLimited(r, maxBytes)
+	case encodingDeflate:
+		r, err := zlib.NewReader(bytes.NewReader(body))
+		if err != nil {
+			return nil, fmt.Errorf("deflate (zlib) reader: %w", err)
+		}
+		defer r.Close()
+		return readLimited(r, maxBytes)
+	case encodingDeflateRaw:
+		r := flate.NewReader(bytes.NewReader(body))
+		defer r.Close()
+		return readLimited(r, maxBytes)
 	default:
 		return body, nil
 	}
@@ -180,9 +254,13 @@ func newStreamDecompressor(encoding string, maxBytes int64) *streamDecompressor
 	go func() {
 		defer close(outChan)
 		defer close(decoderDone)
+		// gzip/zlib/zstd readers consume their stream header eagerly, so
+		// construction blocks here until the first chunk arrives. That is why the
+		// decoder lives on its own goroutine: newStreamDecompressor must return
+		// before any body byte has been seen.
 		var r io.Reader
 		switch encoding {
-		case "gzip":
+		case encodingGzip:
 			gr, err := gzip.NewReader(input)
 			if err != nil {
 				select {
@@ -193,8 +271,34 @@ func newStreamDecompressor(encoding string, maxBytes int64) *streamDecompressor
 			}
 			defer gr.Close()
 			r = gr
-		case "br":
+		case encodingBr:
 			r = brotli.NewReader(input)
+		case encodingZstd:
+			zr, err := zstd.NewReader(input, zstd.WithDecoderConcurrency(zstdDecoderConcurrency))
+			if err != nil {
+				select {
+				case errChan <- fmt.Errorf("zstd.NewReader: %w", err):
+				default:
+				}
+				return
+			}
+			defer zr.Close()
+			r = zr
+		case encodingDeflate:
+			zr, err := zlib.NewReader(input)
+			if err != nil {
+				select {
+				case errChan <- fmt.Errorf("zlib.NewReader: %w", err):
+				default:
+				}
+				return
+			}
+			defer zr.Close()
+			r = zr
+		case encodingDeflateRaw:
+			fr := flate.NewReader(input)
+			defer fr.Close()
+			r = fr
 		default:
 			r = input
 		}
@@ -344,32 +448,143 @@ func (sd *streamDecompressor) Close() {
 	}
 }
 
-// recompressBody re-compresses body bytes using the original Content-Encoding.
-// Used to restore compression after policies have processed the decompressed body.
-// Supported encodings: "gzip", "br" (Brotli). Unknown encodings are returned as-is.
-func recompressBody(body []byte, encoding string) ([]byte, error) {
+// ─── Streaming re-compression ────────────────────────────────────────────────
+//
+// A streamed response must be re-compressed as ONE compressed stream spanning
+// the whole body, not one per chunk. Calling recompressBody per chunk produces
+// N independent members: for gzip that is a multi-member stream which most HTTP
+// clients (Go's transport, httpx/urllib3, curl) do not read past the first
+// member, so the client silently sees a truncated body; for brotli, which has
+// no multi-member concatenation at all, the remainder is undecodable.
+//
+// streamCompressor keeps a single writer alive for the lifetime of the response
+// and flushes after each chunk so data still reaches the client incrementally.
+type streamCompressor struct {
+	encoding string
+	buf      bytes.Buffer
+	// w and flush are the encoder for this stream. Every supported codec exposes
+	// Write/Close/Flush, so they are held behind these two fields rather than one
+	// typed field per codec — a per-codec field forces every method here to grow a
+	// new case, and a missed one silently degrades to "no compression applied".
+	w      io.WriteCloser
+	flush  func() error
+	closed bool
+}
+
+// newStreamCompressor returns a compressor for the encoding, or nil when the
+// encoding needs no re-compression (callers then forward bytes unchanged).
+// Encodings must be pre-validated with isRecompressibleEncoding; nil here means
+// "forward untouched", which is only correct for an unencoded body.
+func newStreamCompressor(encoding string) *streamCompressor {
+	sc := &streamCompressor{encoding: encoding}
 	switch encoding {
-	case "gzip":
-		var buf bytes.Buffer
-		w := gzip.NewWriter(&buf)
-		if _, err := w.Write(body); err != nil {
-			return nil, fmt.Errorf("gzip write: %w", err)
+	case encodingGzip:
+		w := gzip.NewWriter(&sc.buf)
+		sc.w, sc.flush = w, w.Flush
+	case encodingBr:
+		w := brotli.NewWriter(&sc.buf)
+		sc.w, sc.flush = w, w.Flush
+	case encodingZstd:
+		// Error is unreachable: it reports invalid encoder options, and the
+		// options here are compile-time constants.
+		w, err := zstd.NewWriter(&sc.buf, zstd.WithEncoderConcurrency(zstdEncoderConcurrency))
+		if err != nil {
+			return nil
 		}
-		if err := w.Close(); err != nil {
-			return nil, fmt.Errorf("gzip close: %w", err)
+		sc.w, sc.flush = w, w.Flush
+	case encodingDeflate:
+		w := zlib.NewWriter(&sc.buf)
+		sc.w, sc.flush = w, w.Flush
+	case encodingDeflateRaw:
+		// Error is unreachable: it reports an out-of-range level, and the level
+		// here is a library constant.
+		w, err := flate.NewWriter(&sc.buf, flate.DefaultCompression)
+		if err != nil {
+			return nil
 		}
-		return buf.Bytes(), nil
-	case "br":
-		var buf bytes.Buffer
-		w := brotli.NewWriter(&buf)
-		if _, err := w.Write(body); err != nil {
-			return nil, fmt.Errorf("brotli write: %w", err)
+		sc.w, sc.flush = w, w.Flush
+	default:
+		return nil
+	}
+	return sc
+}
+
+// Compress writes one chunk into the single ongoing compressed stream and
+// returns the bytes produced so far. When endOfStream is set the stream is
+// finalised (footer/checksum written) and the compressor must not be reused.
+//
+// A flush is emitted per chunk so the client receives data incrementally; this
+// costs a few bytes of framing per chunk versus a single whole-body compress,
+// which is the correct trade for a streaming response.
+func (sc *streamCompressor) Compress(body []byte, endOfStream bool) ([]byte, error) {
+	if sc.closed {
+		return nil, fmt.Errorf("%s stream compressor already closed", sc.encoding)
+	}
+	sc.buf.Reset()
+
+	if len(body) > 0 {
+		if _, err := sc.w.Write(body); err != nil {
+			return nil, fmt.Errorf("%s write: %w", sc.encoding, err)
 		}
-		if err := w.Close(); err != nil {
-			return nil, fmt.Errorf("brotli close: %w", err)
+	}
+	if endOfStream {
+		if err := sc.w.Close(); err != nil {
+			return nil, fmt.Errorf("%s close: %w", sc.encoding, err)
 		}
-		return buf.Bytes(), nil
+		sc.closed = true
+	} else if err := sc.flush(); err != nil {
+		return nil, fmt.Errorf("%s flush: %w", sc.encoding, err)
+	}
+
+	out := make([]byte, sc.buf.Len())
+	copy(out, sc.buf.Bytes())
+	return out, nil
+}
+
+// Close finalises the stream on error paths where endOfStream never arrives.
+func (sc *streamCompressor) Close() {
+	if sc.closed {
+		return
+	}
+	sc.closed = true
+	_ = sc.w.Close()
+}
+
+// isRecompressibleEncoding reports whether the kernel can decompress and
+// re-compress this Content-Encoding. Anything else is rejected outright by
+// execution_context.go: the kernel neither runs body policies on bytes they
+// cannot read nor forwards a body it could not have inspected.
+func isRecompressibleEncoding(encoding string) bool {
+	switch encoding {
+	case encodingGzip, encodingBr, encodingZstd, encodingDeflate, encodingDeflateRaw:
+		return true
 	default:
-		return body, nil
+		return false
+	}
+}
+
+// recompressBody re-compresses body bytes using the original Content-Encoding.
+// Used for the BUFFERED response path, where the whole body is compressed in a
+// single call. Streaming responses must use streamCompressor instead so the
+// response is one compressed stream rather than one per chunk.
+// Supported encodings: gzip, br, zstd, and both deflate variants. Unknown
+// encodings are returned as-is; call sites are gated by isRecompressibleEncoding
+// so that case is unreachable for a body policies actually touched.
+func recompressBody(body []byte, encoding string) ([]byte, error) {
+	// Reuse the streaming encoder in a single write+finalise, so the buffered and
+	// streaming paths cannot drift apart on which encodings they support or how
+	// each one is framed.
+	sc := newStreamCompressor(encoding)
+	if sc == nil {
+		// An encoding this function does not encode at all: return the body
+		// untouched, as it always has.
+		if !isRecompressibleEncoding(encoding) {
+			return body, nil
+		}
+		// A supported encoding whose codec refused its options. Returning the body
+		// here would emit plaintext under a compressed Content-Encoding header —
+		// the exact corruption this file exists to prevent.
+		return nil, fmt.Errorf("no compressor available for encoding %q", encoding)
 	}
+	return sc.Compress(body, true)
 }