Review feedback

Author: alpe

block/internal/da/client.go

@@ -113,7 +113,7 @@ func (c *client) Submit(ctx context.Context, data [][]byte, _ float64, namespace
 		c.logger.Debug().
 			Int("original_size", len(raw)).
 			Int("compressed_size", len(compressed)).
-			Float64("ratio", float64(len(compressed))/float64(len(raw))).
+			Float64("ratio", float64(len(compressed))/float64(max(len(raw), 1))).
 			Int("level", int(compLevel)).
 			Msg("compressed blob for DA submission")
 
@@ -311,13 +311,14 @@ func (c *client) Retrieve(ctx context.Context, height uint64, namespace []byte)
 	data := make([]datypes.Blob, len(blobs))
 	for i, b := range blobs {
 		ids[i] = blobrpc.MakeID(height, b.Commitment)
-		decompressed, decompErr := da.Decompress(b.Data())
+		decompressed, decompErr := da.Decompress(ctx, b.Data())
 		if decompErr != nil {
 			return datypes.ResultRetrieve{
 				BaseResult: datypes.BaseResult{
-					Code:    datypes.StatusError,
-					Message: fmt.Sprintf("decompress blob %d at height %d: %v", i, height, decompErr),
-					Height:  height,
+					Code:      datypes.StatusError,
+					Message:   fmt.Sprintf("decompress blob %d at height %d: %v", i, height, decompErr),
+					Height:    height,
+					Timestamp: blockTime,
 				},
 			}
 		}
@@ -399,7 +400,7 @@ func (c *client) Get(ctx context.Context, ids []datypes.ID, namespace []byte) ([
 		if b == nil {
 			continue
 		}
-		decompressed, decompErr := da.Decompress(b.Data())
+		decompressed, decompErr := da.Decompress(ctx, b.Data())
 		if decompErr != nil {
 			return nil, fmt.Errorf("decompress blob: %w", decompErr)
 		}

pkg/da/compression.go

@@ -1,7 +1,10 @@
 package da
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"time"
 
 	"github.com/klauspost/compress/zstd"
 )
@@ -10,6 +13,20 @@ import (
 // ASCII "ZSTD" = 0x5A 0x53 0x54 0x44.
 var magic = []byte{0x5A, 0x53, 0x54, 0x44}
 
+// maxDecompressedSize is the maximum allowed decompressed output size.
+// Matches the WithDecoderMaxMemory cap and provides early rejection
+// by inspecting the zstd frame header before allocating anything.
+const maxDecompressedSize = 7 * 1024 * 1024 // 7 MiB
+
+// decompressTimeout is the hard wall-clock cap on a single decompression.
+// This guards against CPU-based decompression bombs (crafted inputs that
+// are slow to decode) independently of the caller's context deadline.
+const decompressTimeout = 500 * time.Millisecond
+
+// ErrDecompressedSizeExceeded is returned when the zstd frame header
+// declares a decompressed size that exceeds the allowed limit.
+var ErrDecompressedSizeExceeded = errors.New("compression: declared decompressed size exceeds limit")
+
 // CompressionLevel controls the speed/ratio trade-off for blob compression.
 type CompressionLevel int
 
@@ -44,8 +61,9 @@ func init() {
 		}
 		encoders[i] = enc
 	}
+	const maxDecoderMemory = 7 * 1024 * 1024 // 7 MiB cap
 	var err error
-	decoder, err = zstd.NewReader(nil)
+	decoder, err = zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecoderMemory))
 	if err != nil {
 		panic(fmt.Sprintf("compression: create zstd decoder: %v", err))
 	}
@@ -72,21 +90,51 @@ func Compress(data []byte, level CompressionLevel) ([]byte, error) {
 	return result, nil
 }
 
+// decodeResult carries the output of a DecodeAll goroutine.
+type decodeResult struct {
+	data []byte
+	err  error
+}
+
 // Decompress decompresses data that was compressed with Compress.
 // If the data does not have the magic prefix, it is returned as-is
 // (backward-compatible passthrough for uncompressed blobs).
-func Decompress(data []byte) ([]byte, error) {
+func Decompress(ctx context.Context, data []byte) ([]byte, error) {
 	if !IsCompressed(data) {
 		return data, nil
 	}
 
-	// Strip magic prefix and decompress
-	decompressed, err := decoder.DecodeAll(data[len(magic):], nil)
-	if err != nil {
-		return nil, fmt.Errorf("compression: zstd decompress: %w", err)
+	payload := data[len(magic):]
+
+	// Layer 1: Parse frame header to check declared decompressed size
+	// before allocating anything. This is a zero-cost upfront rejection.
+	var hdr zstd.Header
+	if err := hdr.Decode(payload); err == nil && hdr.HasFCS {
+		if hdr.FrameContentSize > maxDecompressedSize {
+			return nil, fmt.Errorf("%w: %d bytes declared, %d allowed",
+				ErrDecompressedSizeExceeded, hdr.FrameContentSize, maxDecompressedSize)
+		}
 	}
 
-	return decompressed, nil
+	// Layer 3: Apply the shorter of caller deadline and our hard cap.
+	ctx, cancel := context.WithTimeout(ctx, decompressTimeout)
+	defer cancel()
+
+	ch := make(chan decodeResult, 1)
+	go func() {
+		out, err := decoder.DecodeAll(payload, nil)
+		ch <- decodeResult{data: out, err: err}
+	}()
+
+	select {
+	case res := <-ch:
+		if res.err != nil {
+			return nil, fmt.Errorf("zstd decompress: %w", res.err)
+		}
+		return res.data, nil
+	case <-ctx.Done():
+		return nil, fmt.Errorf("zstd decompress timeout: %w", ctx.Err())
+	}
 }
 
 // IsCompressed reports whether data starts with the compression magic prefix.

pkg/da/compression_test.go

@@ -2,9 +2,11 @@ package da
 
 import (
 	"bytes"
+	"context"
 	"crypto/rand"
 	"testing"
 
+	"github.com/klauspost/compress/zstd"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -39,7 +41,7 @@ func TestCompressDecompress_RoundTrip(t *testing.T) {
 
 			assert.True(t, IsCompressed(compressed), "compressed data should have magic prefix")
 
-			decompressed, err := Decompress(compressed)
+			decompressed, err := Decompress(context.Background(), compressed)
 			require.NoError(t, err)
 
 			assert.Equal(t, tt.data, decompressed, "round-trip should preserve data")
@@ -67,7 +69,7 @@ func TestCompress_AllLevelsRoundTrip(t *testing.T) {
 
 			sizes = append(sizes, len(compressed))
 
-			decompressed, err := Decompress(compressed)
+			decompressed, err := Decompress(context.Background(), compressed)
 			require.NoError(t, err)
 			assert.Equal(t, data, decompressed)
 
@@ -94,32 +96,32 @@ func TestCompress_Empty(t *testing.T) {
 func TestDecompress_UncompressedPassthrough(t *testing.T) {
 	// Data without magic prefix should pass through unchanged
 	raw := []byte("this is uncompressed protobuf data")
-	result, err := Decompress(raw)
+	result, err := Decompress(context.Background(), raw)
 	require.NoError(t, err)
 	assert.Equal(t, raw, result)
 }
 
 func TestDecompress_Empty(t *testing.T) {
-	result, err := Decompress(nil)
+	result, err := Decompress(context.Background(), nil)
 	require.NoError(t, err)
 	assert.Nil(t, result)
 
-	result, err = Decompress([]byte{})
+	result, err = Decompress(context.Background(), []byte{})
 	require.NoError(t, err)
 	assert.Empty(t, result)
 }
 
 func TestDecompress_ShortData(t *testing.T) {
 	// Data shorter than magic prefix should pass through
-	result, err := Decompress([]byte{0x5A, 0x53})
+	result, err := Decompress(context.Background(), []byte{0x5A, 0x53})
 	require.NoError(t, err)
 	assert.Equal(t, []byte{0x5A, 0x53}, result)
 }
 
 func TestDecompress_CorruptCompressedData(t *testing.T) {
 	// Magic prefix followed by invalid zstd data
 	corrupt := append([]byte{0x5A, 0x53, 0x54, 0x44}, []byte("not valid zstd")...)
-	_, err := Decompress(corrupt)
+	_, err := Decompress(context.Background(), corrupt)
 	assert.Error(t, err, "should fail on corrupt compressed data")
 }
 
@@ -163,14 +165,14 @@ func TestCompress_RandomDataStillWorks(t *testing.T) {
 	compressed, err := Compress(data, LevelFastest)
 	require.NoError(t, err)
 
-	decompressed, err := Decompress(compressed)
+	decompressed, err := Decompress(context.Background(), compressed)
 	require.NoError(t, err)
 	assert.Equal(t, data, decompressed)
 }
 
 func TestDecompress_DataStartingWithMagicButUncompressed(t *testing.T) {
 	fakeCompressed := append([]byte{0x5A, 0x53, 0x54, 0x44}, bytes.Repeat([]byte{0x00}, 100)...)
-	_, err := Decompress(fakeCompressed)
+	_, err := Decompress(context.Background(), fakeCompressed)
 	assert.Error(t, err, "data starting with magic but containing invalid zstd should error")
 }
 
@@ -180,7 +182,43 @@ func TestCompress_InvalidLevel(t *testing.T) {
 	compressed, err := Compress(data, CompressionLevel(99))
 	require.NoError(t, err)
 
-	decompressed, err := Decompress(compressed)
+	decompressed, err := Decompress(context.Background(), compressed)
 	require.NoError(t, err)
 	assert.Equal(t, data, decompressed)
 }
+
+func TestDecompress_ContextCanceled(t *testing.T) {
+	data := []byte("test data for context cancellation")
+	compressed, err := Compress(data, LevelDefault)
+	require.NoError(t, err)
+
+	// Pre-canceled context should cause Decompress to return an error.
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err = Decompress(ctx, compressed)
+	assert.Error(t, err, "decompress with canceled context should fail")
+	assert.ErrorIs(t, err, context.Canceled)
+}
+
+func TestDecompress_OversizedFrameRejected(t *testing.T) {
+	// Build a fake zstd frame header that declares 100 MiB decompressed size.
+	// This should be rejected by the frame header pre-check before any
+	// decompression occurs.
+	hdr := zstd.Header{
+		SingleSegment:    true,
+		HasFCS:           true,
+		FrameContentSize: 100 * 1024 * 1024, // 100 MiB — way over the 7 MiB limit
+	}
+	frame, err := hdr.AppendTo(nil)
+	require.NoError(t, err)
+
+	// Prepend our custom magic prefix
+	blob := make([]byte, len(magic)+len(frame))
+	copy(blob, magic)
+	copy(blob[len(magic):], frame)
+
+	_, err = Decompress(context.Background(), blob)
+	assert.Error(t, err, "should reject blob declaring oversized decompressed output")
+	assert.ErrorIs(t, err, ErrDecompressedSizeExceeded)
+}