AI codebase review

fix: cap snapshot node count relative to file size

The 64x decompressed-body limit still let an 804 KB crafted snapshot
(repeating blocks of records) expand to 4.2 M nodes and ~500 MB of
heap, ~600x amplification; larger files exhaust memory. LoadFile now
allows at most one node per byte of file (1 Mi floor); real snapshots
use ~15 bytes per node. Limit errors on the node count are no longer
misreported as "no root node".

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

internal/snapshot/hostile_test.go

@@ -6,6 +6,8 @@ import (
 	"crypto/sha256"
 	"encoding/binary"
 	"io"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -129,3 +131,30 @@ func TestHostileMetadataSanitisedInSummary(t *testing.T) {
 		t.Fatalf("raw control characters in summary: %q", b.String())
 	}
 }
+
+// Blocks of one directory plus K-1 files repeat byte for byte, so a small
+// file describes millions of nodes; LoadFile must refuse before allocating.
+func TestNodeAmplificationLimited(t *testing.T) {
+	const k, n = 1000, 1_200_001
+	data := craft(n, `{"root":"/r"}`, func(i int, w *writer) {
+		switch {
+		case i == 0:
+			node(w, 0, inventory.KindDir, "", 0)
+		case i == 1:
+			node(w, 1, inventory.KindDir, "d", 0)
+		case (i-1)%k == 0:
+			node(w, k, inventory.KindDir, "d", 0)
+		default:
+			node(w, uint64((i-1)%k), inventory.KindFile, "f", 0)
+		}
+	})
+	if _, err := Load(bytes.NewReader(data)); err != nil {
+		t.Fatalf("crafted tree is otherwise valid: %v", err)
+	}
+	p := filepath.Join(t.TempDir(), "amp.msz")
+	os.WriteFile(p, data, 0o644)
+	_, err := LoadFile(p)
+	if err == nil || !strings.Contains(err.Error(), "node count") {
+		t.Fatalf("%d nodes from a %d-byte file accepted: %v", n, len(data), err)
+	}
+}

internal/snapshot/snapshot.go

@@ -333,12 +333,22 @@ func (l *limitedReader) Read(p []byte) (int, error) {
 // (real snapshots compress ~6×), with a 64 MiB floor for tiny files.
 func bodyLimit(fileSize int64) int64 { return max(64<<20, fileSize*64) }
 
+// nodeLimit caps the node count for a file of fileSize bytes. Real
+// snapshots spend ~15 compressed bytes per node; a crafted one describes a
+// node in a fraction of a byte, so without this a 1 MB file can demand
+// gigabytes of memory while staying within bodyLimit.
+func nodeLimit(fileSize int64) uint64 { return uint64(max(1<<20, fileSize)) }
+
 // Load reads and validates a snapshot, allowing up to DefaultMaxBody bytes
 // of decompressed data.
 func Load(in io.Reader) (*inventory.Tree, error) { return LoadLimit(in, DefaultMaxBody) }
 
 // LoadLimit is Load with an explicit decompressed-size limit.
 func LoadLimit(in io.Reader, maxBody int64) (*inventory.Tree, error) {
+	return load(in, maxBody, MaxNodes)
+}
+
+func load(in io.Reader, maxBody int64, maxNodes uint64) (*inventory.Tree, error) {
 	br := bufio.NewReaderSize(in, 1<<20)
 	hdr := make([]byte, len(brand.SnapshotMagic)+4)
 	if _, err := io.ReadFull(br, hdr); err != nil {
@@ -389,8 +399,11 @@ func LoadLimit(in io.Reader, maxBody int64) (*inventory.Tree, error) {
 		return nil, r.err
 	}
 	extIndex := t.ExtIndexer()
-	nNodes := r.uvarint(MaxNodes, "node count")
+	nNodes := r.uvarint(maxNodes, "node count")
 	depth := []uint16{0} // per node, for the depth limit; parents precede children
+	if r.err != nil {
+		return nil, r.err
+	}
 	if nNodes == 0 {
 		return nil, errors.New("snapshot: no root node")
 	}
@@ -507,9 +520,9 @@ func LoadFile(path string) (*inventory.Tree, error) {
 		return nil, err
 	}
 	defer f.Close()
-	limit := int64(DefaultMaxBody)
+	limit, nodes := int64(DefaultMaxBody), uint64(MaxNodes)
 	if fi, err := f.Stat(); err == nil {
-		limit = bodyLimit(fi.Size())
+		limit, nodes = bodyLimit(fi.Size()), nodeLimit(fi.Size())
 	}
-	return LoadLimit(f, limit)
+	return load(f, limit, nodes)
 }