[Store] fix: prevent cross-process memcpy segfault when MC_STORE_MEMCPY auto-enables

mooncake-store/include/transfer_task.h

@@ -410,8 +410,24 @@ class TransferSubmitter {
         const std::unordered_map<std::string, std::vector<Slice>>&
             batched_slices);
 
+    /**
+     * @brief Pure comparison helper: returns true iff both endpoints are
+     * non-empty and identical. Exposed for unit testing of the locality
+     * decision without instantiating a full TransferEngine.
+     *
+     * Two endpoints identify the same process only when their ip:port (or
+     * full hostname) match exactly; same-host different-process pairs share
+     * an IP but not a port and must NOT be treated as locally addressable.
+     */
+    static bool isSameProcessEndpoint(const std::string& handle_endpoint,
+                                      const std::string& local_endpoint);
+
    private:
     TransferEngine& engine_;
+    // Cached at construction: the local transport endpoint never changes for
+    // the lifetime of the TransferSubmitter, so we avoid calling
+    // engine_.getLocalIpAndPort() (which allocates a string) on every transfer.
+    const std::string local_endpoint_;
     std::unique_ptr<MemcpyWorkerPool> memcpy_pool_;
     std::unique_ptr<FilereadWorkerPool> fileread_pool_;
     bool memcpy_enabled_;

mooncake-store/src/transfer_task.cpp

@@ -451,6 +451,7 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine,
                                      const std::string& local_hostname,
                                      TransferMetric* transfer_metric)
     : engine_(engine),
+      local_endpoint_(engine.getLocalIpAndPort()),
       memcpy_pool_(std::make_unique<MemcpyWorkerPool>()),
       fileread_pool_(std::make_unique<FilereadWorkerPool>(backend)),
       local_hostname_(local_hostname),
@@ -816,28 +817,67 @@ TransferStrategy TransferSubmitter::selectStrategy(
     return TransferStrategy::TRANSFER_ENGINE;
 }
 
-bool TransferSubmitter::isLocalTransfer(
-    const AllocatedBuffer::Descriptor& handle) const {
-    if (handle.transport_endpoint_.empty()) return false;
+namespace {
+// Helper function to extract IP address from endpoint string (ip:port format).
+// Supports both IPv4 (ip:port) and IPv6 ([ipv6]:port) formats.
+std::string extractIpAddress(const std::string& endpoint) {
+    if (endpoint.empty()) {
+        return "";
+    }
 
-    // Metadata-service descriptors use the client hostname as the segment ID.
-    // If it matches this client's hostname, the buffer address is local.
-    if (!local_hostname_.empty() &&
-        local_hostname_ == handle.transport_endpoint_) {
-        return true;
+    // Handle IPv6 format: [ipv6]:port
+    if (endpoint[0] == '[') {
+        size_t closing_bracket = endpoint.find(']');
+        if (closing_bracket == std::string::npos) {
+            LOG(WARNING) << "Invalid IPv6 endpoint format: " << endpoint;
+            return "";
+        }
+        return endpoint.substr(1, closing_bracket - 1);
+    }
+
+    // Handle IPv4 or hostname:port format.
+    size_t colon_pos = endpoint.rfind(':');
+    if (colon_pos != std::string::npos) {
+        return endpoint.substr(0, colon_pos);
     }
 
-    // P2P descriptors use the transfer engine endpoint as the segment ID.
-    // If it matches this engine's endpoint, the buffer address is local.
-    std::string local_ep = engine_.getLocalIpAndPort();
-    if (!local_ep.empty() && handle.transport_endpoint_ == local_ep) {
+    // No colon found, return the whole string (might be just IP or hostname).
+    return endpoint;
+}
+}  // namespace
+
+bool TransferSubmitter::isSameProcessEndpoint(
+    const std::string& handle_endpoint, const std::string& local_endpoint) {
+    // Local memcpy requires that handle.buffer_address_ is a virtual address
+    // valid in THIS process. Same host is not enough: two processes on the
+    // same host share an IP but have distinct virtual address spaces, so a
+    // memcpy on a peer process's address would segfault. Require the full
+    // transport endpoint to match, which uniquely identifies the owning
+    // process.
+    if (handle_endpoint.empty() || local_endpoint.empty()) {
+        return false;
+    }
+    if (handle_endpoint == local_endpoint) {
         return true;
     }
 
-    // Without a local endpoint we cannot prove locality; disable memcpy.
+    const std::string handle_ip = extractIpAddress(handle_endpoint);
+    const std::string local_ip = extractIpAddress(local_endpoint);
+    if (!handle_ip.empty() && handle_ip == local_ip) {
+        VLOG(2) << "Disabling local memcpy for same-host endpoints with "
+                   "different process endpoints: handle="
+                << handle_endpoint << ", local=" << local_endpoint;
+    }
+
     return false;
 }
 
+bool TransferSubmitter::isLocalTransfer(
+    const AllocatedBuffer::Descriptor& handle) const {
+    return isSameProcessEndpoint(handle.transport_endpoint_, local_hostname_) ||
+           isSameProcessEndpoint(handle.transport_endpoint_, local_endpoint_);
+}
+
 bool TransferSubmitter::validateTransferParams(
     const AllocatedBuffer::Descriptor& handle,
     const std::vector<Slice>& slices) const {

mooncake-store/tests/transfer_task_test.cpp

@@ -144,6 +144,37 @@ TEST_F(TransferTaskTest, MemcpyWorkerPoolMultipleOperations) {
     }
 }
 
+// Test the locality decision used by TransferSubmitter::isLocalTransfer.
+// Same-host different-process pairs share an IP but have distinct ports;
+// they must NOT be treated as locally addressable, otherwise memcpy in the
+// caller process would dereference a virtual address belonging to a peer
+// process and segfault.
+TEST_F(TransferTaskTest, IsSameProcessEndpoint) {
+    // Empty inputs -> not same-process (cannot prove locality).
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint("", ""));
+    EXPECT_FALSE(
+        TransferSubmitter::isSameProcessEndpoint("", "192.168.1.10:12345"));
+    EXPECT_FALSE(
+        TransferSubmitter::isSameProcessEndpoint("192.168.1.10:12345", ""));
+
+    // Identical ip:port -> same process.
+    EXPECT_TRUE(TransferSubmitter::isSameProcessEndpoint("192.168.1.10:12345",
+                                                         "192.168.1.10:12345"));
+
+    // Same host, different port -> different process, NOT local.
+    // This is the regression case fixed by this change.
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint(
+        "192.168.1.10:12345", "192.168.1.10:12346"));
+
+    // Different hosts -> not local.
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint(
+        "192.168.1.10:12345", "192.168.1.11:12345"));
+
+    // Hostname endpoints (non-P2P metadata mode) compare as full strings.
+    EXPECT_TRUE(TransferSubmitter::isSameProcessEndpoint("host-a", "host-a"));
+    EXPECT_FALSE(TransferSubmitter::isSameProcessEndpoint("host-a", "host-b"));
+}
+
 // Test TransferStrategy enum and stream operator
 TEST_F(TransferTaskTest, TransferStrategyEnum) {
     // Test enum values

scripts/run_tests.sh

@@ -40,7 +40,11 @@ DUMMY_TEST_PID_2=$!
 wait $DUMMY_TEST_PID_1 $DUMMY_TEST_PID_2
 kill $CLIENT_PID || true
 
-pip install torch numpy safetensors packaging
+pip install numpy safetensors packaging
+# Keep the test torch aligned with the EP/PG variants packaged into the CI wheel.
+pip install "${MOONCAKE_TEST_TORCH_SPEC:-torch==2.11.0+cu128}" \
+    --index-url "${MOONCAKE_TEST_TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu128}" \
+    --extra-index-url https://pypi.org/simple
 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_put_get_tensor.py
 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_safetensor_functions.py
 kill $MASTER_PID || true