backup: harden snapshot encoder core per M1 review (single-use guard + edge tests)

claude review on PR #825 found no blocking issues; addressed the two
meaningful findings plus the cheap nits to harden the foundation the
M2-M5 adapter feed loops build on:

- WriteTo is now single-use: a second call fails closed with
  ErrSnapshotBuilderReused rather than silently re-emitting the
  already-written entries (a footgun once adapters feed the builder
  in loops). Caller audit: the only callers of WriteTo/newSnapshotBuilder
  are tests; the production caller lands with the M6 CLI, so no
  existing call site relies on reuse.
- Add now size-checks from the user-buffer lengths BEFORE allocating
  the framed key/value, so an oversize record fails closed without a
  wasted allocation.
- countingWriter.Write uses the idiomatic if-err-nil guard for
  consistency with the rest of the package.

Tests: empty-builder header-only round-trip (fresh-cluster restore
path), WriteTo reuse rejection, and a deterministic WriteTo byte-count
assertion in the single-entry round-trip. Extracted assertWriteByteCount
/ assertSingleEntry helpers to keep the round-trip test under the cyclop
budget.

The gemini medium findings (unbounded slice / seen-map memory, reflect
in binary.Write) are below the in-memory-sort bound the design already
tracks as the external-sort follow-up milestone; not addressed here.

Author: bootjp

internal/backup/encode.go

@@ -53,6 +53,14 @@ var ErrEncodeDuplicateKey = errors.New("backup: duplicate encoded key in snapsho
 var ErrEncodeKeyTooLarge = errors.New("backup: encoded key length exceeds limit")
 var ErrEncodeValueTooLarge = errors.New("backup: encoded value length exceeds limit")
 
+// ErrSnapshotBuilderReused is returned by WriteTo when it is called
+// more than once on the same builder. A builder is single-use (one
+// per encode run); a second WriteTo would re-emit the already-written
+// entries, producing a valid-but-unintended stream. Enforced so the
+// per-adapter feed loops in later milestones cannot silently double-
+// emit (claude review on PR #825).
+var ErrSnapshotBuilderReused = errors.New("backup: snapshotBuilder.WriteTo called more than once")
+
 // ErrLastCommitTSRegression is returned by ResolveCommitTS when a
 // `--last-commit-ts` override is older than MANIFEST.last_commit_ts.
 // Seeding the restored node's HLC ceiling below a timestamp already
@@ -123,6 +131,7 @@ type snapshotBuilder struct {
 	commitTS uint64
 	entries  []encodedKV
 	seen     map[string]struct{}
+	written  bool
 }
 
 // newSnapshotBuilder constructs a builder that stamps every key with
@@ -141,21 +150,24 @@ func newSnapshotBuilder(commitTS uint64) *snapshotBuilder {
 // order-dependent `.fsm`. userKey/userValue are copied — callers may
 // reuse their buffers after Add returns.
 func (b *snapshotBuilder) Add(userKey, userValue []byte, expireAt uint64) error {
-	key := encodeMVCCKey(userKey, b.commitTS)
-	if uint64(len(key)) > MaxSnapshotEncodedKeySize {
+	// Size-check from the user-buffer lengths before allocating the
+	// framed buffers — encKey = userKey + snapshotTSSize, encVal =
+	// snapshotValueHeaderSize + userValue — so an oversize record
+	// fails closed without a wasted allocation (claude review nit).
+	if uint64(len(userKey))+snapshotTSSize > MaxSnapshotEncodedKeySize {
 		return errors.Wrapf(ErrEncodeKeyTooLarge,
-			"length %d > %d", len(key), MaxSnapshotEncodedKeySize)
+			"length %d > %d", len(userKey)+snapshotTSSize, MaxSnapshotEncodedKeySize)
 	}
-	val := encodeMVCCValue(userValue, expireAt)
-	if uint64(len(val)) > MaxSnapshotEncodedValueSize {
+	if uint64(len(userValue))+snapshotValueHeaderSize > MaxSnapshotEncodedValueSize {
 		return errors.Wrapf(ErrEncodeValueTooLarge,
-			"length %d > %d", len(val), MaxSnapshotEncodedValueSize)
+			"length %d > %d", len(userValue)+snapshotValueHeaderSize, MaxSnapshotEncodedValueSize)
 	}
+	key := encodeMVCCKey(userKey, b.commitTS)
 	if _, dup := b.seen[string(key)]; dup {
 		return errors.Wrapf(ErrEncodeDuplicateKey, "userKey %q", userKey)
 	}
 	b.seen[string(key)] = struct{}{}
-	b.entries = append(b.entries, encodedKV{key: key, val: val})
+	b.entries = append(b.entries, encodedKV{key: key, val: encodeMVCCValue(userValue, expireAt)})
 	return nil
 }
 
@@ -169,6 +181,10 @@ func (b *snapshotBuilder) Len() int { return len(b.entries) }
 // format terminates on a clean EOF (store/snapshot_pebble.go WriteTo,
 // internal/backup/snapshot_reader.go ReadSnapshot).
 func (b *snapshotBuilder) WriteTo(w io.Writer) (int64, error) {
+	if b.written {
+		return 0, errors.WithStack(ErrSnapshotBuilderReused)
+	}
+	b.written = true
 	sort.Slice(b.entries, func(i, j int) bool {
 		return bytes.Compare(b.entries[i].key, b.entries[j].key) < 0
 	})
@@ -214,7 +230,10 @@ type countingWriter struct {
 func (w *countingWriter) Write(p []byte) (int, error) {
 	n, err := w.w.Write(p)
 	w.n += int64(n)
-	return n, errors.WithStack(err)
+	if err != nil {
+		return n, errors.WithStack(err)
+	}
+	return n, nil
 }
 
 // RoundTripEntry is one live record recovered by DecodeLiveEntries.

internal/backup/encode_test.go

@@ -101,9 +101,16 @@ func TestSnapshotBuilderRoundTrip(t *testing.T) {
 	}
 
 	var buf bytes.Buffer
-	if _, err := b.WriteTo(&buf); err != nil {
+	n, err := b.WriteTo(&buf)
+	if err != nil {
 		t.Fatalf("WriteTo: %v", err)
 	}
+	// Byte count is deterministic: magic + lastCommitTS + per-entry
+	// (8-byte len prefix + bytes) for both key and value.
+	wantN := int64(len(PebbleSnapshotMagic) + 8 +
+		(8 + len(userKey) + snapshotTSSize) +
+		(8 + snapshotValueHeaderSize + len(userValue)))
+	assertWriteByteCount(t, n, &buf, wantN)
 
 	entries, hdr, err := DecodeLiveEntries(&buf)
 	if err != nil {
@@ -112,17 +119,35 @@ func TestSnapshotBuilderRoundTrip(t *testing.T) {
 	if hdr.LastCommitTS != encodeFixtureTS {
 		t.Fatalf("header LastCommitTS = %#x, want %#x", hdr.LastCommitTS, encodeFixtureTS)
 	}
-	if len(entries) != 1 {
-		t.Fatalf("got %d entries, want 1", len(entries))
+	assertSingleEntry(t, entries, userKey, userValue, expireAt)
+}
+
+// assertWriteByteCount checks WriteTo's returned count matches the
+// expected value and the actual buffer length.
+func assertWriteByteCount(t *testing.T, n int64, buf *bytes.Buffer, want int64) {
+	t.Helper()
+	if n != want {
+		t.Fatalf("WriteTo n = %d, want %d", n, want)
 	}
-	if !bytes.Equal(entries[0].UserKey, userKey) {
-		t.Fatalf("UserKey = %q, want %q", entries[0].UserKey, userKey)
+	if n != int64(buf.Len()) {
+		t.Fatalf("WriteTo n = %d, but buf has %d bytes", n, buf.Len())
 	}
-	if !bytes.Equal(entries[0].UserValue, userValue) {
-		t.Fatalf("UserValue = %q, want %q", entries[0].UserValue, userValue)
+}
+
+// assertSingleEntry checks the decoded set has exactly one entry
+// matching the given key/value/expiry.
+func assertSingleEntry(t *testing.T, entries []RoundTripEntry, key, val []byte, exp uint64) {
+	t.Helper()
+	if len(entries) != 1 {
+		t.Fatalf("got %d entries, want 1", len(entries))
 	}
-	if entries[0].ExpireAt != expireAt {
-		t.Fatalf("ExpireAt = %d, want %d", entries[0].ExpireAt, expireAt)
+	switch {
+	case !bytes.Equal(entries[0].UserKey, key):
+		t.Fatalf("UserKey = %q, want %q", entries[0].UserKey, key)
+	case !bytes.Equal(entries[0].UserValue, val):
+		t.Fatalf("UserValue = %q, want %q", entries[0].UserValue, val)
+	case entries[0].ExpireAt != exp:
+		t.Fatalf("ExpireAt = %d, want %d", entries[0].ExpireAt, exp)
 	}
 }
 
@@ -202,6 +227,48 @@ func TestSnapshotBuilderRejectsOversizeKey(t *testing.T) {
 	}
 }
 
+// TestSnapshotBuilderEmptyRoundTrip pins the empty-builder path (the
+// "fresh cluster / empty shard" restore case): WriteTo emits a valid
+// header-only stream that DecodeLiveEntries reads back as zero entries
+// with the header timestamp intact.
+func TestSnapshotBuilderEmptyRoundTrip(t *testing.T) {
+	t.Parallel()
+	b := newSnapshotBuilder(encodeFixtureTS)
+	var buf bytes.Buffer
+	if _, err := b.WriteTo(&buf); err != nil {
+		t.Fatalf("WriteTo: %v", err)
+	}
+	entries, hdr, err := DecodeLiveEntries(&buf)
+	if err != nil {
+		t.Fatalf("DecodeLiveEntries: %v", err)
+	}
+	if hdr.LastCommitTS != encodeFixtureTS {
+		t.Fatalf("hdr.LastCommitTS = %#x, want %#x", hdr.LastCommitTS, encodeFixtureTS)
+	}
+	if len(entries) != 0 {
+		t.Fatalf("got %d entries, want 0", len(entries))
+	}
+}
+
+// TestSnapshotBuilderRejectsReuse pins that a builder is single-use:
+// a second WriteTo fails closed with ErrSnapshotBuilderReused rather
+// than silently re-emitting the already-written entries.
+func TestSnapshotBuilderRejectsReuse(t *testing.T) {
+	t.Parallel()
+	b := newSnapshotBuilder(encodeFixtureTS)
+	if err := b.Add([]byte("!redis|str|k"), []byte("v"), 0); err != nil {
+		t.Fatalf("Add: %v", err)
+	}
+	var buf bytes.Buffer
+	if _, err := b.WriteTo(&buf); err != nil {
+		t.Fatalf("first WriteTo: %v", err)
+	}
+	_, err := b.WriteTo(&buf)
+	if !errors.Is(err, ErrSnapshotBuilderReused) {
+		t.Fatalf("second WriteTo err = %v, want ErrSnapshotBuilderReused", err)
+	}
+}
+
 // TestResolveCommitTS pins the --last-commit-ts override semantics
 // (design §"MVCC re-encoding"): no override yields the manifest value;
 // an override >= manifest is accepted; an override < manifest fails