fix(kv): bound batch decoding allocations

commit: 7f3e337f1b9b6eba2805b34d4fec7afce83c85e0
author: brice · 2mo ago · 2026-07-25 20:17 UTC
parent: 3faaee9
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

2 files · +23 -5

M pkg/kv/codec.go +10 -5


@@ -12,11 +12,12 @@ const (
12	12		opDel
13	13		opBatch
14	14
15		-	headerLen = 8
16		-	payloadPrefixLen = 3
17		-	compactThreshold = 64
18		-	compactWasteRatio = 0.5
19		-	maxKeyLen = 1<<16 - 1
	15	+	headerLen = 8
	16	+	payloadPrefixLen = 3
	17	+	batchEntryPrefixLen = 7
	18	+	compactThreshold = 64
	19	+	compactWasteRatio = 0.5
	20	+	maxKeyLen = 1<<16 - 1
20	21		)
21	22
22	23		var (

@@ -151,6 +152,10 @@ func decodeBatch(payload []byte) ([]record, error) {
151	152		return nil, errBadRecord
152	153		}
153	154		off := 5
	155	+	maxCount := uint64(len(payload)-off) / batchEntryPrefixLen
	156	+	if uint64(count) > maxCount {
	157	+	return nil, errBadRecord
	158	+	}
154	159		records := make([]record, 0, int(count))
155	160		seen := make(map[string]struct{}, int(count))
156	161		for range count {

 	opDel
 	opBatch
-	headerLen         = 8
+	headerLen           = 8
-	payloadPrefixLen  = 3
+	payloadPrefixLen    = 3
-	compactThreshold  = 64
+	batchEntryPrefixLen = 7
-	compactWasteRatio = 0.5
+	compactThreshold    = 64
-	maxKeyLen         = 1<<16 - 1
+	compactWasteRatio   = 0.5
+	maxKeyLen           = 1<<16 - 1
 )
 var (
 		return nil, errBadRecord
 	}
 	off := 5
+	maxCount := uint64(len(payload)-off) / batchEntryPrefixLen
+	if uint64(count) > maxCount {
+		return nil, errBadRecord
+	}
 	records := make([]record, 0, int(count))
 	seen := make(map[string]struct{}, int(count))
 	for range count {

M pkg/kv/kv_test.go +13 -0


@@ -256,6 +256,7 @@ func TestReplayStrict(t *testing.T) {
256	256		{name: "final-crc", wal: corruptRecord(append(append([]byte{}, valid...), batch...), len(valid)+4), corrupt: true},
257	257		{name: "middle-crc", wal: corruptRecord(append(append([]byte{}, valid...), batch...), 4), corrupt: true},
258	258		{name: "invalid-op", wal: rawRecord([]byte{0xff, 0, 0}), corrupt: true},
	259	+	{name: "impossible-batch-count", wal: rawRecord([]byte{opBatch, 0xff, 0xff, 0xff, 0xff}), corrupt: true},
259	260		{name: "duplicate-key-in-batch", wal: rawRecord(rawBatchPayload([]BatchChange{{Key: []byte("x")}, {Key: []byte("x"), Delete: true}})), corrupt: true},
260	261		{name: "trailing-batch-garbage", wal: rawRecord(append(rawBatchPayload([]BatchChange{{Key: []byte("x")}}), 0)), corrupt: true},
261	262		}

@@ -293,6 +294,18 @@ func TestReplayStrict(t *testing.T) {
293	294		}
294	295		}
295	296
	297	+	func TestDecodeBatchRejectsImpossibleCountBeforeAllocating(t *testing.T) {
	298	+	payload := []byte{opBatch, 0xff, 0xff, 0xff, 0xff}
	299	+	allocs := testing.AllocsPerRun(10, func() {
	300	+	if _, err := decodeBatch(payload); !errors.Is(err, errBadRecord) {
	301	+	t.Fatalf("decodeBatch error = %v, want errBadRecord", err)
	302	+	}
	303	+	})
	304	+	if allocs != 0 {
	305	+	t.Fatalf("decodeBatch allocations = %v, want 0", allocs)
	306	+	}
	307	+	}
	308	+
296	309		func TestBatchAtomicity(t *testing.T) {
297	310		path := privatePath(t)
298	311		s, err := Open(path)


@@ -256,6 +256,7 @@ func TestReplayStrict(t *testing.T) {
256	{name: "final-crc", wal: corruptRecord(append(append([]byte{}, valid...), batch...), len(valid)+4), corrupt: true},	256		{name: "final-crc", wal: corruptRecord(append(append([]byte{}, valid...), batch...), len(valid)+4), corrupt: true},
257	{name: "middle-crc", wal: corruptRecord(append(append([]byte{}, valid...), batch...), 4), corrupt: true},	257		{name: "middle-crc", wal: corruptRecord(append(append([]byte{}, valid...), batch...), 4), corrupt: true},
258	{name: "invalid-op", wal: rawRecord([]byte{0xff, 0, 0}), corrupt: true},	258		{name: "invalid-op", wal: rawRecord([]byte{0xff, 0, 0}), corrupt: true},
		259	+	{name: "impossible-batch-count", wal: rawRecord([]byte{opBatch, 0xff, 0xff, 0xff, 0xff}), corrupt: true},
259	{name: "duplicate-key-in-batch", wal: rawRecord(rawBatchPayload([]BatchChange{{Key: []byte("x")}, {Key: []byte("x"), Delete: true}})), corrupt: true},	260		{name: "duplicate-key-in-batch", wal: rawRecord(rawBatchPayload([]BatchChange{{Key: []byte("x")}, {Key: []byte("x"), Delete: true}})), corrupt: true},
260	{name: "trailing-batch-garbage", wal: rawRecord(append(rawBatchPayload([]BatchChange{{Key: []byte("x")}}), 0)), corrupt: true},	261		{name: "trailing-batch-garbage", wal: rawRecord(append(rawBatchPayload([]BatchChange{{Key: []byte("x")}}), 0)), corrupt: true},
261	}	262		}

@@ -293,6 +294,18 @@ func TestReplayStrict(t *testing.T) {
293	}	294		}
294	}	295		}
295		296
		297	+	func TestDecodeBatchRejectsImpossibleCountBeforeAllocating(t *testing.T) {
		298	+	payload := []byte{opBatch, 0xff, 0xff, 0xff, 0xff}
		299	+	allocs := testing.AllocsPerRun(10, func() {
		300	+	if _, err := decodeBatch(payload); !errors.Is(err, errBadRecord) {
		301	+	t.Fatalf("decodeBatch error = %v, want errBadRecord", err)
		302	+	}
		303	+	})
		304	+	if allocs != 0 {
		305	+	t.Fatalf("decodeBatch allocations = %v, want 0", allocs)
		306	+	}
		307	+	}
		308	+
296	func TestBatchAtomicity(t *testing.T) {	309		func TestBatchAtomicity(t *testing.T) {
297	path := privatePath(t)	310		path := privatePath(t)
298	s, err := Open(path)	311		s, err := Open(path)