fix(snapshot): bound inflated object payloads on restore

Inflation allocated the declared raw length up front, so a corrupt object header could force a 256 MiB allocation. It now grows the buffer from a 1 MiB start as the stream produces bytes. Restoring one generation also caps its total inflated payloads at 256 MiB, the budget uncompressed checkpoints had, since the on-disk read cap no longer bounds memory.

commit: 03130e93c0f13011c53aceb3c78817ff0a45e313
author: brice · 14d ago · 2026-10-08 05:39 UTC
committer: brice · 14d ago · 2026-10-08 05:46 UTC
parent: 4604c71
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

4 files · +37 -6

M internal/snapshotcodec/codec.go +4 -0


@@ -34,6 +34,10 @@ const (
34	34		maxSnapshotDecodedAllocation = 256 << 20
35	35		)
36	36
	37	+	// MaxGenerationPayloadBytes bounds the inflated object payloads of one
	38	+	// incremental generation, the same budget an uncompressed checkpoint had.
	39	+	const MaxGenerationPayloadBytes = 256 << 20
	40	+
37	41		// Unmarshal first measures the complete manifest and its VT blobs without
38	42		// allocating. After accepting aggregate budgets, it validates references and
39	43		// blob semantics before decoding the session.


@@ -34,6 +34,10 @@ const (
34	maxSnapshotDecodedAllocation = 256 << 20	34		maxSnapshotDecodedAllocation = 256 << 20
35	)	35		)
36		36
		37	+	// MaxGenerationPayloadBytes bounds the inflated object payloads of one
		38	+	// incremental generation, the same budget an uncompressed checkpoint had.
		39	+	const MaxGenerationPayloadBytes = 256 << 20
		40	+
37	// Unmarshal first measures the complete manifest and its VT blobs without	41		// Unmarshal first measures the complete manifest and its VT blobs without
38	// allocating. After accepting aggregate budgets, it validates references and	42		// allocating. After accepting aggregate budgets, it validates references and
39	// blob semantics before decoding the session.	43		// blob semantics before decoding the session.

M internal/snapshotcodec/object_compress.go +12 -3


@@ -12,7 +12,12 @@ import (
12	12		// Object payloads are stored zlib-compressed. History cells are highly
13	13		// repetitive, so a sealed 10k-row history shrinks several times on disk.
14	14		// BestSpeed keeps snapshot encoding cheap; Go's deflate output is
15		-	// deterministic for one level, so content addressing stays stable.
	15	+	// deterministic for one toolchain and level, so content addressing stays
	16	+	// stable between checkpoints of one build.
	17	+
	18	+	// inflateInitialCap bounds the up-front allocation for one payload, so a
	19	+	// corrupt declared length cannot allocate more than the stream produces.
	20	+	const inflateInitialCap = 1 << 20
16	21
17	22		var objectCompressorPool = sync.Pool{New: func() any {
18	23		w, err := zlib.NewWriterLevel(io.Discard, zlib.BestSpeed)

@@ -53,10 +58,14 @@ func inflateObjectPayload(compressed []byte, size uint32) ([]byte, error) {
53	58		return nil, invalid(err)
54	59		}
55	60		defer objectDecompressorPool.Put(r)
56		-	payload := make([]byte, size)
57		-	if _, err := io.ReadFull(r, payload); err != nil {
	61	+	out := bytes.NewBuffer(make([]byte, 0, min(int(size), inflateInitialCap)))
	62	+	if _, err := io.Copy(out, io.LimitReader(r, int64(size))); err != nil {
58	63		return nil, invalid(err)
59	64		}
	65	+	if out.Len() != int(size) {
	66	+	return nil, invalid(io.ErrUnexpectedEOF)
	67	+	}
	68	+	payload := out.Bytes()
60	69		var extra [1]byte
61	70		if n, err := r.Read(extra[:]); n != 0 \|\| err != io.EOF {
62	71		if err == nil \|\| err == io.EOF {

 // Object payloads are stored zlib-compressed. History cells are highly
 // repetitive, so a sealed 10k-row history shrinks several times on disk.
 // BestSpeed keeps snapshot encoding cheap; Go's deflate output is
-// deterministic for one level, so content addressing stays stable.
+// deterministic for one toolchain and level, so content addressing stays
+// stable between checkpoints of one build.
+// inflateInitialCap bounds the up-front allocation for one payload, so a
+// corrupt declared length cannot allocate more than the stream produces.
+const inflateInitialCap = 1 << 20
 var objectCompressorPool = sync.Pool{New: func() any {
 	w, err := zlib.NewWriterLevel(io.Discard, zlib.BestSpeed)
 		return nil, invalid(err)
 	}
 	defer objectDecompressorPool.Put(r)
-	payload := make([]byte, size)
+	out := bytes.NewBuffer(make([]byte, 0, min(int(size), inflateInitialCap)))
-	if _, err := io.ReadFull(r, payload); err != nil {
+	if _, err := io.Copy(out, io.LimitReader(r, int64(size))); err != nil {
 		return nil, invalid(err)
 	}
+	if out.Len() != int(size) {
+		return nil, invalid(io.ErrUnexpectedEOF)
+	}
+	payload := out.Bytes()
 	var extra [1]byte
 	if n, err := r.Read(extra[:]); n != 0 || err != io.EOF {
 		if err == nil || err == io.EOF {

M internal/snapshotcodec/object_compress_test.go +4 -0


@@ -46,6 +46,10 @@ func TestObjectCompressionRejectsMalformedPayloads(t *testing.T) {
46	46		binary.BigEndian.PutUint32(b[1:5], uint32(len(payload)+1))
47	47		return b
48	48		}, want: ErrInvalidData},
	49	+	{name: "declared maximum on a small stream", mutate: func(b []byte) []byte {
	50	+	binary.BigEndian.PutUint32(b[1:5], maxObjectPayloadSize)
	51	+	return b
	52	+	}, want: ErrInvalidData},
49	53		{name: "truncated stream", mutate: func(b []byte) []byte { return b[:len(b)-1] }, want: ErrInvalidData},
50	54		{name: "corrupt stream", mutate: func(b []byte) []byte { b[len(b)-6] ^= 0xff; return b }, want: ErrInvalidData},
51	55		{name: "trailing garbage", mutate: func(b []byte) []byte { return append(b, 0) }, want: ErrTrailingBytes},


@@ -46,6 +46,10 @@ func TestObjectCompressionRejectsMalformedPayloads(t *testing.T) {
46	binary.BigEndian.PutUint32(b[1:5], uint32(len(payload)+1))	46		binary.BigEndian.PutUint32(b[1:5], uint32(len(payload)+1))
47	return b	47		return b
48	}, want: ErrInvalidData},	48		}, want: ErrInvalidData},
		49	+	{name: "declared maximum on a small stream", mutate: func(b []byte) []byte {
		50	+	binary.BigEndian.PutUint32(b[1:5], maxObjectPayloadSize)
		51	+	return b
		52	+	}, want: ErrInvalidData},
49	{name: "truncated stream", mutate: func(b []byte) []byte { return b[:len(b)-1] }, want: ErrInvalidData},	53		{name: "truncated stream", mutate: func(b []byte) []byte { return b[:len(b)-1] }, want: ErrInvalidData},
50	{name: "corrupt stream", mutate: func(b []byte) []byte { b[len(b)-6] ^= 0xff; return b }, want: ErrInvalidData},	54		{name: "corrupt stream", mutate: func(b []byte) []byte { b[len(b)-6] ^= 0xff; return b }, want: ErrInvalidData},
51	{name: "trailing garbage", mutate: func(b []byte) []byte { return append(b, 0) }, want: ErrTrailingBytes},	55		{name: "trailing garbage", mutate: func(b []byte) []byte { return append(b, 0) }, want: ErrTrailingBytes},

M internal/usecase/daemon/snapshot_restore_incremental.go +17 -3


@@ -395,21 +395,35 @@ func sessionFromGeneration(generation ports.SnapshotGeneration) (snapcodec.Sessi
395	395		return snapcodec.Session{}, fmt.Errorf("snapshot: generation identity mismatch")
396	396		}
397	397		result := snapcodec.Session{Name: manifest.Name, CreatedAt: manifest.CreatedAt, Active: manifest.Active, Tabs: make([]snapcodec.Tab, 0, len(manifest.Tabs))}
	398	+	// Objects are compressed, so the on-disk read budget no longer bounds the
	399	+	// inflated generation; budget its payloads as one uncompressed checkpoint.
	400	+	budget := snapcodec.MaxGenerationPayloadBytes
	401	+	object := func(ref snapcodec.ObjectRef, kind snapcodec.ObjectKind) ([]byte, error) {
	402	+	data, err := generationObject(generation, ref, kind)
	403	+	if err != nil {
	404	+	return nil, err
	405	+	}
	406	+	if len(data) > budget {
	407	+	return nil, fmt.Errorf("snapshot: generation payload too large")
	408	+	}
	409	+	budget -= len(data)
	410	+	return data, nil
	411	+	}
398	412		for _, tab := range manifest.Tabs {
399	413		outTab := snapcodec.Tab{StableID: tab.StableID, Cols: tab.Cols, Rows: tab.Rows, NextPaneID: tab.NextPaneID, Focus: tab.Focus, Tree: tab.Tree, Panes: make([]snapcodec.Pane, 0, len(tab.Panes))}
400	414		for _, pane := range tab.Panes {
401	415		outPane := snapcodec.Pane{ID: pane.ID, StableID: pane.StableID, Cwd: pane.Cwd, Process: pane.Process}
402	416		for _, ref := range pane.Sealed {
403		-	data, err := generationObject(generation, ref, snapcodec.HistoryChunk)
	417	+	data, err := object(ref, snapcodec.HistoryChunk)
404	418		if err != nil {
405	419		return snapcodec.Session{}, err
406	420		}
407	421		outPane.SealedChunks = append(outPane.SealedChunks, data)
408	422		}
409		-	if outPane.Tail, err = generationObject(generation, pane.Tail, snapcodec.HistoryTail); err != nil {
	423	+	if outPane.Tail, err = object(pane.Tail, snapcodec.HistoryTail); err != nil {
410	424		return snapcodec.Session{}, err
411	425		}
412		-	if outPane.Transcript, err = generationObject(generation, pane.Transcript, snapcodec.RecoveryTranscript); err != nil {
	426	+	if outPane.Transcript, err = object(pane.Transcript, snapcodec.RecoveryTranscript); err != nil {
413	427		return snapcodec.Session{}, err
414	428		}
415	429		outTab.Panes = append(outTab.Panes, outPane)

 		return snapcodec.Session{}, fmt.Errorf("snapshot: generation identity mismatch")
 	}
 	result := snapcodec.Session{Name: manifest.Name, CreatedAt: manifest.CreatedAt, Active: manifest.Active, Tabs: make([]snapcodec.Tab, 0, len(manifest.Tabs))}
+	// Objects are compressed, so the on-disk read budget no longer bounds the
+	// inflated generation; budget its payloads as one uncompressed checkpoint.
+	budget := snapcodec.MaxGenerationPayloadBytes
+	object := func(ref snapcodec.ObjectRef, kind snapcodec.ObjectKind) ([]byte, error) {
+		data, err := generationObject(generation, ref, kind)
+		if err != nil {
+			return nil, err
+		}
+		if len(data) > budget {
+			return nil, fmt.Errorf("snapshot: generation payload too large")
+		}
+		budget -= len(data)
+		return data, nil
+	}
 	for _, tab := range manifest.Tabs {
 		outTab := snapcodec.Tab{StableID: tab.StableID, Cols: tab.Cols, Rows: tab.Rows, NextPaneID: tab.NextPaneID, Focus: tab.Focus, Tree: tab.Tree, Panes: make([]snapcodec.Pane, 0, len(tab.Panes))}
 		for _, pane := range tab.Panes {
 			outPane := snapcodec.Pane{ID: pane.ID, StableID: pane.StableID, Cwd: pane.Cwd, Process: pane.Process}
 			for _, ref := range pane.Sealed {
-				data, err := generationObject(generation, ref, snapcodec.HistoryChunk)
+				data, err := object(ref, snapcodec.HistoryChunk)
 				if err != nil {
 					return snapcodec.Session{}, err
 				}
 				outPane.SealedChunks = append(outPane.SealedChunks, data)
 			}
-			if outPane.Tail, err = generationObject(generation, pane.Tail, snapcodec.HistoryTail); err != nil {
+			if outPane.Tail, err = object(pane.Tail, snapcodec.HistoryTail); err != nil {
 				return snapcodec.Session{}, err
 			}
-			if outPane.Transcript, err = generationObject(generation, pane.Transcript, snapcodec.RecoveryTranscript); err != nil {
+			if outPane.Transcript, err = object(pane.Transcript, snapcodec.RecoveryTranscript); err != nil {
 				return snapcodec.Session{}, err
 			}
 			outTab.Panes = append(outTab.Panes, outPane)