From 28aaa0e114c1d1928d529a1aa6ab67226daab59a Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 2 Oct 2026 17:55:51 +0800 Subject: [PATCH 1/2] [Bugfix][TransferEngine] Bound the handshake daemon's pending-notify queue The handshake daemon accepts notify frames from any peer that can reach the RPC port, and receivePeerNotify appended every one of them to an unbounded in-memory vector. A trivial flood grows the host's RSS without limit (#4445). Cap the queue at MC_HANDSHAKE_MAX_NOTIFY_ENTRIES (default 1024, max 1M) and evict the oldest undrained entry once full, so a flood degrades pending notifications instead of the process. Warn at most once per 1000 evictions to keep the flood from turning into log spam. Adds a socket-level regression test that floods cap + 64 frames and verifies the oldest 64 are evicted while the survivors keep arrival order. Fixes #4445 Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- mooncake-transfer-engine/include/common.h | 31 ++++++++++++ .../src/transfer_metadata.cpp | 11 +++++ .../tests/transfer_metadata_test.cpp | 49 +++++++++++++++++++ 3 files changed, 91 insertions(+) diff --git a/mooncake-transfer-engine/include/common.h b/mooncake-transfer-engine/include/common.h index 63fabdfdb9..7bbd13373f 100644 --- a/mooncake-transfer-engine/include/common.h +++ b/mooncake-transfer-engine/include/common.h @@ -541,6 +541,37 @@ static inline size_t getHandshakeMaxLength() { return max_length; } +// Constants for the handshake daemon's pending-notify bound. +constexpr size_t kDefaultHandshakeMaxNotifyEntries = 1024; +constexpr size_t kMaxHandshakeMaxNotifyEntries = 1ULL << 20; + +// Load the pending-notify bound from the environment variable. +static inline size_t loadHandshakeMaxNotifyEntries() { + const char *env = std::getenv("MC_HANDSHAKE_MAX_NOTIFY_ENTRIES"); + if (env != nullptr) { + std::string_view env_sv(env); + size_t val = 0; + auto [ptr, ec] = + std::from_chars(env_sv.data(), env_sv.data() + env_sv.size(), val); + if (ec == std::errc() && ptr == env_sv.data() + env_sv.size() && + val >= 1 && val <= kMaxHandshakeMaxNotifyEntries) { + LOG(INFO) << "MC_HANDSHAKE_MAX_NOTIFY_ENTRIES set to " << val; + return val; + } + LOG(WARNING) << "Invalid MC_HANDSHAKE_MAX_NOTIFY_ENTRIES value: " << env + << ", valid range: 1 to " << kMaxHandshakeMaxNotifyEntries + << ", using default " << kDefaultHandshakeMaxNotifyEntries; + } + return kDefaultHandshakeMaxNotifyEntries; +} + +// Get the pending-notify bound. +// Loaded once from MC_HANDSHAKE_MAX_NOTIFY_ENTRIES at first call. +static inline size_t getHandshakeMaxNotifyEntries() { + static size_t max_entries = loadHandshakeMaxNotifyEntries(); + return max_entries; +} + static inline std::pair readString(int fd) { HandShakeRequestType type = HandShakeRequestType::Invalid; diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index a8f87e63fb..c12d94d317 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -402,6 +402,17 @@ int TransferMetadata::receivePeerNotify(const Json::Value &peer_json, RWSpinlock::WriteGuard guard(notify_lock_); TransferMetadata::NotifyDesc peer_notify, local_reply; TransferNotifyUtil::decode(peer_json, peer_notify); + const size_t max_entries = getHandshakeMaxNotifyEntries(); + if (notifys.size() >= max_entries) { + // The daemon accepts notify frames from any peer, so an unbounded + // vector lets a flood grow RSS without limit (#4445). Evict the + // oldest undrained entry: a flood degrades pending notifications + // instead of the process. + notifys.erase(notifys.begin()); + LOG_EVERY_N(WARNING, 1000) + << "pending notify queue is full (" << max_entries + << " entries); evicting the oldest entry"; + } notifys.push_back(peer_notify); // reply local_reply.name = ""; diff --git a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp index 1002f3d9df..18bf4b8c61 100644 --- a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp +++ b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp @@ -1057,6 +1057,55 @@ TEST(HandshakeFrameTest, RejectsInvalidLength) { close(fds[1]); } +// Every notify frame the daemon accepts lands in an in-memory vector that +// used to have no bound, so any peer that can reach the RPC port could grow +// the host's RSS forever (#4445). The queue now caps at +// getHandshakeMaxNotifyEntries() and evicts the oldest undrained entry once +// full. Flood cap + 64 frames over the real socket, then drain. +TEST(HandshakeNotifyTest, BoundedQueueEvictsOldestUnderFlood) { + TransferMetadata server(P2PHANDSHAKE); + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_NE(port, 0); + + TransferMetadata::RpcMetaDesc rpc{}; + rpc.ip_or_host_name = "127.0.0.1"; + rpc.rpc_port = port; + rpc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry("notify-bounded-test", rpc), 0); + + const size_t cap = getHandshakeMaxNotifyEntries(); + const size_t total = cap + 64; + for (size_t i = 0; i < total; ++i) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + ASSERT_NE(fd, -1); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + ASSERT_EQ( + connect(fd, reinterpret_cast(&addr), sizeof(addr)), 0); + Json::Value body; + body["name"] = "atk"; + body["notify_msg"] = "m" + std::to_string(i); + ASSERT_EQ(writeString(fd, HandShakeRequestType::Notify, + Json::FastWriter{}.write(body)), + 0); + // The daemon answers each accepted frame before half-closing, so + // reading the reply orders frame i before frame i+1. + char reply[128]; + ASSERT_GT(read(fd, reply, sizeof(reply)), 0); + close(fd); + } + + std::vector drained; + ASSERT_EQ(server.getNotifies(drained), 0); + ASSERT_EQ(drained.size(), cap); + // The oldest 64 frames were evicted; the survivors keep arrival order. + EXPECT_EQ(drained.front().notify_msg, "m64"); + EXPECT_EQ(drained.back().notify_msg, "m" + std::to_string(total - 1)); +} + namespace { // Hardware-free in-memory MetadataStoragePlugin. get() returns false for keys From 74f04e7a6b53d0b90df21df854f61dbc74ca2e4d Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 2 Oct 2026 20:31:36 +0800 Subject: [PATCH 2/2] test: read the full notify reply frame before closing The flood loop's bare read() can return after the frame header alone. Closing then resets the socket while the daemon is still writing the payload, and its next write() kills the test process with SIGPIPE, as seen on the CTest leg. Drain the frame with readString instead, which also asserts the daemon answered with a Notify reply. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- .../tests/transfer_metadata_test.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp index 18bf4b8c61..2ebb1eefc0 100644 --- a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp +++ b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp @@ -1092,9 +1092,13 @@ TEST(HandshakeNotifyTest, BoundedQueueEvictsOldestUnderFlood) { Json::FastWriter{}.write(body)), 0); // The daemon answers each accepted frame before half-closing, so - // reading the reply orders frame i before frame i+1. - char reply[128]; - ASSERT_GT(read(fd, reply, sizeof(reply)), 0); + // reading the reply orders frame i before frame i+1. Read the whole + // frame: a bare read() can return after the header alone, and closing + // while the daemon is still writing the payload resets the socket + // and kills the process with SIGPIPE on its next write. + auto [rtype, rjson] = readString(fd); + ASSERT_EQ(rtype, HandShakeRequestType::Notify); + ASSERT_FALSE(rjson.empty()); close(fd); }