diff --git a/wsflate/writer.go b/wsflate/writer.go index 0342ccb..64e4a48 100644 --- a/wsflate/writer.go +++ b/wsflate/writer.go @@ -54,7 +54,10 @@ type Writer struct { ctor func(io.Writer) Compressor c Compressor cbuf cbuf - err error + // tailChecked reports whether checkTail() already consumed the stream + // tail for the current (non-reset) message. + tailChecked bool + err error } // NewWriter returns a new Writer. @@ -74,6 +77,7 @@ func NewWriter(w io.Writer, ctor func(io.Writer) Compressor) *Writer { // Any not flushed data will be lost. func (w *Writer) Reset(dest io.Writer) { w.err = nil + w.tailChecked = false w.cbuf.reset(dest) if x, ok := w.c.(WriteResetter); ok { x.Reset(&w.cbuf) @@ -119,11 +123,53 @@ func (w *Writer) Err() error { return w.err } +// tailBFINALFixed is the last 4 bytes cbuf holds when a compressor ends the +// stream with an empty final fixed-Huffman block (BFINAL=1, BTYPE=01) instead +// of the traditional empty non-final stored block. +// +// Starting with Go 1.27 (golang/go commit 388145c, "compress/flate: improve +// compression speed"), compress/flate finishes a stream with a 3-byte empty +// final fixed block (0x03 0x00 0x00: the 0x03 header has BFINAL|BTYPE=01 and +// 5 zero padding bits) instead of the 4-byte empty stored block (0x00 0x00 +// 0xff 0xff) used before. cbuf holds back the last 4 bytes, so the new +// tail reads as {0xff, 0xff, 0x03, 0x00}. +// +// Both tails are valid per RFC 7692 section 7.2.3.4 ("Using a DEFLATE Block +// with BFINAL Set to 1"), which allows ending a message on a BFINAL block +// rather than on an empty non-final stored block. +var tailBFINALFixed = [4]byte{0xff, 0xff, 0x03, 0x00} + func (w *Writer) checkTail() { - if w.err == nil && w.cbuf.buf != compressionTail { + if w.err != nil { + return + } + switch { + case w.tailChecked: + // Tail was already consumed by a previous checkTail() call (e.g. + // Flush() followed by Close()). Nothing left to do. + case w.cbuf.buf == compressionTail: + // Classic tail: the compressor ended on an empty non-final stored + // block. cbuf's 4 held bytes are the 0x00 0x00 0xff 0xff that the + // receiver appends back per RFC 7692 section 7.2.2, so they are + // dropped (nothing is written to dest). + case w.cbuf.buf == tailBFINALFixed: + // Go 1.27 tail: the stream ends with an empty final fixed-Huffman + // block (0x03 0x00) right after the empty non-final stored block. + // cbuf has already flushed the first 3 bytes of that stored block, + // so pass all 4 held bytes (FF FF 03 00) through to the destination + // to reconstruct the full, self-terminating stream: + // … [00 00 00 FF FF] [03 00] + // (empty stored block + final fixed block). The receiver then + // appends 0x00 0x00 0xff 0xff per RFC 7692 7.2.2, which lands after + // the BFINAL block and is simply ignored by the decompressor. + w.cbuf.flush(w.cbuf.buf[:]) + w.cbuf.buf = [4]byte{} + w.cbuf.n = 0 + w.tailChecked = true + default: w.err = fmt.Errorf( - "wsflate: bad compressor: unexpected stream tail: %#x vs %#x", - w.cbuf.buf, compressionTail, + "wsflate: bad compressor: unexpected stream tail: %#x (want %#x or %#x)", + w.cbuf.buf, compressionTail, tailBFINALFixed, ) } } diff --git a/wsflate/writer_tail_test.go b/wsflate/writer_tail_test.go new file mode 100644 index 0000000..ef61a37 --- /dev/null +++ b/wsflate/writer_tail_test.go @@ -0,0 +1,147 @@ +package wsflate + +import ( + "bytes" + "compress/flate" + "io" + "testing" +) + +// blockCompressor is a Compressor that emits an exact, pre-recorded byte +// stream (the blocks, in order) then writes any literal Write() input +// through unchanged. It lets tests pin the exact tail bytes independent of +// the Go toolchain's compress/flate encoder. +type blockCompressor struct { + dst io.Writer + blocks [][]byte + pos int +} + +func (f *blockCompressor) emit() { + for f.pos < len(f.blocks) { + f.dst.Write(f.blocks[f.pos]) + f.pos++ + } +} + +func (f *blockCompressor) Write(p []byte) (int, error) { + f.emit() + f.dst.Write(p) + return len(p), nil +} + +func (f *blockCompressor) Flush() error { f.emit(); return nil } + +func (f *blockCompressor) Close() error { f.emit(); return nil } + +// TestCheckTailClassic pins the legacy behaviour: a compressor that ends on +// an empty non-final stored block produces a 0x00 0x00 0xff 0xff held tail, +// which must be dropped (nothing extra written to the destination) and must +// not raise an error. +func TestCheckTailClassic(t *testing.T) { + // Stream: an empty final fixed block (03 00) followed by an empty + // non-final stored block (00 00 00 FF FF). The last 4 bytes held by + // cbuf are therefore 00 00 FF FF (compressionTail). + var dest bytes.Buffer + w := NewWriter(&dest, func(w io.Writer) Compressor { + return &blockCompressor{ + dst: w, + blocks: [][]byte{{0x03, 0x00}, {0x00, 0x00, 0x00, 0xff, 0xff}}, + } + }) + w.Write(nil) + if err := w.Flush(); err != nil { + t.Fatalf("Flush() error: %v", err) + } + // Flush() then Close() (as CompressTo does) must not double-handle the tail. + if err := w.Close(); err != nil { + t.Fatalf("Close() error: %v", err) + } + if err := w.Err(); err != nil { + t.Fatalf("Writer.Err() = %v; want nil", err) + } + // cbuf already flushed the first 3 bytes of the stored block; the final + // 4 (00 00 ff ff) must be dropped, leaving dest = 03 00 00. + got := dest.Bytes() + if want := []byte{0x03, 0x00, 0x00}; !bytes.Equal(got, want) { + t.Fatalf("dest = % X; want % X", got, want) + } +} + +// TestCheckTailBFINALFixed pins the Go 1.27 behaviour: a compressor that ends +// on an empty final fixed-Huffman block (BFINAL=1, BTYPE=01) produces a +// FF FF 03 00 held tail. checkTail must pass all 4 held bytes through so the +// on-wire payload is a complete, self-terminating DEFLATE stream, and must +// not raise an error. This branch is exercised here on every toolchain, not +// only on Go 1.27. +func TestCheckTailBFINALFixed(t *testing.T) { + // Stream ends: ... [00 00 00 FF FF] [03 00] + // (empty non-final stored block)(empty final fixed block) + // cbuf holds the last 4 bytes: FF FF 03 00 (tailBFINALFixed). + var dest bytes.Buffer + w := NewWriter(&dest, func(w io.Writer) Compressor { + return &blockCompressor{ + dst: w, + blocks: [][]byte{{0x03, 0x00}, {0x00, 0x00, 0x00, 0xff, 0xff}, {0x03, 0x00}}, + } + }) + w.Write(nil) + if err := w.Flush(); err != nil { + t.Fatalf("Flush() error: %v", err) + } + if err := w.Close(); err != nil { + t.Fatalf("Close() error: %v", err) + } + if err := w.Err(); err != nil { + t.Fatalf("Writer.Err() = %v; want nil", err) + } + // Full reconstructed stream: 03 00 00 00 00 FF FF 03 00 + // (9 bytes). The 03 00 final block makes it self-terminating. + got := dest.Bytes() + if want := []byte{0x03, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0x03, 0x00}; !bytes.Equal(got, want) { + t.Fatalf("dest = % X; want % X", got, want) + } + // The on-wire payload must round-trip through an RFC 7692 receiver + // (append 00 00 FF FF, decompress). + var out bytes.Buffer + fr := flate.NewReader(io.MultiReader( + bytes.NewReader(got), bytes.NewReader([]byte{0, 0, 0xff, 0xff}), + )) + if _, err := io.Copy(&out, fr); err != nil { + t.Fatalf("decompressing payload: %v", err) + } +} + +// TestCheckTailUnexpected ensures tails matching neither shape are still +// rejected, preserving the original validation. +func TestCheckTailUnexpected(t *testing.T) { + var dest bytes.Buffer + w := NewWriter(&dest, func(w io.Writer) Compressor { + return &blockCompressor{ + dst: w, + blocks: [][]byte{{0xaa, 0xbb, 0xcc, 0xdd}}, + } + }) + w.Write(nil) + w.Flush() + if w.Err() == nil { + t.Fatal("Writer.Err() = nil; want non-nil for unexpected tail") + } +} + +// TestWriterRoundTripDefault is the end-to-end guard that the default +// compress/flate based helper still round-trips on the current toolchain. +func TestWriterRoundTripDefault(t *testing.T) { + data := []byte("hello, flate!") + var compressed bytes.Buffer + if err := DefaultHelper.CompressTo(&compressed, data); err != nil { + t.Fatalf("CompressTo() error: %v", err) + } + var out bytes.Buffer + if err := DefaultHelper.DecompressTo(&out, compressed.Bytes()); err != nil { + t.Fatalf("DecompressTo() error: %v", err) + } + if !bytes.Equal(out.Bytes(), data) { + t.Fatalf("round-trip mismatch: got %q; want %q", out.String(), string(data)) + } +} diff --git a/wsflate/zz_edge_test.go b/wsflate/zz_edge_test.go new file mode 100644 index 0000000..6013aaf --- /dev/null +++ b/wsflate/zz_edge_test.go @@ -0,0 +1,135 @@ +package wsflate + +import ( + "bytes" + "compress/flate" + "io" + "strings" + "testing" +) + +func TestEdgeCasesRoundTrip(t *testing.T) { + cases := map[string][]byte{ + "empty": {}, + "short": []byte("hi"), + "one-byte": []byte("a"), + "exact-4": []byte("abcd"), + "big-text": []byte(strings.Repeat("The quick brown fox. ", 1000)), + } + cases["randomish"] = []byte(strings.Repeat("random-lookin data | structure", 300)) + for name, data := range cases { + t.Run(name, func(t *testing.T) { + var buf bytes.Buffer + w := NewWriter(&buf, func(w io.Writer) Compressor { f, _ := flate.NewWriter(w, 9); return f }) + w.Write(data) + w.Flush() + w.Close() + if err := w.Err(); err != nil { + t.Fatalf("Writer.Err(): %v", err) + } + var out bytes.Buffer + r := NewReader(&buf, func(r io.Reader) Decompressor { return flate.NewReader(r) }) + io.Copy(&out, r) + if err := r.Err(); err != nil { + t.Fatalf("Reader.Err(): %v (payload=% X)", err, buf.Bytes()) + } + if !bytes.Equal(out.Bytes(), data) { + t.Fatalf("mismatch: got %d bytes want %d", out.Len(), len(data)) + } + }) + } +} + +// Each message self-contained (DefaultParameters / no context takeover), with +// fragmented writes and mid-message Flush() calls. +func TestEdgeNoContextTakeover(t *testing.T) { + msgs := []string{ + "first message with some content", + "", + "second message " + strings.Repeat("structure ", 40), + "tiny", + "abcd", + } + for i, m := range msgs { + var buf bytes.Buffer + w := NewWriter(&buf, func(w io.Writer) Compressor { f, _ := flate.NewWriter(w, 9); return f }) + data := []byte(m) + const chunk = 5 + for len(data) > 0 { + n := chunk + if n > len(data) { + n = len(data) + } + w.Write(data[:n]) + data = data[n:] + if len(data)%2 == 0 { + if err := w.Flush(); err != nil { + t.Fatalf("msg %d mid-flush: %v", i, err) + } + } + } + w.Flush() + w.Close() + if err := w.Err(); err != nil { + t.Fatalf("msg %d writer err: %v (payload=% X)", i, err, buf.Bytes()) + } + var out bytes.Buffer + r := NewReader(&buf, func(r io.Reader) Decompressor { return flate.NewReader(r) }) + io.Copy(&out, r) + if err := r.Err(); err != nil { + t.Fatalf("msg %d reader err: %v (payload=% X)", i, err, buf.Bytes()) + } + if out.String() != m { + t.Fatalf("msg %d: got %q want %q", i, out.String(), m) + } + } +} + +// Context takeover: ONE long-lived writer. Each message ends on Flush(); the +// final one on Close(). Per RFC 7692 the per-message payloads are byte-aligned +// continuations of a single stream, so we reconstruct each message's payload +// from the wire offsets and feed all payloads (each + 00 00 FF FF) to ONE +// long-lived reader with a carried-over window. +func TestEdgeContextTakeover(t *testing.T) { + msgs := []string{ + "first message establishes the window with repeatable content ", + "now this message should reference the previous window content ", + "and yet another referencing the window ", + "shorter", + } + var wire bytes.Buffer + w := NewWriter(&wire, func(w io.Writer) Compressor { f, _ := flate.NewWriter(w, 9); return f }) + for i, m := range msgs { + if _, err := w.Write([]byte(m)); err != nil { + t.Fatalf("msg %d write: %v", i, err) + } + if err := w.Flush(); err != nil { + t.Fatalf("msg %d flush: %v", i, err) + } + if err := w.Err(); err != nil { + t.Fatalf("msg %d writer err: %v", i, err) + } + } + if err := w.Close(); err != nil { + t.Fatalf("close: %v", err) + } + if err := w.Err(); err != nil { + t.Fatalf("final writer err: %v (wire=% X)", err, wire.Bytes()) + } + + // Context takeover = one continuous byte-aligned stream, terminated by the + // final BFINAL block. Feed the whole wire buffer to a single long-lived + // reader (window carried); no per-message suffix is added. + var out bytes.Buffer + r := NewReader(&wire, func(r io.Reader) Decompressor { return flate.NewReader(r) }) + if _, err := io.Copy(&out, r); err != nil { + t.Fatalf("read: %v (wire=% X)", err, wire.Bytes()) + } + if err := r.Err(); err != nil { + t.Fatalf("reader err: %v", err) + } + want := strings.Join(msgs, "") + if out.String() != want { + t.Fatalf("context-takeover concatenation mismatch:\n got %q\nwant %q", out.String(), want) + } +}