fix: close LZ4 writer before reading compressed buffer (#41) (#44)

* fix: close LZ4 writer before reading compressed buffer (#41)

The LZ4 writer was being closed with defer, but compressed.Bytes() was
called before the defer executed. This caused truncated data since
Close() flushes the final compressed blocks.

Applied fix to all 9 storage providers and added regression tests for
providers without external dependencies (simplefs, otter, badger, nuts).

* fix(ci): golangci-lint

---------

Co-authored-by: darkweak <darkweak@protonmail.com>

Author: cdaguerre

badger/badger.go

@@ -234,16 +234,20 @@ func (provider *Badger) SetMultiLevel(baseKey, variedKey string, value []byte, v
 		compressed := new(bytes.Buffer)
 		writer := lz4.NewWriter(compressed)
 
-		defer func() {
+		if _, err = writer.ReadFrom(bytes.NewReader(value)); err != nil {
 			_ = writer.Close()
-		}()
 
-		if _, err = writer.ReadFrom(bytes.NewReader(value)); err != nil {
 			provider.logger.Errorf("Impossible to compress the key %s into Badger, %v", variedKey, err)
 
 			return err
 		}
 
+		if err = writer.Close(); err != nil {
+			provider.logger.Errorf("Impossible to close the compressor for key %s into Badger, %v", variedKey, err)
+
+			return err
+		}
+
 		err = btx.SetEntry(badger.NewEntry([]byte(variedKey), compressed.Bytes()).WithTTL(duration + provider.stale))
 		if err != nil {
 			provider.logger.Errorf("Impossible to set the key %s into Badger, %v", variedKey, err)

badger/badger_test.go

@@ -3,11 +3,14 @@
 package badger_test
 
 import (
+	"bytes"
+	"net/http"
 	"testing"
 	"time"
 
 	"github.com/darkweak/storages/badger"
 	"github.com/darkweak/storages/core"
+	"github.com/pierrec/lz4/v4"
 	"go.uber.org/zap"
 )
 
@@ -132,3 +135,50 @@ func TestBadger_Init(t *testing.T) {
 		t.Error("Impossible to init Badger provider")
 	}
 }
+
+// TestBadger_SetMultiLevel_LargeValue tests that large values are not truncated
+// when stored and retrieved via SetMultiLevel. This reproduces issue #41.
+// See: https://github.com/darkweak/storages/issues/41
+func TestBadger_SetMultiLevel_LargeValue(t *testing.T) {
+	client, err := getBadgerInstance()
+	if err != nil {
+		t.Fatalf("Failed to create badger instance: %v", err)
+	}
+
+	_ = client.Init()
+
+	// Create a large value (5MB) to simulate the issue where large cached
+	// responses were being truncated
+	largeSize := 5 * 1024 * 1024
+
+	largeValue := make([]byte, largeSize)
+	for i := range largeValue {
+		largeValue[i] = byte(i % 256)
+	}
+
+	key := "large_value_test"
+
+	err = client.SetMultiLevel(key, key, largeValue, http.Header{}, "", time.Minute, key)
+	if err != nil {
+		t.Fatalf("Failed to set large value: %v", err)
+	}
+
+	// Retrieve and decompress
+	compressed := client.Get(key)
+	if compressed == nil {
+		t.Fatal("Retrieved value is nil")
+	}
+
+	reader := lz4.NewReader(bytes.NewBuffer(compressed))
+
+	decompressed := new(bytes.Buffer)
+	if _, err := reader.WriteTo(decompressed); err != nil {
+		t.Fatalf("Failed to decompress: %v", err)
+	}
+
+	retrieved := decompressed.Bytes()
+	if len(retrieved) != len(largeValue) {
+		t.Errorf("Data truncation: expected %d bytes, got %d bytes (%.2f%%)",
+			len(largeValue), len(retrieved), float64(len(retrieved))/float64(len(largeValue))*100)
+	}
+}

etcd/etcd.go

@@ -219,16 +219,20 @@ func (provider *Etcd) SetMultiLevel(baseKey, variedKey string, value []byte, var
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
-	defer func() {
+	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		_ = writer.Close()
-	}()
 
-	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		provider.logger.Errorf("Impossible to compress the key %s into Etcd, %v", variedKey, err)
 
 		return err
 	}
 
+	if err := writer.Close(); err != nil {
+		provider.logger.Errorf("Impossible to close the compressor for key %s into Etcd, %v", variedKey, err)
+
+		return err
+	}
+
 	rs, err := provider.Grant(context.TODO(), int64(duration.Seconds()))
 	if err == nil {
 		_, err = provider.Put(provider.ctx, variedKey, compressed.String(), clientv3.WithLease(rs.ID))

go-redis/go-redis.go

@@ -201,16 +201,20 @@ func (provider *Redis) SetMultiLevel(baseKey, variedKey string, value []byte, va
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
-	defer func() {
+	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		_ = writer.Close()
-	}()
 
-	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		provider.logger.Errorf("Impossible to compress the key %s into Redis, %v", variedKey, err)
 
 		return err
 	}
 
+	if err := writer.Close(); err != nil {
+		provider.logger.Errorf("Impossible to close the compressor for key %s into Redis, %v", variedKey, err)
+
+		return err
+	}
+
 	if err := provider.Set(provider.hashtags+variedKey, compressed.Bytes(), duration); err != nil {
 		provider.logger.Errorf("Impossible to set value into Redis, %v", err)
 

nats/nats.go

@@ -267,16 +267,20 @@ func (provider *Nats) SetMultiLevel(baseKey, variedKey string, value []byte, var
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
-	defer func() {
+	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		_ = writer.Close()
-	}()
 
-	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		provider.logger.Errorf("Impossible to compress the key %s into Nats: %v", variedKey, err)
 
 		return err
 	}
 
+	if err := writer.Close(); err != nil {
+		provider.logger.Errorf("Impossible to close the compressor for key %s into Nats, %v", variedKey, err)
+
+		return err
+	}
+
 	property := item{
 		invalidAt: now.Add(duration + provider.stale),
 		value:     compressed.Bytes(),

nuts/nuts.go

@@ -298,16 +298,20 @@ func (provider *Nuts) SetMultiLevel(baseKey, variedKey string, value []byte, var
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
-	defer func() {
+	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		_ = writer.Close()
-	}()
 
-	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		provider.logger.Errorf("Impossible to compress the key %s into Nuts, %v", variedKey, err)
 
 		return err
 	}
 
+	if err := writer.Close(); err != nil {
+		provider.logger.Errorf("Impossible to close the compressor for key %s into Nuts, %v", variedKey, err)
+
+		return err
+	}
+
 	_ = provider.Update(func(tx *nutsdb.Tx) error {
 		return tx.NewBucket(nutsdb.DataStructureBTree, bucket)
 	})

nuts/nuts_test.go

@@ -1,11 +1,14 @@
 package nuts_test
 
 import (
+	"bytes"
+	"net/http"
 	"testing"
 	"time"
 
 	"github.com/darkweak/storages/core"
 	"github.com/darkweak/storages/nuts"
+	"github.com/pierrec/lz4/v4"
 	"go.uber.org/zap"
 )
 
@@ -107,3 +110,50 @@ func TestNuts_Init(t *testing.T) {
 		t.Error("Impossible to init Nuts provider")
 	}
 }
+
+// TestNuts_SetMultiLevel_LargeValue tests that large values are not truncated
+// when stored and retrieved via SetMultiLevel. This reproduces issue #41.
+// See: https://github.com/darkweak/storages/issues/41
+func TestNuts_SetMultiLevel_LargeValue(t *testing.T) {
+	client, err := getNutsInstance()
+	if err != nil {
+		t.Fatalf("Failed to create nuts instance: %v", err)
+	}
+
+	_ = client.Init()
+
+	// Create a large value (5MB) to simulate the issue where large cached
+	// responses were being truncated
+	largeSize := 5 * 1024 * 1024
+
+	largeValue := make([]byte, largeSize)
+	for i := range largeValue {
+		largeValue[i] = byte(i % 256)
+	}
+
+	key := "large_value_test"
+
+	err = client.SetMultiLevel(key, key, largeValue, http.Header{}, "", time.Minute, key)
+	if err != nil {
+		t.Fatalf("Failed to set large value: %v", err)
+	}
+
+	// Retrieve and decompress
+	compressed := client.Get(key)
+	if compressed == nil {
+		t.Fatal("Retrieved value is nil")
+	}
+
+	reader := lz4.NewReader(bytes.NewBuffer(compressed))
+
+	decompressed := new(bytes.Buffer)
+	if _, err := reader.WriteTo(decompressed); err != nil {
+		t.Fatalf("Failed to decompress: %v", err)
+	}
+
+	retrieved := decompressed.Bytes()
+	if len(retrieved) != len(largeValue) {
+		t.Errorf("Data truncation: expected %d bytes, got %d bytes (%.2f%%)",
+			len(largeValue), len(retrieved), float64(len(retrieved))/float64(len(largeValue))*100)
+	}
+}

olric/olric.go

@@ -261,16 +261,20 @@ func (provider *Olric) SetMultiLevel(baseKey, variedKey string, value []byte, va
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
-	defer func() {
+	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		_ = writer.Close()
-	}()
 
-	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		provider.logger.Errorf("Impossible to compress the key %s into Olric, %v", variedKey, err)
 
 		return err
 	}
 
+	if err := writer.Close(); err != nil {
+		provider.logger.Errorf("Impossible to close the compressor for key %s into Olric, %v", variedKey, err)
+
+		return err
+	}
+
 	if err := dmap.Put(context.Background(), variedKey, compressed.Bytes(), olric.EX(duration)); err != nil {
 		provider.logger.Errorf("Impossible to set value into Olric, %v", err)
 

otter/otter.go

@@ -149,16 +149,20 @@ func (provider *Otter) SetMultiLevel(baseKey, variedKey string, value []byte, va
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
-	defer func() {
+	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		_ = writer.Close()
-	}()
 
-	if _, err := writer.ReadFrom(bytes.NewReader(value)); err != nil {
 		provider.logger.Errorf("Impossible to compress the key %s into Otter, %v", variedKey, err)
 
 		return err
 	}
 
+	if err := writer.Close(); err != nil {
+		provider.logger.Errorf("Impossible to close the compressor for key %s into Otter, %v", variedKey, err)
+
+		return err
+	}
+
 	inserted := provider.cache.Set(variedKey, compressed.Bytes(), duration)
 	if !inserted {
 		provider.logger.Errorf("Impossible to set value into Otter, too large for the cost function")

otter/otter_test.go

@@ -1,11 +1,14 @@
 package otter_test
 
 import (
+	"bytes"
+	"net/http"
 	"testing"
 	"time"
 
 	"github.com/darkweak/storages/core"
 	"github.com/darkweak/storages/otter"
+	"github.com/pierrec/lz4/v4"
 	"go.uber.org/zap"
 )
 
@@ -133,3 +136,50 @@ func TestOtter_Init(t *testing.T) {
 		t.Error("Impossible to init Otter provider")
 	}
 }
+
+// TestOtter_SetMultiLevel_LargeValue tests that large values are not truncated
+// when stored and retrieved via SetMultiLevel. This reproduces issue #41.
+// See: https://github.com/darkweak/storages/issues/41
+func TestOtter_SetMultiLevel_LargeValue(t *testing.T) {
+	client, err := getOtterInstance()
+	if err != nil {
+		t.Fatalf("Failed to create otter instance: %v", err)
+	}
+
+	_ = client.Init()
+
+	// Create a large value (5MB) to simulate the issue where large cached
+	// responses were being truncated
+	largeSize := 5 * 1024 * 1024
+
+	largeValue := make([]byte, largeSize)
+	for i := range largeValue {
+		largeValue[i] = byte(i % 256)
+	}
+
+	key := "large_value_test"
+
+	err = client.SetMultiLevel(key, key, largeValue, http.Header{}, "", time.Minute, key)
+	if err != nil {
+		t.Fatalf("Failed to set large value: %v", err)
+	}
+
+	// Retrieve and decompress
+	compressed := client.Get(key)
+	if compressed == nil {
+		t.Fatal("Retrieved value is nil")
+	}
+
+	reader := lz4.NewReader(bytes.NewBuffer(compressed))
+
+	decompressed := new(bytes.Buffer)
+	if _, err := reader.WriteTo(decompressed); err != nil {
+		t.Fatalf("Failed to decompress: %v", err)
+	}
+
+	retrieved := decompressed.Bytes()
+	if len(retrieved) != len(largeValue) {
+		t.Errorf("Data truncation: expected %d bytes, got %d bytes (%.2f%%)",
+			len(largeValue), len(retrieved), float64(len(retrieved))/float64(len(largeValue))*100)
+	}
+}