From 22b886318dff2d79ab1c71d143897f9aaca37cb1 Mon Sep 17 00:00:00 2001 From: sb123sb123 <152394158+sb123sb123@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:44:55 +0800 Subject: [PATCH 1/3] wsflate: handle short compressor EOF blocks --- wsflate/writer.go | 27 ++++++++++++++++++++++-- wsflate/writer_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/wsflate/writer.go b/wsflate/writer.go index 0342ccb..77452dc 100644 --- a/wsflate/writer.go +++ b/wsflate/writer.go @@ -55,6 +55,10 @@ type Writer struct { c Compressor cbuf cbuf err error + // flushed reports whether cbuf currently holds the validated + // per-message sync-flush marker. A compressor may append a final block + // when it is closed, but that block is not part of a PMCE message. + flushed bool } // NewWriter returns a new Writer. @@ -74,6 +78,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.flushed = false w.cbuf.reset(dest) if x, ok := w.c.(WriteResetter); ok { x.Reset(&w.cbuf) @@ -87,6 +92,7 @@ func (w *Writer) Write(p []byte) (n int, err error) { if w.err != nil { return 0, w.err } + w.flushed = false n, w.err = w.c.Write(p) return n, w.err } @@ -96,8 +102,12 @@ func (w *Writer) Flush() error { if w.err != nil { return w.err } + w.flushed = false w.err = w.c.Flush() w.checkTail() + if w.err == nil { + w.flushed = true + } return w.err } @@ -108,9 +118,22 @@ func (w *Writer) Close() error { return w.err } if c, ok := w.c.(io.Closer); ok { - w.err = c.Close() + if w.flushed { + // PMCE requires the sync-flush marker to be omitted from the + // message. Keep it buffered while closing the compressor so a + // final block emitted by newer flate implementations cannot + // move any of those bytes to the destination. + dst := w.cbuf.dst + w.cbuf.dst = io.Discard + w.err = c.Close() + w.cbuf.dst = dst + } else { + w.err = c.Close() + } + } + if !w.flushed { + w.checkTail() } - w.checkTail() return w.err } diff --git a/wsflate/writer_test.go b/wsflate/writer_test.go index f088b83..5e7aad2 100644 --- a/wsflate/writer_test.go +++ b/wsflate/writer_test.go @@ -14,6 +14,53 @@ import ( "github.com/gobwas/ws" ) +type shortFinalBlockCompressor struct { + dst io.Writer + closed bool +} + +func (c *shortFinalBlockCompressor) Write(p []byte) (int, error) { + return c.dst.Write(p) +} + +func (c *shortFinalBlockCompressor) Flush() error { + _, err := c.dst.Write(compressionTail[:]) + return err +} + +func (c *shortFinalBlockCompressor) Close() error { + c.closed = true + // Go 1.27's flate writer uses a short fixed-Huffman EOF block instead + // of the four-byte stored block emitted by older versions. + _, err := c.dst.Write([]byte{0x03, 0x00}) + return err +} + +func TestWriterCloseAfterFlushWithShortFinalBlock(t *testing.T) { + var buf bytes.Buffer + var compressor *shortFinalBlockCompressor + w := NewWriter(&buf, func(dst io.Writer) Compressor { + compressor = &shortFinalBlockCompressor{dst: dst} + return compressor + }) + + if _, err := w.Write([]byte("payload")); err != nil { + t.Fatalf("unexpected Write error: %v", err) + } + if err := w.Flush(); err != nil { + t.Fatalf("unexpected Flush error: %v", err) + } + if err := w.Close(); err != nil { + t.Fatalf("unexpected Close error: %v", err) + } + if !compressor.closed { + t.Fatal("Close did not close the compressor") + } + if got, want := buf.String(), "payload"; got != want { + t.Fatalf("unexpected output: %q; want %q", got, want) + } +} + func TestWriter(t *testing.T) { var buf bytes.Buffer w := NewWriter(&buf, func(w io.Writer) Compressor { From 09de78c84487fb0f01d4cdf4950ee2047137c39b Mon Sep 17 00:00:00 2001 From: sb123sb123 <152394158+sb123sb123@users.noreply.github.com> Date: Tue, 29 Sep 2026 08:25:00 +0000 Subject: [PATCH 2/3] Simplify successful flush state assignment --- wsflate/writer.go | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/wsflate/writer.go b/wsflate/writer.go index 77452dc..5e4d690 100644 --- a/wsflate/writer.go +++ b/wsflate/writer.go @@ -105,9 +105,7 @@ func (w *Writer) Flush() error { w.flushed = false w.err = w.c.Flush() w.checkTail() - if w.err == nil { - w.flushed = true - } + w.flushed = w.err == nil return w.err } From 738dc659f43784e54966566d717b6bf530fbe0cb Mon Sep 17 00:00:00 2001 From: Sergey Kamardin Date: Sun, 25 Oct 2026 08:16:44 +0000 Subject: [PATCH 3/3] wsflate: make Writer.Close() flush and ignore compressor output A PMCE message is complete once the compressor has been sync flushed and the trailing 0x00 0x00 0xff 0xff stripped (RFC 7692, 7.2.1). Whatever a compressor emits when it is closed (a final DEFLATE block) is not part of the message, and its exact bytes are an implementation detail: Go 1.26 emitted an empty stored block ending in the same four bytes as a sync flush, so dropping the last four bytes after Close() happened to work; Go 1.27 emits a short fixed-Huffman block, which broke the tail check. Build on the previous commit, which keeps the compressor's Close() output out of the message when Close() follows Flush(), and extend it to the other call sequences: - Close() flushes first unless nothing was written since the last successful Flush(), the way flate, gzip and zlib Writers do. This covers callers that finish a message with a bare Close(), including the autobahn example server. - Close() then always closes the compressor into io.Discard, so the tail check runs only in Flush() and does not depend on how a particular compressor terminates its stream. Flush(), Close() and Flush()+Close() now produce byte-identical messages with compress/flate from both Go 1.26 and Go 1.27, which the new tests assert; the Autobahn suite passes against a Go 1.27 server. Fixes #221. --- wsflate/writer.go | 60 +++++++------- wsflate/writer_test.go | 179 ++++++++++++++++++++++++++++++++++------- 2 files changed, 185 insertions(+), 54 deletions(-) diff --git a/wsflate/writer.go b/wsflate/writer.go index 5e4d690..13c2db4 100644 --- a/wsflate/writer.go +++ b/wsflate/writer.go @@ -31,7 +31,11 @@ type WriteResetter interface { // Essentially Writer is a thin wrapper around Compressor interface to meet // PMCE specs. // -// After all data has been written client should call Flush() method. +// After all data of a message has been written client should call Flush(), +// which completes the message as PMCE requires. Close() flushes the Writer, +// if it has not been flushed yet, and closes the Compressor if it implements +// io.Closer. Whatever the Compressor writes while being closed is ignored. +// // If any error occurs after writing to or flushing a Writer, all subsequent // calls to Write(), Flush() or Close() will return the error. // @@ -50,14 +54,11 @@ type Writer struct { // 4. To make Reader and Writer API the same. That is, there is no Reset() // method for flate.Reader already, so we need to provide it as a wrapper // (see point #3), or drop the Reader.Reset() method. - dest io.Writer - ctor func(io.Writer) Compressor - c Compressor - cbuf cbuf - err error - // flushed reports whether cbuf currently holds the validated - // per-message sync-flush marker. A compressor may append a final block - // when it is closed, but that block is not part of a PMCE message. + dest io.Writer + ctor func(io.Writer) Compressor + c Compressor + cbuf cbuf + err error flushed bool } @@ -97,40 +98,45 @@ func (w *Writer) Write(p []byte) (n int, err error) { return n, w.err } -// Flush writes any pending data into w.Dest. +// Flush writes any pending data into w.Dest and completes the message. +// +// That is, it makes the Compressor emit an empty stored DEFLATE block (sync +// flush) and strips its trailing 0x00 0x00 0xff 0xff from the output, as +// required by RFC 7692 section 7.2.1. The Compressor must end its Flush() +// output with those four bytes; otherwise Flush() returns an error. func (w *Writer) Flush() error { if w.err != nil { return w.err } - w.flushed = false w.err = w.c.Flush() w.checkTail() w.flushed = w.err == nil return w.err } -// Close closes Writer and a Compressor instance used under the hood (if it -// implements io.Closer interface). +// Close closes the Writer by flushing any unwritten data to the underlying +// io.Writer (see Flush()) and closing the Compressor if it implements +// io.Closer. It does not close the underlying io.Writer. +// +// The bytes the Compressor writes while being closed (usually a final DEFLATE +// block) are not part of a PMCE message and are discarded. func (w *Writer) Close() error { if w.err != nil { return w.err } - if c, ok := w.c.(io.Closer); ok { - if w.flushed { - // PMCE requires the sync-flush marker to be omitted from the - // message. Keep it buffered while closing the compressor so a - // final block emitted by newer flate implementations cannot - // move any of those bytes to the destination. - dst := w.cbuf.dst - w.cbuf.dst = io.Discard - w.err = c.Close() - w.cbuf.dst = dst - } else { - w.err = c.Close() + if !w.flushed { + if err := w.Flush(); err != nil { + return err } } - if !w.flushed { - w.checkTail() + if c, ok := w.c.(io.Closer); ok { + // Keep the sync-flush marker held in cbuf while closing the + // compressor, so that whatever it writes can not push any of those + // bytes to the destination. + dst := w.cbuf.dst + w.cbuf.dst = io.Discard + w.err = c.Close() + w.cbuf.dst = dst } return w.err } diff --git a/wsflate/writer_test.go b/wsflate/writer_test.go index 5e7aad2..842cba7 100644 --- a/wsflate/writer_test.go +++ b/wsflate/writer_test.go @@ -14,50 +14,175 @@ import ( "github.com/gobwas/ws" ) -type shortFinalBlockCompressor struct { +// fakeCompressor is a Compressor which passes data through as is, writes +// given bytes on Flush() and Close(), and remembers whether it was closed. +type fakeCompressor struct { dst io.Writer + flush []byte + close []byte closed bool } -func (c *shortFinalBlockCompressor) Write(p []byte) (int, error) { +func (c *fakeCompressor) Write(p []byte) (int, error) { return c.dst.Write(p) } -func (c *shortFinalBlockCompressor) Flush() error { - _, err := c.dst.Write(compressionTail[:]) +func (c *fakeCompressor) Flush() error { + _, err := c.dst.Write(c.flush) return err } -func (c *shortFinalBlockCompressor) Close() error { +func (c *fakeCompressor) Close() error { c.closed = true - // Go 1.27's flate writer uses a short fixed-Huffman EOF block instead - // of the four-byte stored block emitted by older versions. - _, err := c.dst.Write([]byte{0x03, 0x00}) + _, err := c.dst.Write(c.close) return err } -func TestWriterCloseAfterFlushWithShortFinalBlock(t *testing.T) { - var buf bytes.Buffer - var compressor *shortFinalBlockCompressor - w := NewWriter(&buf, func(dst io.Writer) Compressor { - compressor = &shortFinalBlockCompressor{dst: dst} - return compressor - }) - - if _, err := w.Write([]byte("payload")); err != nil { - t.Fatalf("unexpected Write error: %v", err) - } - if err := w.Flush(); err != nil { - t.Fatalf("unexpected Flush error: %v", err) +func TestWriterClose(t *testing.T) { + // Go 1.27's flate writer terminates the stream with a short fixed-Huffman + // block instead of the empty stored block emitted by older versions. Both + // must stay out of the message. + for _, test := range []struct { + name string + close []byte + flush bool + bad bool + }{ + { + name: "stored final block, flushed", + close: []byte{0x01, 0x00, 0x00, 0xff, 0xff}, + flush: true, + }, + { + name: "stored final block, not flushed", + close: []byte{0x01, 0x00, 0x00, 0xff, 0xff}, + }, + { + name: "fixed final block, flushed", + close: []byte{0x03, 0x00}, + flush: true, + }, + { + name: "fixed final block, not flushed", + close: []byte{0x03, 0x00}, + }, + { + name: "no final block, not flushed", + close: nil, + }, + { + name: "bad compressor", + bad: true, + }, + } { + t.Run(test.name, func(t *testing.T) { + var buf bytes.Buffer + var compressor *fakeCompressor + w := NewWriter(&buf, func(dst io.Writer) Compressor { + compressor = &fakeCompressor{ + dst: dst, + flush: compressionTail[:], + close: test.close, + } + if test.bad { + // Does not sync flush. + compressor.flush = nil + } + return compressor + }) + if _, err := w.Write([]byte("payload")); err != nil { + t.Fatalf("unexpected Write() error: %v", err) + } + if test.flush { + if err := w.Flush(); err != nil { + t.Fatalf("unexpected Flush() error: %v", err) + } + } + err := w.Close() + if test.bad { + if err == nil { + t.Fatalf("expected Close() error for bad compressor") + } + return + } + if err != nil { + t.Fatalf("unexpected Close() error: %v", err) + } + if !compressor.closed { + t.Fatalf("Close() did not close the compressor") + } + if got, want := buf.String(), "payload"; got != want { + t.Fatalf("unexpected output: %q; want %q", got, want) + } + }) } - if err := w.Close(); err != nil { - t.Fatalf("unexpected Close error: %v", err) +} + +// TestWriterCloseFlate checks that with the real compress/flate the message +// bytes do not depend on whether Flush(), Close() or both were called. +func TestWriterCloseFlate(t *testing.T) { + type config struct { + flush bool + close bool } - if !compressor.closed { - t.Fatal("Close did not close the compressor") + var ( + data = []byte("hello, flate!") + configs = []config{ + { + flush: true, + }, + { + close: true, + }, + { + flush: true, + close: true, + }, + } + // Message bytes produced by each config and, inversely, the number + // of configs produced each message; all of them must agree. + output = make(map[config]string, len(configs)) + seen = make(map[string]int, len(configs)) + ) + for _, c := range configs { + var buf bytes.Buffer + w := NewWriter(&buf, func(w io.Writer) Compressor { + fw, _ := flate.NewWriter(w, 9) + return fw + }) + if _, err := w.Write(data); err != nil { + t.Fatalf("%+v: unexpected Write() error: %v", c, err) + } + if c.flush { + if err := w.Flush(); err != nil { + t.Fatalf("%+v: unexpected Flush() error: %v", c, err) + } + } + if c.close { + if err := w.Close(); err != nil { + t.Fatalf("%+v: unexpected Close() error: %v", c, err) + } + } + output[c] = buf.String() + seen[output[c]]++ + + // Must be decompressible as a PMCE message. + r := NewReader(&buf, func(r io.Reader) Decompressor { + return flate.NewReader(r) + }) + p, err := io.ReadAll(r) + if err != nil { + t.Fatalf("%+v: unexpected Reader error: %v", c, err) + } + if !bytes.Equal(p, data) { + t.Fatalf("%+v: unexpected bytes: %#q; want %#q", c, p, data) + } } - if got, want := buf.String(), "payload"; got != want { - t.Fatalf("unexpected output: %q; want %q", got, want) + if len(seen) > 1 { + for _, c := range configs { + t.Errorf("%+v: %x (seen %d times)", c, output[c], seen[output[c]]) + } + t.Fatalf("Flush()/Close() combinations produced %d different outputs", len(seen)) } }