fix(snapshot): inflate object payloads into an exact-size buffer

commit: 351738de7a983ac81132dd2e297eef50dc2f3b90
author: brice · 14d ago · 2026-10-08 08:53 UTC
committer: brice · 14d ago · 2026-10-08 09:31 UTC
parent: f0493e8
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

2 files · +67 -10

M internal/snapshotcodec/object_compress.go +26 -10


@@ -58,18 +58,10 @@ func inflateObjectPayload(compressed []byte, size uint32) ([]byte, error) {
58	58		return nil, invalid(err)
59	59		}
60	60		defer objectDecompressorPool.Put(r)
61		-	out := bytes.NewBuffer(make([]byte, 0, min(int(size), inflateInitialCap)))
62		-	if _, err := io.Copy(out, io.LimitReader(r, int64(size))); err != nil {
	61	+	payload, err := readExact(r, int(size))
	62	+	if err != nil {
63	63		return nil, invalid(err)
64	64		}
65		-	if out.Len() != int(size) {
66		-	return nil, invalid(io.ErrUnexpectedEOF)
67		-	}
68		-	payload := out.Bytes()
69		-	if cap(payload) > len(payload)+len(payload)/4 {
70		-	// Restored payloads are retained; drop the growth slack.
71		-	payload = bytes.Clone(payload)
72		-	}
73	65		var extra [1]byte
74	66		if n, err := r.Read(extra[:]); n != 0 \|\| err != io.EOF {
75	67		if err == nil \|\| err == io.EOF {

@@ -83,6 +75,30 @@ func inflateObjectPayload(compressed []byte, size uint32) ([]byte, error) {
83	75		return payload, nil
84	76		}
85	77
	78	+	// readExact reads exactly size bytes. It doubles the buffer as data arrives,
	79	+	// starting at inflateInitialCap and capping every step at size, so a corrupt
	80	+	// declared length allocates at most about twice what the stream produces and
	81	+	// the returned payload has no growth slack.
	82	+	func readExact(r io.Reader, size int) ([]byte, error) {
	83	+	buf := make([]byte, min(size, inflateInitialCap))
	84	+	for n := 0; ; {
	85	+	m, err := io.ReadFull(r, buf[n:])
	86	+	n += m
	87	+	if err != nil {
	88	+	if err == io.EOF {
	89	+	err = io.ErrUnexpectedEOF
	90	+	}
	91	+	return nil, err
	92	+	}
	93	+	if n == size {
	94	+	return buf, nil
	95	+	}
	96	+	grown := make([]byte, min(2*len(buf), size))
	97	+	copy(grown, buf)
	98	+	buf = grown
	99	+	}
	100	+	}
	101	+
86	102		func objectDecompressor(source io.Reader) (io.ReadCloser, error) {
87	103		if pooled, ok := objectDecompressorPool.Get().(io.ReadCloser); ok {
88	104		if err := pooled.(zlib.Resetter).Reset(source, nil); err != nil {

 		return nil, invalid(err)
 	}
 	defer objectDecompressorPool.Put(r)
-	out := bytes.NewBuffer(make([]byte, 0, min(int(size), inflateInitialCap)))
+	payload, err := readExact(r, int(size))
-	if _, err := io.Copy(out, io.LimitReader(r, int64(size))); err != nil {
+	if err != nil {
 		return nil, invalid(err)
 	}
-	if out.Len() != int(size) {
-		return nil, invalid(io.ErrUnexpectedEOF)
-	}
-	payload := out.Bytes()
-	if cap(payload) > len(payload)+len(payload)/4 {
-		// Restored payloads are retained; drop the growth slack.
-		payload = bytes.Clone(payload)
-	}
 	var extra [1]byte
 	if n, err := r.Read(extra[:]); n != 0 || err != io.EOF {
 		if err == nil || err == io.EOF {
 	return payload, nil
 }
+// readExact reads exactly size bytes. It doubles the buffer as data arrives,
+// starting at inflateInitialCap and capping every step at size, so a corrupt
+// declared length allocates at most about twice what the stream produces and
+// the returned payload has no growth slack.
+func readExact(r io.Reader, size int) ([]byte, error) {
+	buf := make([]byte, min(size, inflateInitialCap))
+	for n := 0; ; {
+		m, err := io.ReadFull(r, buf[n:])
+		n += m
+		if err != nil {
+			if err == io.EOF {
+				err = io.ErrUnexpectedEOF
+			}
+			return nil, err
+		}
+		if n == size {
+			return buf, nil
+		}
+		grown := make([]byte, min(2*len(buf), size))
+		copy(grown, buf)
+		buf = grown
+	}
+}
 func objectDecompressor(source io.Reader) (io.ReadCloser, error) {
 	if pooled, ok := objectDecompressorPool.Get().(io.ReadCloser); ok {
 		if err := pooled.(zlib.Resetter).Reset(source, nil); err != nil {

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


@@ -5,6 +5,7 @@ import (
5	5		"encoding/binary"
6	6		"errors"
7	7		"hash/crc32"
	8	+	"runtime"
8	9		"testing"
9	10		)
10	11

@@ -65,3 +66,43 @@ func TestObjectCompressionRejectsMalformedPayloads(t *testing.T) {
65	66		})
66	67		}
67	68		}
	69	+
	70	+	// TestInflateObjectPayloadAllocation pins the inflate memory bound: a payload
	71	+	// costs about its own size, and a corrupt maximum declared length on a small
	72	+	// stream never allocates the declared size.
	73	+	func TestInflateObjectPayloadAllocation(t *testing.T) {
	74	+	for _, tc := range []struct {
	75	+	name string
	76	+	size int
	77	+	declared uint32
	78	+	maxBytes int64
	79	+	}{
	80	+	{name: "small payload", size: 64 << 10, maxBytes: 64<<10 + 64<<10},
	81	+	{name: "payload at the initial cap", size: inflateInitialCap, maxBytes: inflateInitialCap + 64<<10},
	82	+	{name: "payload above the initial cap", size: 3 << 20, maxBytes: 2*(3<<20) + 64<<10},
	83	+	{name: "declared maximum on a small stream", size: 64 << 10, declared: maxObjectPayloadSize, maxBytes: 2*inflateInitialCap + 64<<10},
	84	+	} {
	85	+	t.Run(tc.name, func(t *testing.T) {
	86	+	payload := bytes.Repeat([]byte("history row "), tc.size/12+1)[:tc.size]
	87	+	compressed, err := compressObjectPayload(nil, payload)
	88	+	if err != nil {
	89	+	t.Fatal(err)
	90	+	}
	91	+	declared := uint32(tc.size)
	92	+	if tc.declared != 0 {
	93	+	declared = tc.declared
	94	+	}
	95	+	_, _ = inflateObjectPayload(compressed, declared) // warm the decompressor pool
	96	+	const runs = 8
	97	+	var before, after runtime.MemStats
	98	+	runtime.ReadMemStats(&before)
	99	+	for range runs {
	100	+	_, _ = inflateObjectPayload(compressed, declared)
	101	+	}
	102	+	runtime.ReadMemStats(&after)
	103	+	if got := int64(after.TotalAlloc-before.TotalAlloc) / runs; got > tc.maxBytes {
	104	+	t.Fatalf("inflate allocated %d bytes/op, want at most %d", got, tc.maxBytes)
	105	+	}
	106	+	})
	107	+	}
	108	+	}


@@ -5,6 +5,7 @@ import (
5	"encoding/binary"	5		"encoding/binary"
6	"errors"	6		"errors"
7	"hash/crc32"	7		"hash/crc32"
		8	+	"runtime"
8	"testing"	9		"testing"
9	)	10		)
10		11

@@ -65,3 +66,43 @@ func TestObjectCompressionRejectsMalformedPayloads(t *testing.T) {
65	})	66		})
66	}	67		}
67	}	68		}
		69	+
		70	+	// TestInflateObjectPayloadAllocation pins the inflate memory bound: a payload
		71	+	// costs about its own size, and a corrupt maximum declared length on a small
		72	+	// stream never allocates the declared size.
		73	+	func TestInflateObjectPayloadAllocation(t *testing.T) {
		74	+	for _, tc := range []struct {
		75	+	name string
		76	+	size int
		77	+	declared uint32
		78	+	maxBytes int64
		79	+	}{
		80	+	{name: "small payload", size: 64 << 10, maxBytes: 64<<10 + 64<<10},
		81	+	{name: "payload at the initial cap", size: inflateInitialCap, maxBytes: inflateInitialCap + 64<<10},
		82	+	{name: "payload above the initial cap", size: 3 << 20, maxBytes: 2*(3<<20) + 64<<10},
		83	+	{name: "declared maximum on a small stream", size: 64 << 10, declared: maxObjectPayloadSize, maxBytes: 2*inflateInitialCap + 64<<10},
		84	+	} {
		85	+	t.Run(tc.name, func(t *testing.T) {
		86	+	payload := bytes.Repeat([]byte("history row "), tc.size/12+1)[:tc.size]
		87	+	compressed, err := compressObjectPayload(nil, payload)
		88	+	if err != nil {
		89	+	t.Fatal(err)
		90	+	}
		91	+	declared := uint32(tc.size)
		92	+	if tc.declared != 0 {
		93	+	declared = tc.declared
		94	+	}
		95	+	_, _ = inflateObjectPayload(compressed, declared) // warm the decompressor pool
		96	+	const runs = 8
		97	+	var before, after runtime.MemStats
		98	+	runtime.ReadMemStats(&before)
		99	+	for range runs {
		100	+	_, _ = inflateObjectPayload(compressed, declared)
		101	+	}
		102	+	runtime.ReadMemStats(&after)
		103	+	if got := int64(after.TotalAlloc-before.TotalAlloc) / runs; got > tc.maxBytes {
		104	+	t.Fatalf("inflate allocated %d bytes/op, want at most %d", got, tc.maxBytes)
		105	+	}
		106	+	})
		107	+	}
		108	+	}