wsflate: handle short compressor EOF blocks

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
 }
 

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 {