address #3232 review: narrow fix to Snapshot clone

The shared *Snapshot pointer is the real cause of #3231: the shallow
copy of the raft message meant re-slicing raftMsg.Snapshot.Data also
re-sliced m.Snapshot.Data, shrinking its capacity until a later
iteration panicked with "slice bounds out of range".

Cloning the Snapshot struct alone is sufficient. Drop the 64 KiB
payload-size clamp (which could not produce a sendable message when
struct overhead already exceeds GRPCMaxMsgSize) and the redundant
fullData capture. Replace the tests with a single regression test
that feeds a realistic MsgSnap of 3*GRPCMaxMsgSize and asserts the
input message's Snapshot.Data is untouched after splitting.

Signed-off-by: Sylvere Richard <sylvere.richard@gmail.com>

Author: nowheresly

manager/state/raft/transport/peer.go

@@ -145,15 +145,8 @@ func raftMessageStructSize(m *raftpb.Message) int {
 
 // Returns the max allowable payload based on MaxRaftMsgSize and
 // the struct size for the given raftpb.Message.
-// If the struct overhead exceeds GRPCMaxMsgSize (e.g. very large
-// snapshot metadata), the payload size is clamped to a minimum of
-// 64 KiB to guarantee forward progress when splitting.
 func raftMessagePayloadSize(m *raftpb.Message) int {
-	s := GRPCMaxMsgSize - raftMessageStructSize(m)
-	if s < 64<<10 {
-		s = 64 << 10
-	}
-	return s
+	return GRPCMaxMsgSize - raftMessageStructSize(m)
 }
 
 // Split a large raft message into smaller messages.
@@ -171,11 +164,6 @@ func splitSnapshotData(_ context.Context, m *raftpb.Message) []api.StreamRaftMes
 	// Get the max payload size.
 	payloadSize := raftMessagePayloadSize(m)
 
-	// Capture the full snapshot data before the loop, because
-	// we need to copy the Snapshot struct for each chunk to avoid
-	// mutating m.Snapshot.Data through the shared pointer.
-	fullData := m.Snapshot.Data
-
 	// split the snapshot into smaller messages.
 	for snapDataIndex := 0; snapDataIndex < size; {
 		chunkSize := size - snapDataIndex
@@ -184,13 +172,13 @@ func splitSnapshotData(_ context.Context, m *raftpb.Message) []api.StreamRaftMes
 		}
 
 		raftMsg := *m
-		// Copy the Snapshot to avoid mutating the original message's
-		// Snapshot.Data through the shared pointer.
+		// Clone Snapshot so that re-slicing Snapshot.Data below
+		// does not mutate m.Snapshot.Data through the shared pointer.
 		snap := *m.Snapshot
 		raftMsg.Snapshot = &snap
 
 		// sub-slice for this snapshot chunk.
-		raftMsg.Snapshot.Data = fullData[snapDataIndex : snapDataIndex+chunkSize]
+		raftMsg.Snapshot.Data = m.Snapshot.Data[snapDataIndex : snapDataIndex+chunkSize]
 
 		snapDataIndex += chunkSize
 

manager/state/raft/transport/transport_test.go

@@ -105,71 +105,49 @@ func testSend(ctx context.Context, c *mockCluster, from uint64, to []uint64, msg
 	}
 }
 
-func TestSplitSnapshotData(t *testing.T) {
+// TestSplitSnapshotDataDoesNotMutateInput is a regression test for #3231.
+// Before the fix, splitSnapshotData did a shallow copy of the raft message
+// and re-sliced the shared Snapshot.Data on each iteration, shrinking the
+// original slice's capacity and eventually panicking with
+// "slice bounds out of range".
+func TestSplitSnapshotDataDoesNotMutateInput(t *testing.T) {
 	ctx := context.Background()
 
-	t.Run("NormalSnapshot", func(t *testing.T) {
-		m := newSnapshotMessage(1, 2)
-		msgs := splitSnapshotData(ctx, &m)
-		assert.NotEmpty(t, msgs)
-		// reassemble and verify
-		var assembled []byte
-		for _, msg := range msgs {
-			assembled = append(assembled, msg.Message.Snapshot.Data...)
-		}
-		assert.Equal(t, m.Snapshot.Data, assembled)
-	})
-
-	t.Run("LargeMetadataDoesNotPanic", func(t *testing.T) {
-		// Simulate a snapshot where struct overhead exceeds GRPCMaxMsgSize.
-		// This is the scenario from the bug: many cluster objects cause
-		// large raft message metadata, making payloadSize negative without
-		// the fix.
-		data := make([]byte, 5<<20) // 5 MiB snapshot data
-		for i := range data {
-			data[i] = byte(i % 256)
-		}
-		m := raftpb.Message{
-			Type: raftpb.MsgSnap,
-			From: 1,
-			To:   2,
-			Snapshot: &raftpb.Snapshot{
-				Data: data,
-				Metadata: raftpb.SnapshotMetadata{
-					Index: uint64(len(data)),
-					// Large ConfState to push struct size above GRPCMaxMsgSize
-					ConfState: raftpb.ConfState{
-						Voters: make([]uint64, 1<<20),
-					},
-				},
-			},
-		}
-		// Must not panic
-		msgs := splitSnapshotData(ctx, &m)
-		assert.NotEmpty(t, msgs)
-		// reassemble and verify
-		var assembled []byte
-		for _, msg := range msgs {
-			assembled = append(assembled, msg.Message.Snapshot.Data...)
-		}
-		assert.Equal(t, data, assembled)
-	})
-}
-
-func TestRaftMessagePayloadSizeMinimum(t *testing.T) {
-	// When struct overhead exceeds GRPCMaxMsgSize, payloadSize must
-	// be clamped to the minimum (64 KiB) instead of going negative.
-	// Use Entries to inflate the struct size well beyond GRPCMaxMsgSize.
-	bigEntry := raftpb.Entry{Data: make([]byte, GRPCMaxMsgSize)}
-	m := &raftpb.Message{
-		Type:    raftpb.MsgSnap,
-		Entries: []raftpb.Entry{bigEntry},
+	// Build a MsgSnap whose Snapshot.Data clearly exceeds GRPCMaxMsgSize so
+	// that the split loop runs multiple iterations (where the bug manifests).
+	const dataSize = 3 * GRPCMaxMsgSize
+	data := make([]byte, dataSize)
+	for i := range data {
+		data[i] = byte(i % (1 << 8))
+	}
+	m := raftpb.Message{
+		Type: raftpb.MsgSnap,
+		From: 1,
+		To:   2,
 		Snapshot: &raftpb.Snapshot{
-			Data: nil,
+			Data: data,
+			Metadata: raftpb.SnapshotMetadata{
+				Index: uint64(len(data)),
+			},
 		},
 	}
-	ps := raftMessagePayloadSize(m)
-	assert.Equal(t, 64<<10, ps, "payloadSize should be clamped to minimum 64 KiB")
+	origData := m.Snapshot.Data
+	origLen, origCap := len(origData), cap(origData)
+
+	msgs := splitSnapshotData(ctx, &m)
+	require.Greater(t, len(msgs), 1, "data larger than GRPCMaxMsgSize must split into multiple chunks")
+
+	// Chunks must reassemble to the original data.
+	var assembled []byte
+	for _, msg := range msgs {
+		assembled = append(assembled, msg.Message.Snapshot.Data...)
+	}
+	assert.Equal(t, data, assembled)
+
+	// The input message's Snapshot.Data must be untouched (regression guard).
+	assert.Equal(t, origLen, len(m.Snapshot.Data))
+	assert.Equal(t, origCap, cap(m.Snapshot.Data))
+	assert.Equal(t, data, m.Snapshot.Data)
 }
 
 func TestSend(t *testing.T) {