Handle invalid buffers being put back into the Block pool

internal/lz4block/blocks.go

@@ -61,15 +61,24 @@ func (b BlockSizeIndex) Get() []byte {
 }
 
 func (b BlockSizeIndex) Put(buf []byte) {
-	switch b {
+	// Safeguard: do not allow invalid buffers.
+	switch c := uint32(cap(buf)); b {
 	case 4:
-		BlockPool64K.Put(buf)
+		if c == Block64Kb {
+			BlockPool64K.Put(buf[:c])
+		}
 	case 5:
-		BlockPool256K.Put(buf)
+		if c == Block256Kb {
+			BlockPool256K.Put(buf[:c])
+		}
 	case 6:
-		BlockPool1M.Put(buf)
+		if c == Block1Mb {
+			BlockPool1M.Put(buf[:c])
+		}
 	case 7:
-		BlockPool4M.Put(buf)
+		if c == Block4Mb {
+			BlockPool4M.Put(buf[:c])
+		}
 	}
 }
 

internal/lz4stream/frame.go

@@ -260,11 +260,14 @@ type FrameDataBlock struct {
 }
 
 func (b *FrameDataBlock) CloseW(f *Frame) {
-	size := f.Descriptor.Flags.BlockSizeIndex()
-	size.Put(b.data)
-	b.Data = nil
-	b.data = nil
-	b.src = nil
+	if b.data != nil {
+		// Block was not already closed.
+		size := f.Descriptor.Flags.BlockSizeIndex()
+		size.Put(b.data)
+		b.Data = nil
+		b.data = nil
+		b.src = nil
+	}
 }
 
 // Block compression errors are ignored since the buffer is sized appropriately.

reader_test.go

@@ -91,6 +91,8 @@ func TestReader_Reset(t *testing.T) {
 
 	buf.Reset()
 	src.Reset(data)
+	// Another time to maybe trigger some edge case.
+	src.Reset(data)
 	zr.Reset(src)
 	if _, err := io.Copy(buf, zr); err != nil {
 		t.Fatal(err)

writer_test.go

@@ -98,6 +98,8 @@ func TestWriter_Reset(t *testing.T) {
 	buf.Reset()
 	src.Reset(data)
 	zw.Reset(buf)
+	zw.Reset(buf)
+	// Another time to maybe trigger some edge case.
 	if _, err := io.Copy(zw, src); err != nil {
 		t.Fatal(err)
 	}