fix(compress): read zlib wrapped deflate, and close the body

internal/compress/compress_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"compress/flate"
 	"compress/gzip"
+	"compress/zlib"
 	"io"
 	"strings"
 	"testing"
@@ -124,3 +125,92 @@ func TestNewCompressReader(t *testing.T) {
 		})
 	}
 }
+
+func zlibData(t *testing.T, data string) io.ReadCloser {
+	t.Helper()
+	var buf bytes.Buffer
+	zw := zlib.NewWriter(&buf)
+	if _, err := zw.Write([]byte(data)); err != nil {
+		t.Fatal(err)
+	}
+	zw.Close()
+	return io.NopCloser(&buf)
+}
+
+// Content-Encoding: deflate is the zlib format of RFC 1950. Raw RFC 1951
+// deflate is the common variant. Both have to read.
+func TestDeflateReaderReadsBothVariants(t *testing.T) {
+	const original = "hello deflate world, long enough to compress"
+	tests := []struct {
+		name string
+		body func(*testing.T, string) io.ReadCloser
+	}{
+		{"zlib", zlibData},
+		{"raw", deflateData},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dr := NewDeflateReader(tt.body(t, original))
+			defer dr.Close()
+
+			out, err := io.ReadAll(dr)
+			if err != nil {
+				t.Fatalf("ReadAll failed: %v", err)
+			}
+			if string(out) != original {
+				t.Fatalf("got %q, want %q", string(out), original)
+			}
+		})
+	}
+}
+
+// The header check is compared against compress/zlib itself over every
+// possible two-byte header, so it cannot drift from the reference and no
+// header has to be worked out by hand. A zlib header naming a preset
+// dictionary is still a zlib header; the reader surfaces zlib's own error
+// for it instead of reading the body as raw deflate.
+func TestLooksLikeZlibAgreesWithCompressZlib(t *testing.T) {
+	var accepted int
+	for cmf := 0; cmf < 256; cmf++ {
+		for flg := 0; flg < 256; flg++ {
+			header := []byte{byte(cmf), byte(flg)}
+			_, err := zlib.NewReader(bytes.NewReader(append(header, make([]byte, 8)...)))
+			want := err == nil || err == zlib.ErrDictionary
+			if want {
+				accepted++
+			}
+			if got := looksLikeZlib(byte(cmf), byte(flg)); got != want {
+				t.Errorf("header %#02x %#02x: looksLikeZlib = %v, compress/zlib says %v", cmf, flg, got, want)
+			}
+		}
+	}
+	if accepted == 0 {
+		t.Fatal("compress/zlib accepted no header at all, the check proves nothing")
+	}
+}
+
+func TestDeflateReaderCloseClosesTheBody(t *testing.T) {
+	body := deflateData(t, "hello deflate world")
+	tracked := &closeTracker{ReadCloser: body}
+	dr := NewDeflateReader(tracked)
+
+	if _, err := io.ReadAll(dr); err != nil {
+		t.Fatalf("ReadAll failed: %v", err)
+	}
+	if err := dr.Close(); err != nil {
+		t.Fatalf("Close failed: %v", err)
+	}
+	if !tracked.closed {
+		t.Fatal("the response body was never closed, so the connection is never released")
+	}
+}
+
+type closeTracker struct {
+	io.ReadCloser
+	closed bool
+}
+
+func (c *closeTracker) Close() error {
+	c.closed = true
+	return c.ReadCloser.Close()
+}

internal/compress/deflate_reader.go

@@ -1,7 +1,9 @@
 package compress
 
 import (
+	"bytes"
 	"compress/flate"
+	"compress/zlib"
 	"io"
 )
 
@@ -15,19 +17,53 @@ func NewDeflateReader(body io.ReadCloser) *DeflateReader {
 	return &DeflateReader{Body: body}
 }
 
+// looksLikeZlib reports whether the two bytes are an RFC 1950 zlib header:
+// the compression method is deflate, the window size is in range, and the
+// pair is a multiple of 31. A raw deflate stream can satisfy all three by
+// chance, so this is a heuristic, but it is the one browsers use and it is
+// the only way to tell the two apart.
+func looksLikeZlib(cmf, flg byte) bool {
+	if cmf&0x0f != 8 { // CM, 8 is deflate
+		return false
+	}
+	if cmf>>4 > 7 { // CINFO, window size
+		return false
+	}
+	return (uint16(cmf)<<8|uint16(flg))%31 == 0 // FCHECK
+}
+
 func (df *DeflateReader) Read(p []byte) (n int, err error) {
 	if df.derr != nil {
 		return 0, df.derr
 	}
 	if df.dr == nil {
-		df.dr = flate.NewReader(df.Body)
+		// Content-Encoding: deflate is the zlib format of RFC 1950, but raw
+		// RFC 1951 deflate is common enough that both have to be read. Sniff
+		// the header and put the bytes back.
+		var head [2]byte
+		got, herr := io.ReadFull(df.Body, head[:])
+		if herr != nil && herr != io.EOF && herr != io.ErrUnexpectedEOF {
+			df.derr = herr
+			return 0, herr
+		}
+		body := io.MultiReader(bytes.NewReader(head[:got]), df.Body)
+		if got == len(head) && looksLikeZlib(head[0], head[1]) {
+			zr, zerr := zlib.NewReader(body)
+			if zerr != nil {
+				df.derr = zerr
+				return 0, zerr
+			}
+			df.dr = zr
+		} else {
+			df.dr = flate.NewReader(body)
+		}
 	}
 	return df.dr.Read(p)
 }
 
 func (df *DeflateReader) Close() error {
 	if df.dr != nil {
-		return df.dr.Close()
+		df.dr.Close()
 	}
 	return df.Body.Close()
 }