bring back old behavior, add tests to verify

Author: djeebus

packages/orchestrator/pkg/sandbox/block/cache.go

@@ -258,6 +258,11 @@ func (c *Cache) isCached(off, length int64) bool {
 	// Cap if the length goes beyond the cache size, so we don't check for blocks that are out of bounds.
 	end := min(off+length, c.size)
 
+	// Handle zero-length or empty range (vacuously true - no blocks to check)
+	if end <= off {
+		return true
+	}
+
 	// Check all blocks in the range
 	startBlock := uint(header.BlockIdx(off, c.blockSize))
 	endBlock := uint(header.BlockIdx(end-1, c.blockSize))
@@ -275,6 +280,11 @@ func (c *Cache) isCached(off, length int64) bool {
 }
 
 func (c *Cache) setIsCached(off, length int64) {
+	// Handle zero-length - nothing to mark
+	if length <= 0 {
+		return
+	}
+
 	startBlock := uint(header.BlockIdx(off, c.blockSize))
 	endBlock := uint(header.BlockIdx(off+length-1, c.blockSize))
 

packages/orchestrator/pkg/sandbox/block/cache_dirty_test.go

@@ -0,0 +1,198 @@
+package block
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+func TestIsCached_ZeroSizeCache(t *testing.T) {
+	t.Parallel()
+
+	// Create a zero-size cache (triggers the size == 0 branch in NewCache)
+	cache, err := NewCache(0, header.PageSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// With size=0, any isCached call should return false (or handle gracefully)
+	// The bug: end = min(off+length, 0) = 0, then end-1 = -1 causes underflow
+	// Original behavior: returned true for empty range (no blocks to check)
+
+	// This should not panic and should return a sensible value
+	result := cache.isCached(0, 100)
+	// With zero-size cache, nothing can be cached
+	assert.False(t, result, "zero-size cache should report nothing as cached")
+}
+
+func TestIsCached_ZeroLength(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Zero-length query should return true (vacuously true - no blocks to check)
+	// This matches the original sync.Map behavior where BlocksOffsets(0, blockSize)
+	// returned an empty slice, so the loop didn't execute and returned true.
+	result := cache.isCached(0, 0)
+	assert.True(t, result, "zero-length isCached should return true (vacuously)")
+
+	result = cache.isCached(blockSize*5, 0)
+	assert.True(t, result, "zero-length isCached at any offset should return true")
+}
+
+func TestSetIsCached_ZeroLength(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Zero-length setIsCached should be a no-op
+	// The bug: off=0, length=0 -> off+length-1 = -1 -> BlockIdx(-1, blockSize) = 0
+	// This incorrectly marks block 0 as dirty
+	cache.setIsCached(0, 0)
+
+	// Block 0 should NOT be marked as cached after a zero-length set
+	// We need to check directly since isCached(0, 0) returns true vacuously
+	result := cache.isCached(0, blockSize)
+	assert.False(t, result, "zero-length setIsCached should not mark any blocks as dirty")
+}
+
+func TestSetIsCached_ZeroLengthAtNonZeroOffset(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Zero-length at offset in block 5 should not mark anything
+	cache.setIsCached(blockSize*5, 0)
+
+	// No blocks should be marked as cached
+	for i := int64(0); i < 10; i++ {
+		result := cache.isCached(i*blockSize, blockSize)
+		assert.False(t, result, "block %d should not be cached after zero-length setIsCached", i)
+	}
+}
+
+func TestWriteAt_EmptyData(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Writing empty data should not mark any blocks as dirty
+	n, err := cache.WriteAt([]byte{}, 0)
+	require.NoError(t, err)
+	assert.Equal(t, 0, n)
+
+	// Block 0 should NOT be marked as cached
+	result := cache.isCached(0, blockSize)
+	assert.False(t, result, "empty WriteAt should not mark block 0 as dirty")
+}
+
+func TestWriteAt_PartiallyBeyondCacheSize(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Writing at the last block with data extending beyond cache size
+	// should only write up to cache size and mark appropriate blocks
+	data := make([]byte, blockSize*2) // 2 blocks of data
+	n, err := cache.WriteAt(data, blockSize*9) // Start at block 9, would extend to block 11
+	require.NoError(t, err)
+	assert.Equal(t, int(blockSize), n, "should only write up to cache size")
+
+	// Only block 9 should be marked as cached (the portion that fit)
+	for i := int64(0); i < 10; i++ {
+		if i == 9 {
+			assert.True(t, cache.isCached(i*blockSize, blockSize), "block 9 should be cached")
+		} else {
+			assert.False(t, cache.isCached(i*blockSize, blockSize), "block %d should not be cached", i)
+		}
+	}
+}
+
+func TestIsCached_OffsetBeyondSize(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Query beyond cache size should return false
+	result := cache.isCached(cacheSize+blockSize, blockSize)
+	assert.False(t, result, "isCached beyond cache size should return false")
+}
+
+func TestDirtyTracking_NormalOperation(t *testing.T) {
+	t.Parallel()
+
+	const blockSize = int64(header.PageSize)
+	const cacheSize = blockSize * 10
+
+	cache, err := NewCache(cacheSize, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	defer cache.Close()
+
+	// Initially nothing is cached
+	for i := int64(0); i < 10; i++ {
+		assert.False(t, cache.isCached(i*blockSize, blockSize), "block %d should not be initially cached", i)
+	}
+
+	// Write to block 3
+	data := make([]byte, blockSize)
+	n, err := cache.WriteAt(data, blockSize*3)
+	require.NoError(t, err)
+	assert.Equal(t, int(blockSize), n)
+
+	// Only block 3 should be cached
+	for i := int64(0); i < 10; i++ {
+		if i == 3 {
+			assert.True(t, cache.isCached(i*blockSize, blockSize), "block 3 should be cached")
+		} else {
+			assert.False(t, cache.isCached(i*blockSize, blockSize), "block %d should not be cached", i)
+		}
+	}
+
+	// Write spanning blocks 5-7
+	multiBlockData := make([]byte, blockSize*3)
+	n, err = cache.WriteAt(multiBlockData, blockSize*5)
+	require.NoError(t, err)
+	assert.Equal(t, int(blockSize*3), n)
+
+	// Blocks 3, 5, 6, 7 should be cached
+	expected := map[int64]bool{3: true, 5: true, 6: true, 7: true}
+	for i := int64(0); i < 10; i++ {
+		if expected[i] {
+			assert.True(t, cache.isCached(i*blockSize, blockSize), "block %d should be cached", i)
+		} else {
+			assert.False(t, cache.isCached(i*blockSize, blockSize), "block %d should not be cached", i)
+		}
+	}
+}