wsflate: handle short compressor EOF blocks

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,10 @@ 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
 }
 
@@ -108,9 +116,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
 }
 

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 {