Merge pull request #240 from lkinley-rythmos/v4

Fix Reader stale data, Writer buffer leak, and zero-length write

Author: pierrec

reader.go

@@ -207,6 +207,7 @@ func (r *Reader) read(buf []byte) (int, error) {
 	}
 	r.cum += uint32(len(dst))
 	if direct {
+		r.data = r.data[:0]
 		return len(dst), nil
 	}
 	r.data = dst

reader_test.go

@@ -361,3 +361,108 @@ func TestReader_WriteTo(t *testing.T) {
 		t.Fatal("result does not match original")
 	}
 }
+
+// TestReader_DirectModeStaleData verifies that a zero-length uncompressed block
+// in direct mode does not cause stale pool data to be returned. Before the fix,
+// r.data was not cleared after a direct-mode decompress, so a subsequent
+// zero-length block would copy stale data from r.data.
+func TestReader_DirectModeStaleData(t *testing.T) {
+	// Build a minimal LZ4 frame containing:
+	//   1. A normal uncompressed block with known data
+	//   2. A zero-length uncompressed block (0x80000000)
+	//   3. End-of-stream marker (0x00000000)
+	//
+	// First, compress real data to get a valid frame, then append the
+	// zero-length block before the end marker.
+	payload := []byte("hello, world! this is test data for direct mode stale check.")
+	var compressed bytes.Buffer
+	zw := lz4.NewWriter(&compressed)
+	if err := zw.Apply(lz4.ConcurrencyOption(-1), lz4.ChecksumOption(false)); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := zw.Write(payload); err != nil {
+		t.Fatal(err)
+	}
+	if err := zw.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	// The frame ends with: [end mark: 00 00 00 00].
+	// Insert a zero-length uncompressed block (0x80000000) before the end mark.
+	raw := compressed.Bytes()
+	endMark := raw[len(raw)-4:]       // last 4 bytes = end mark
+	prefix := raw[:len(raw)-4]        // everything before end mark
+	var frame bytes.Buffer
+	frame.Write(prefix)
+	frame.Write([]byte{0x00, 0x00, 0x00, 0x80}) // zero-length uncompressed block
+	frame.Write(endMark)                         // end mark
+
+	// Read with a large buffer to force the direct path (buf >= block size)
+	// and concurrency=1. The zero-length block triggers the bn==0 fallback
+	// in Reader.Read, which would copy stale r.data if the fix is missing.
+	zr := lz4.NewReader(&frame)
+	if err := zr.Apply(lz4.ConcurrencyOption(-1)); err != nil {
+		t.Fatal(err)
+	}
+	var result bytes.Buffer
+	buf := make([]byte, 64*1024) // >= default block size (64KB)
+	for {
+		n, readErr := zr.Read(buf)
+		if n > 0 {
+			result.Write(buf[:n])
+		}
+		if readErr == io.EOF {
+			break
+		}
+		if readErr != nil {
+			t.Fatal(readErr)
+		}
+	}
+	if !bytes.Equal(result.Bytes(), payload) {
+		t.Fatalf("decompressed data mismatch: got %d bytes %q, want %d bytes %q",
+			result.Len(), result.Bytes(), len(payload), payload)
+	}
+}
+
+func TestReader_BufferIssue(t *testing.T) {
+	for _, opts := range [][]lz4.Option{
+		nil,
+		_o(lz4.ConcurrencyOption(2)),
+	} {
+		label := fmt.Sprintf("%v", opts)
+		t.Run(label, func(t *testing.T) {
+			pr, pw := io.Pipe()
+			go func(w *io.PipeWriter) {
+				defer w.Close()
+				file, err := os.Open("testdata/bundle.00001.part00000.lz4")
+				if err != nil {
+					w.CloseWithError(err)
+					return
+				}
+				defer file.Close()
+				io.Copy(w, file)
+			}(pw)
+			data := make([]byte, 1024*1024*4)
+			lz4Reader := lz4.NewReader(pr)
+			if err := lz4Reader.Apply(opts...); err != nil {
+				t.Fatal(err)
+			}
+			var total int
+			for {
+				n, readErr := lz4Reader.Read(data)
+				if n > 0 {
+					total += n
+				}
+				if readErr == io.EOF {
+					break
+				}
+				if readErr != nil {
+					t.Fatal(readErr)
+				}
+			}
+			if total != 128*1024*1024 {
+				t.Fatalf("incorrect number of bytes: got %d, want %d", total, 128*1024*1024)
+			}
+		})
+	}
+}

writer.go

@@ -188,6 +188,8 @@ func (w *Writer) Close() error {
 //
 // w.Close must be called before Reset or pending data may be dropped.
 func (w *Writer) Reset(writer io.Writer) {
+	lz4block.Put(w.data)
+	w.data = nil
 	w.frame.Reset(w.num)
 	w.state.reset()
 	w.src = writer
@@ -225,11 +227,13 @@ func (w *Writer) ReadFrom(r io.Reader) (n int64, err error) {
 			return
 		}
 		n += int64(rn)
-		err = w.write(data[:rn], true)
-		if err != nil {
-			return
+		if rn > 0 {
+			err = w.write(data[:rn], true)
+			if err != nil {
+				return
+			}
+			w.handler(rn)
 		}
-		w.handler(rn)
 		if !done && !w.isNotConcurrent() {
 			// The buffer will be returned automatically by go routines (safe=true)
 			// so get a new one fo the next round.

writer_test.go

@@ -380,3 +380,85 @@ func TestWriterConcurrency(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+// TestWriter_ResetWithoutClose verifies that Reset returns the internal buffer
+// to the pool even when Close was not called. Before the fix, the buffer would
+// leak because Reset did not call lz4block.Put.
+func TestWriter_ResetWithoutClose(t *testing.T) {
+	data := []byte(strings.Repeat("hello world ", 1000))
+	buf := new(bytes.Buffer)
+	zw := lz4.NewWriter(buf)
+
+	// Write some data but do NOT close.
+	if _, err := zw.Write(data); err != nil {
+		t.Fatal(err)
+	}
+
+	// Reset without Close — this should not panic or leak.
+	buf.Reset()
+	zw.Reset(buf)
+
+	// The writer should still be fully functional after Reset.
+	if _, err := zw.Write(data); err != nil {
+		t.Fatal(err)
+	}
+	if err := zw.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify the output decompresses correctly.
+	out := new(bytes.Buffer)
+	if _, err := io.Copy(out, lz4.NewReader(buf)); err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(out.Bytes(), data) {
+		t.Fatal("decompressed data does not match original after Reset without Close")
+	}
+}
+
+// zeroThenDataReader returns 0, nil on the first Read call, then delegates
+// to the underlying reader. This simulates an io.Reader that occasionally
+// returns zero bytes without error (allowed by the io.Reader contract).
+type zeroThenDataReader struct {
+	r     io.Reader
+	zeros int
+}
+
+func (z *zeroThenDataReader) Read(p []byte) (int, error) {
+	if z.zeros > 0 {
+		z.zeros--
+		return 0, nil
+	}
+	return z.r.Read(p)
+}
+
+// TestWriter_ReadFromZeroLengthRead verifies that ReadFrom correctly handles
+// an io.Reader that returns 0 bytes without error. Before the fix, a zero-length
+// read would still call write and handler with empty data.
+func TestWriter_ReadFromZeroLengthRead(t *testing.T) {
+	data := []byte(strings.Repeat("test data for ReadFrom ", 500))
+
+	buf := new(bytes.Buffer)
+	zw := lz4.NewWriter(buf)
+	src := &zeroThenDataReader{r: bytes.NewReader(data), zeros: 3}
+
+	n, err := zw.ReadFrom(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if int(n) != len(data) {
+		t.Fatalf("ReadFrom byte count: got %d, want %d", n, len(data))
+	}
+	if err := zw.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify decompressed output matches.
+	out := new(bytes.Buffer)
+	if _, err := io.Copy(out, lz4.NewReader(buf)); err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(out.Bytes(), data) {
+		t.Fatal("decompressed data does not match original after ReadFrom with zero-length reads")
+	}
+}