fix: invalid mapping normalization

packages/orchestrator/cmd/inspect-header/main.go

@@ -48,6 +48,12 @@ func main() {
 		log.Fatalf("failed to deserialize header: %s", err)
 	}
 
+	// Validate mappings
+	err = header.ValidateMappings(h.Mapping, h.Metadata.Size, h.Metadata.BlockSize)
+	if err != nil {
+		fmt.Printf("\n⚠️  WARNING: Mapping validation failed!\n%s\n\n", err)
+	}
+
 	fmt.Printf("\nMETADATA\n")
 	fmt.Printf("========\n")
 	fmt.Printf("Storage            %s/%s\n", storage.GetDetails(), storagePath)

packages/orchestrator/internal/sandbox/sandbox.go

@@ -28,6 +28,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/metadata"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
@@ -833,6 +834,15 @@ func pauseProcessMemory(
 		return nil, nil, fmt.Errorf("failed to create memfile header: %w", err)
 	}
 
+	err = header.ValidateMappings(memfileHeader.Mapping, memfileHeader.Metadata.Size, memfileHeader.Metadata.BlockSize)
+	if err != nil {
+		if memfileHeader.IsNormalizeFixApplied() {
+			return nil, nil, fmt.Errorf("invalid memfile header mappings: %w", err)
+		}
+
+		zap.L().Warn("memfile header mappings are invalid, but normalize fix is not applied", zap.Error(err), logger.WithBuildID(memfileHeader.Metadata.BuildId.String()))
+	}
+
 	return memfileDiff, memfileHeader, nil
 }
 
@@ -889,6 +899,15 @@ func pauseProcessRootfs(
 		return nil, nil, fmt.Errorf("failed to create rootfs header: %w", err)
 	}
 
+	err = header.ValidateMappings(rootfsHeader.Mapping, rootfsHeader.Metadata.Size, rootfsHeader.Metadata.BlockSize)
+	if err != nil {
+		if rootfsHeader.IsNormalizeFixApplied() {
+			return nil, nil, fmt.Errorf("invalid rootfs header mappings: %w", err)
+		}
+
+		zap.L().Warn("rootfs header mappings are invalid, but normalize fix is not applied", zap.Error(err), logger.WithBuildID(rootfsHeader.Metadata.BuildId.String()))
+	}
+
 	return rootfsDiff, rootfsHeader, nil
 }
 

packages/shared/pkg/storage/header/header.go

@@ -5,8 +5,13 @@ import (
 
 	"github.com/bits-and-blooms/bitset"
 	"github.com/google/uuid"
+	"go.uber.org/zap"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
 
+const NormalizeFixVersion = 3
+
 type Header struct {
 	Metadata    *Metadata
 	blockStarts *bitset.BitSet
@@ -49,16 +54,61 @@ func NewHeader(metadata *Metadata, mapping []*BuildMap) (*Header, error) {
 	}, nil
 }
 
+// IsNormalizeFixApplied is a helper method to soft fail for older versions of the header where fix for normalization was not applied.
+// This should be removed in the future.
+func (t *Header) IsNormalizeFixApplied() bool {
+	return t.Metadata.Version >= NormalizeFixVersion
+}
+
 func (t *Header) GetShiftedMapping(offset int64) (mappedOffset int64, mappedLength int64, buildID *uuid.UUID, err error) {
 	mapping, shift, err := t.getMapping(offset)
 	if err != nil {
 		return 0, 0, nil, err
 	}
 
-	return int64(mapping.BuildStorageOffset) + shift, int64(mapping.Length) - shift, &mapping.BuildId, nil
+	mappedOffset = int64(mapping.BuildStorageOffset) + shift
+	mappedLength = int64(mapping.Length) - shift
+	buildID = &mapping.BuildId
+
+	if mappedLength < 0 {
+		if t.IsNormalizeFixApplied() {
+			return 0, 0, nil, fmt.Errorf("mapped length for offset %d is negative: %d", offset, mappedLength)
+		}
+
+		zap.L().Warn("mapped length is negative, but normalize fix is not applied",
+			zap.Int64("offset", offset),
+			zap.Int64("mappedLength", mappedLength),
+			logger.WithBuildID(mapping.BuildId.String()),
+		)
+	}
+
+	return mappedOffset, mappedLength, buildID, nil
 }
 
 func (t *Header) getMapping(offset int64) (*BuildMap, int64, error) {
+	if offset < 0 || offset >= int64(t.Metadata.Size) {
+		if t.IsNormalizeFixApplied() {
+			return nil, 0, fmt.Errorf("offset %d is out of bounds (size: %d)", offset, t.Metadata.Size)
+		}
+
+		zap.L().Warn("offset is out of bounds, but normalize fix is not applied",
+			zap.Int64("offset", offset),
+			zap.Int64("size", int64(t.Metadata.Size)),
+			logger.WithBuildID(t.Metadata.BuildId.String()),
+		)
+	}
+	if offset%int64(t.Metadata.BlockSize) != 0 {
+		if t.IsNormalizeFixApplied() {
+			return nil, 0, fmt.Errorf("offset %d is not aligned to block size %d", offset, t.Metadata.BlockSize)
+		}
+
+		zap.L().Warn("offset is not aligned to block size, but normalize fix is not applied",
+			zap.Int64("offset", offset),
+			zap.Int64("blockSize", int64(t.Metadata.BlockSize)),
+			logger.WithBuildID(t.Metadata.BuildId.String()),
+		)
+	}
+
 	block := BlockIdx(offset, int64(t.Metadata.BlockSize))
 
 	start, ok := t.blockStarts.PreviousSet(uint(block))
@@ -73,5 +123,21 @@ func (t *Header) getMapping(offset int64) (*BuildMap, int64, error) {
 
 	shift := (block - int64(start)) * int64(t.Metadata.BlockSize)
 
+	// Verify that the offset falls within this mapping's range
+	if shift >= int64(mapping.Length) {
+		if t.IsNormalizeFixApplied() {
+			return nil, 0, fmt.Errorf("offset %d (block %d) is beyond the end of mapping at offset %d (ends at %d)",
+				offset, block, mapping.Offset, mapping.Offset+mapping.Length)
+		}
+
+		zap.L().Warn("offset is beyond the end of mapping, but normalize fix is not applied",
+			zap.Int64("offset", offset),
+			zap.Int64("block", block),
+			zap.Uint64("mappingOffset", mapping.Offset),
+			zap.Uint64("mappingEnd", mapping.Offset+mapping.Length),
+			logger.WithBuildID(t.Metadata.BuildId.String()),
+		)
+	}
+
 	return mapping, shift, nil
 }

packages/shared/pkg/storage/header/mapping.go

@@ -19,6 +19,15 @@ type BuildMap struct {
 	BuildStorageOffset uint64
 }
 
+func (mapping *BuildMap) Copy() *BuildMap {
+	return &BuildMap{
+		Offset:             mapping.Offset,
+		Length:             mapping.Length,
+		BuildId:            mapping.BuildId,
+		BuildStorageOffset: mapping.BuildStorageOffset,
+	}
+}
+
 func CreateMapping(
 	buildId *uuid.UUID,
 	dirty *bitset.BitSet,
@@ -237,12 +246,29 @@ func MergeMappings(
 
 // NormalizeMappings joins adjacent mappings that have the same buildId.
 func NormalizeMappings(mappings []*BuildMap) []*BuildMap {
-	for i := 0; i < len(mappings); i++ {
-		if i+1 < len(mappings) && mappings[i].BuildId == mappings[i+1].BuildId {
-			mappings[i].Length += mappings[i+1].Length
-			mappings = append(mappings[:i+1], mappings[i+2:]...)
+	if len(mappings) == 0 {
+		return nil
+	}
+
+	result := make([]*BuildMap, 0, len(mappings))
+
+	// Start with a copy of the first mapping
+	current := mappings[0].Copy()
+
+	for i := 1; i < len(mappings); i++ {
+		mp := mappings[i]
+		if mp.BuildId != current.BuildId {
+			// BuildId changed, add the current map to results and start a new one
+			result = append(result, current)
+			current = mp.Copy() // New copy
+		} else {
+			// Same BuildId, just add the length
+			current.Length += mp.Length
 		}
 	}
 
-	return mappings
+	// Add the last mapping
+	result = append(result, current)
+
+	return result
 }