Skip to content

[Bugfix][P2P-Store] Validate sizeList against the stored payload layout in GetReplica - #4304

Open
yaojiejia wants to merge 1 commit into
kvcache-ai:mainfrom
yaojiejia:fix/p2pstore-getreplica-layout-check
Open

yaojiejia wants to merge 1 commit into
kvcache-ai:mainfrom
yaojiejia:fix/p2pstore-getreplica-layout-check

Conversation

@yaojiejia

@yaojiejia yaojiejia commented Sep 23, 2026 •

Copy link
Copy Markdown

Description

GetReplica only checked that addrList and sizeList had the same non-zero length. It never compared sizeList with the layout stored in etcd for the payload, but doGetReplica and updatePayloadMetadata walk the caller's sizes in MaxShardSize steps and index payload.Shards by position, transferring the stored shard length into each address.

With a mismatched sizeList this meant:

  • more bytes than stored: payload.Shards[taskID] runs past the slice and the process panics;
  • a buffer shorter than a shard: the full stored shard length is written past the buffer's end (a plain out-of-bounds write on TCP; on RDMA it either fails or lands in neighbouring registered memory);
  • fewer total bytes than stored: only the first shards are fetched, GetReplica returns nil, and the node advertises itself as a replica source for a partial copy.

This PR adds validateReplicaLayout, which returns ErrInvalidArgument unless the caller's sizeList matches the stored SizeList entry for entry, MaxShardSize is non-zero, and the implied shard count equals len(payload.Shards). GetReplica calls it right after the first metadata read, before any memory is registered. updatePayloadMetadata calls it after re-reading the payload on a CAS miss, since that path walks the shard list the same way.

Fixes #4303

Module

  • Transfer Engine (mooncake-transfer-engine)
  • Mooncake Store (mooncake-store)
  • Mooncake Conductor (mooncake-conductor)
  • Reshard (mooncake-reshard)
  • Mooncake EP (mooncake-ep)
  • Mooncake PG (mooncake-pg)
  • Integration (mooncake-integration)
  • P2P Store (mooncake-p2p-store)
  • Python Wheel (mooncake-wheel)
  • Common (mooncake-common)
  • Mooncake RL (mooncake-rl)
  • CI/CD
  • Docs
  • Other

Type of Change

  • Bug fix
  • New feature
  • Refactor
  • Breaking change
  • Documentation update
  • Performance improvement
  • Other

How Has This Been Tested?

Added TestGetReplicaRejectsCallerLayoutMismatch in metadata_recheck_test.go, using the existing scriptedMetadata and successfulTransferEngine fakes. It covers six cases: larger than stored, smaller than stored, same total split differently, an extra buffer, a shard list shorter than the layout, and a zero MaxShardSize. Each case asserts ErrInvalidArgument, that no memory was registered, and that no metadata update was issued.

Against the unpatched core.go the new test fails with the crash this PR fixes:

--- FAIL: TestGetReplicaRejectsCallerLayoutMismatch/larger_than_stored
panic: runtime error: index out of range [1] with length 1

With the fix, the whole package passes (10 tests). gofmt -l reports nothing. go vet reports three pre-existing possible misuse of unsafe.Pointer warnings in transfer_engine.go, unrelated to this change.

Note: the package tests need the transfer engine C library to link; they were run against a link-only stub of the C API since the tests only use Go fakes. CI does not currently run the Go tests.

Test commands:

cd mooncake-p2p-store/src/p2pstore
gofmt -l .
go test -count=1 -v ./...

Test results:

  • Unit tests pass
  • Integration tests pass (if applicable)
  • Manual testing done (describe below)

Checklist

  • I have performed a self-review of my own code
  • I have formatted my code using ./scripts/code_format.sh
  • I have run pre-commit on the files changed in this PR and all hooks pass
  • I have updated the documentation (if applicable)
  • I have added tests to prove my changes are effective
  • For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

  • No AI tools were used
  • AI tools were used (specify below)

Claude Code was used to fix the regression test.

@yaojiejia yaojiejia changed the title [Bugfix][P2P-Store] Validate the caller's sizeList against the stored… [Bugfix][P2P-Store] Validate sizeList against the stored payload layout in GetReplica Sep 23, 2026

@he-yufeng he-yufeng left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Read it line by line and ran it in the dev container. This is right and complete.

The validator closes all three failure modes from the issue: exact per-entry equality against the stored SizeList rejects both over- and under-sized buffers before any registration or transfer, the MaxShardSize == 0 arm removes the divide-by-zero, and the shard-count identity catches the nastier case where stored SizeList and stored Shards disagree with each other (the truncatedShards case, which size equality alone would wave through). Placement is correct in both spots: GetReplica right after the first metadata.Get, and updatePayloadMetadata revalidates on the re-fetched payload after a conflict, closing the metadata-changed-mid-transfer window.

The tests assert the fail-closed ordering properly: no memory registration and no metadata update happen when the check fires, which is the part that keeps a rejected call from advertising a partial replica.

Verified locally (dev container, Go 1.25.7, freshly built transfer_engine/base/mooncake_common): TestGetReplicaRejectsCallerLayoutMismatch all six subcases pass on this head, as does the pre-existing layout-change suite; the new test fails on the parent commit without the fix (panic, as the issue predicts). CI is green on the PR.

One non-blocking note for a follow-up, not this PR: the shipped example still derives sizeList from a hardcoded file_size_mb rather than the layout from List(), so it exercises none of this. Not your problem to fix here, just worth recording.

Good to land from my reading.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: P2P Store GetReplica does not validate the caller's sizeList against the stored payload layout

2 participants