Implement stateful per-chunk streaming compression

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 {