diff --git a/internal/compress/compress_test.go b/internal/compress/compress_test.go index 05936e4d..55c6b7c9 100644 --- a/internal/compress/compress_test.go +++ b/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() +} diff --git a/internal/compress/deflate_reader.go b/internal/compress/deflate_reader.go index 44cf4943..8b6c1d45 100644 --- a/internal/compress/deflate_reader.go +++ b/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() }