diff --git a/wsflate/writer.go b/wsflate/writer.go index 0342ccb..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,11 +54,12 @@ 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 + dest io.Writer + ctor func(io.Writer) Compressor + c Compressor + cbuf cbuf + err error + flushed bool } // NewWriter returns a new Writer. @@ -74,6 +79,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,30 +93,51 @@ 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 } -// 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.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 !w.flushed { + if err := w.Flush(); err != nil { + return err + } + } 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 } - w.checkTail() return w.err } diff --git a/wsflate/writer_test.go b/wsflate/writer_test.go index f088b83..842cba7 100644 --- a/wsflate/writer_test.go +++ b/wsflate/writer_test.go @@ -14,6 +14,178 @@ import ( "github.com/gobwas/ws" ) +// 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 *fakeCompressor) Write(p []byte) (int, error) { + return c.dst.Write(p) +} + +func (c *fakeCompressor) Flush() error { + _, err := c.dst.Write(c.flush) + return err +} + +func (c *fakeCompressor) Close() error { + c.closed = true + _, err := c.dst.Write(c.close) + return 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) + } + }) + } +} + +// 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 + } + 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 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)) + } +} + func TestWriter(t *testing.T) { var buf bytes.Buffer w := NewWriter(&buf, func(w io.Writer) Compressor {