fix(replication): confine transfer-service object ids to the ML Flashpoint base directory

docs/user-guide.md

@@ -188,7 +188,9 @@ Also create an `MLFlashpointCheckpointLoader` with those dependencies, and use t
 # Initialize dependencies (shared singletons)
 checkpoint_object_manager = CheckpointObjectManager()
 replication_manager = ReplicationManager()
-replication_manager.initialize(checkpoint_object_manager)
+# base_dir confines replication to the ML Flashpoint base path: peers can only
+# read and write checkpoint objects inside it.
+replication_manager.initialize(checkpoint_object_manager, base_dir=mlflashpoint_base_path)
 
 checkpoint_loader = DefaultMLFlashpointCheckpointLoader(
     checkpoint_object_manager=checkpoint_object_manager,
@@ -262,7 +264,9 @@ from ml_flashpoint.replication.replication_manager import ReplicationManager
 # Initialize dependencies (shared singletons)
 checkpoint_object_manager = CheckpointObjectManager()
 replication_manager = ReplicationManager()
-replication_manager.initialize(checkpoint_object_manager)
+# base_dir confines replication to the ML Flashpoint base path: peers can only
+# read and write checkpoint objects inside it.
+replication_manager.initialize(checkpoint_object_manager, base_dir=mlflashpoint_base_path)
 
 # Instantiate the StorageWriter
 memory_storage_writer = MemoryStorageWriter(

src/ml_flashpoint/adapter/nemo/wrapper_util.py

@@ -98,7 +98,7 @@ def wrap_trainer_and_auto_resume_with_mlflashpoint(
 
     ckpt_obj_manager = CheckpointObjectManager(pool_config=pool_config)
     replication_manager = ReplicationManager()
-    replication_manager.initialize(checkpoint_object_manager=ckpt_obj_manager)
+    replication_manager.initialize(checkpoint_object_manager=ckpt_obj_manager, base_dir=str(flashpoint_base_container))
 
     ckpt_loader = NeMoMLFlashpointCheckpointLoader(
         checkpoint_object_manager=ckpt_obj_manager,

src/ml_flashpoint/replication/replication_manager.py

@@ -15,11 +15,12 @@
 import abc
 import concurrent.futures
 import logging
+import os
 import socket
 import threading
 import time
 from pathlib import Path
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -192,6 +193,7 @@ def initialize(
         retry_config: ReplicationRetryConfig = get_default_retry_config(),
         replication_transfer_service: Optional[transfer_service_ext.TransferService] = None,
         repl_strategy: Optional[ReplicationStrategy] = None,
+        base_dir: Optional[Union[str, os.PathLike]] = None,
     ) -> None:
         """Initializes the ReplicationManager and underlying services.
 
@@ -208,7 +210,24 @@ def initialize(
                 created.
             repl_strategy: An optional replication strategy. If not provided, a
                 default PairwiseReplicationStrategy will be used.
+            base_dir: The ML Flashpoint base directory. The TransferService
+                only reads and writes checkpoint objects inside it: a peer that
+                asks for, or sends, an object id that resolves outside it (after
+                following symlinks) gets an error. Required unless
+                `replication_transfer_service` is provided, in which case that
+                service's own base directory applies and this argument is
+                ignored.
+
+        Raises:
+            ValueError: If `base_dir` is missing and no
+                `replication_transfer_service` is provided.
+            RuntimeError: If the TransferService fails to initialize.
         """
+        if replication_transfer_service is None and not base_dir:
+            raise ValueError(
+                "base_dir is required when no replication_transfer_service is provided: "
+                "the TransferService only transfers objects inside it."
+            )
         self._checkpoint_object_manager = checkpoint_object_manager
         self._listen_port = listen_port
         self._retry_config = retry_config
@@ -220,7 +239,8 @@ def initialize(
             _LOGGER.info("No TransferService provided, initializing a new one...")
             self._transfer_service = transfer_service_ext.TransferService()
             bound_listen_port = self._transfer_service.initialize(
-                listen_port,
+                os.fspath(base_dir),
+                listen_port=listen_port,
                 global_rank=dist.get_rank(),
             )
             if bound_listen_port <= 0:

src/ml_flashpoint/replication/transfer_service/bindings.cpp

@@ -35,13 +35,19 @@ PYBIND11_MODULE(transfer_service_ext, m) {
       .def(py::init<>())
 
       // Synchronous methods remain unchanged (still blocking)
-      .def("initialize", &TransferService::Initialize, py::arg("listen_port"),
-           py::arg("threads") = 16, py::arg("conn_pool_per_peer") = 16,
-           py::arg("global_rank") = -1,
-           "Initializes and starts the C++ transfer service. `threads` is the "
-           "size of each of the two worker pools, one for outbound tasks "
-           "(async_put/async_get) and one for inbound requests from peers, so "
-           "the service runs 2 * threads worker threads.")
+      .def(
+          "initialize", &TransferService::Initialize, py::arg("base_dir"),
+          py::arg("listen_port") = 0, py::arg("threads") = 16,
+          py::arg("conn_pool_per_peer") = 16, py::arg("global_rank") = -1,
+          "Initializes and starts the C++ transfer service. `base_dir` is the "
+          "directory that confines every object id the service reads for a "
+          "peer or writes on its behalf: an id that resolves outside it, after "
+          "following symlinks, is rejected with an error. It must not be "
+          "empty and is created if it does not exist. `threads` is the size "
+          "of each of the two worker pools, one for outbound tasks "
+          "(async_put/async_get) and one for inbound requests from peers, so "
+          "the service runs 2 * threads worker threads. Returns the port the "
+          "service listens on, or -1 on failure.")
 
       .def("shutdown", &TransferService::Shutdown,
            "Stops the C++ transer service.")

src/ml_flashpoint/replication/transfer_service/transfer_helpers.cpp

@@ -15,9 +15,17 @@
 #include "transfer_helpers.h"
 
 #include <array>
+#include <filesystem>
 #include <iomanip>
 #include <random>
 #include <sstream>
+#include <string>
+#include <string_view>
+#include <system_error>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 
 namespace ml_flashpoint::replication::transfer_service {
 
@@ -49,4 +57,37 @@ std::string GenerateUuid() {
   return ss.str();
 }
 
+absl::StatusOr<std::string> ResolvePathWithinBaseDir(
+    std::string_view object_id, const std::filesystem::path& base_dir) {
+  if (object_id.empty()) {
+    return absl::InvalidArgumentError("Object id is empty");
+  }
+  std::error_code ec;
+  // weakly_canonical leaves a relative path relative when none of its
+  // components exist yet, so make it absolute first.
+  const std::filesystem::path absolute =
+      std::filesystem::absolute(std::filesystem::path(object_id), ec);
+  if (ec) {
+    return absl::InternalError(absl::StrCat("Cannot resolve object id '",
+                                            object_id, "': ", ec.message()));
+  }
+  // Follows symlinks in the existing prefix, then normalizes the rest, so a
+  // destination that does not exist yet is handled too.
+  const std::filesystem::path resolved =
+      std::filesystem::weakly_canonical(absolute, ec);
+  if (ec) {
+    return absl::InternalError(absl::StrCat("Cannot resolve object id '",
+                                            object_id, "': ", ec.message()));
+  }
+  // Component-wise comparison: a string prefix check would accept
+  // "/base2/x" for the base "/base".
+  const std::filesystem::path relative = resolved.lexically_relative(base_dir);
+  if (relative.empty() || relative == "." || *relative.begin() == "..") {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Object id '", object_id, "' resolves to '", resolved.native(),
+        "', outside the base directory '", base_dir.native(), "'"));
+  }
+  return resolved.native();
+}
+
 }  // namespace ml_flashpoint::replication::transfer_service

src/ml_flashpoint/replication/transfer_service/transfer_helpers.h

@@ -17,12 +17,34 @@
 #ifndef ML_FLASHPOINT_REPLICATION_TRANSFER_HELPERS_H_
 #define ML_FLASHPOINT_REPLICATION_TRANSFER_HELPERS_H_
 
+#include <filesystem>
 #include <string>
+#include <string_view>
+
+#include "absl/status/statusor.h"
 
 namespace ml_flashpoint::replication::transfer_service {
 
 std::string GenerateUuid();
 
+// Resolves an object id to the path the service may use for it.
+//
+// The id is made absolute (a relative id is resolved against the process's
+// current working directory, as open(2) would), symlinks in the part of it
+// that exists are followed, and the rest is normalized lexically, so "..",
+// "." and symlinks cannot lead outside the base directory. `base_dir` itself
+// is not a valid object path.
+//
+// `base_dir` must be absolute and canonical, as returned by
+// std::filesystem::canonical, so both sides of the comparison name the real
+// location.
+//
+// Returns the resolved path, or InvalidArgument when the id is empty or
+// resolves outside `base_dir`, or another error when the filesystem lookup
+// fails (for example, a name that is too long).
+absl::StatusOr<std::string> ResolvePathWithinBaseDir(
+    std::string_view object_id, const std::filesystem::path& base_dir);
+
 }  // namespace ml_flashpoint::replication::transfer_service
 
 #endif  // ML_FLASHPOINT_REPLICATION_TRANSFER_HELPERS_H_

src/ml_flashpoint/replication/transfer_service/transfer_service.cpp

@@ -34,6 +34,7 @@
 #include <shared_mutex>
 #include <stdexcept>
 #include <string>
+#include <system_error>
 
 #include "absl/log/check.h"
 #include "absl/log/globals.h"
@@ -81,8 +82,9 @@ TransferService::~TransferService() {
   absl::RemoveLogSink(mlf_log_sink_.get());
 }
 
-int TransferService::Initialize(int listen_port, int threads,
-                                int conn_pool_size_per_peer, int global_rank) {
+int TransferService::Initialize(const std::string& base_dir, int listen_port,
+                                int threads, int conn_pool_size_per_peer,
+                                int global_rank) {
   std::lock_guard<std::mutex> lock(init_mutex_);
   if (running_.load()) {
     LOG(INFO) << "Transfer service is already running.";
@@ -99,6 +101,28 @@ int TransferService::Initialize(int listen_port, int threads,
   mlf_log_sink_ = std::make_unique<MLFLogSink>(global_rank);
   absl::AddLogSink(mlf_log_sink_.get());
 
+  // Resolve the directory that confines every object id. There is no default
+  // on purpose: a service that is not told where it may read and write must
+  // not start. canonical() follows symlinks and needs the directory to exist,
+  // so object ids are compared against the directory's real location.
+  if (base_dir.empty()) {
+    LOG(ERROR) << "base_dir must not be empty";
+    return -1;
+  }
+  std::error_code ec;
+  std::filesystem::create_directories(base_dir, ec);
+  if (!ec) {
+    base_dir_ = std::filesystem::canonical(base_dir, ec);
+  }
+  if (!ec && !std::filesystem::is_directory(base_dir_, ec)) {
+    ec = std::make_error_code(std::errc::not_a_directory);
+  }
+  if (ec) {
+    LOG(ERROR) << "Failed to resolve base_dir '" << base_dir
+               << "': " << ec.message();
+    return -1;
+  }
+
   // Get listener fd and bind to listen port
   absl::StatusOr<int> listener_fd_or = SetupListeningSocket(listen_port);
   if (!listener_fd_or.ok()) {
@@ -154,7 +178,8 @@ int TransferService::Initialize(int listen_port, int threads,
   LOG(INFO) << "TransferService initialized and listening on port "
             << listen_port_ << " with threads=" << threads_
             << ", conn_pool_per_peer=" << conn_pool_size_per_peer_
-            << ", global_rank=" << global_rank_;
+            << ", global_rank=" << global_rank_
+            << ", base_dir=" << base_dir_.native();
   return listen_port_;
 }
 
@@ -399,6 +424,17 @@ std::future<TransferResult> TransferService::AsyncGet(
     pending_tasks_[task_id] = {promise, metric_container};
   }
 
+  // Check the local destination now, before anything is sent, so that a bad
+  // id fails fast instead of after a full transfer. HandleDataReceive checks
+  // it again when the data arrives: that is the one place every write goes
+  // through.
+  if (absl::StatusOr<std::string> dest =
+          ResolvePathWithinBaseDir(dest_obj_id, base_dir_);
+      !dest.ok()) {
+    ReportResult(task_id, false, std::string(dest.status().message()));
+    return future;
+  }
+
   task_queue_.enqueue(std::move(task));
 
   return future;
@@ -801,6 +837,18 @@ void TransferService::HandleDataReceive(int client_fd,
     fail("Invalid obj_size received: " + std::to_string(header.obj_size));
     return;
   }
+  // The destination id comes from the peer for a Put, and from the local
+  // caller for a Get. Either way it must resolve into base_dir_. From here on
+  // the resolved path is the one used, so what was checked is what is
+  // written. The sender is already streaming obj_size bytes, so drain them,
+  // as the other early failures do, to keep the connection usable.
+  absl::StatusOr<std::string> dest_path =
+      ResolvePathWithinBaseDir(header.dest_obj_id, base_dir_);
+  if (!dest_path.ok()) {
+    RecvAndDiscard(client_fd, header.obj_size).IgnoreError();
+    fail(std::string(dest_path.status().message()));
+    return;
+  }
   if (is_respond_get_task) {
     UpdateTaskMetrics(
         header.task_id, [obj_size = header.obj_size](TaskMetricContainer& ts) {
@@ -810,8 +858,7 @@ void TransferService::HandleDataReceive(int client_fd,
           }
         });
   }
-  std::string tmp_obj_id =
-      std::string(header.dest_obj_id) + std::string(kTempFileSuffix);
+  std::string tmp_obj_id = *dest_path + std::string(kTempFileSuffix);
   // BufferObject's constructor throws when it cannot create the file.
   std::optional<BufferObject> buffer_obj;
   // Every failure from here on removes the temporary file. obj_size is the
@@ -849,11 +896,11 @@ void TransferService::HandleDataReceive(int client_fd,
   buffer_obj->close();
 
   // Rename the temporary file to the final destination.
-  if (rename(tmp_obj_id.c_str(), header.dest_obj_id) != 0) {
+  if (rename(tmp_obj_id.c_str(), dest_path->c_str()) != 0) {
     const std::string rename_error = std::strerror(errno);
     remove_tmp();
-    fail("Failed to rename temporary file " + tmp_obj_id + " to " +
-         header.dest_obj_id + ": " + rename_error);
+    fail("Failed to rename temporary file " + tmp_obj_id + " to " + *dest_path +
+         ": " + rename_error);
     return;
   }
   if (is_respond_get_task) {
@@ -993,8 +1040,9 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) {
   }
   metric_container.connection_acquired_time = absl::Now();
 
-  // Open the file as a buffer object. The constructor throws if the object
-  // cannot be opened (unreadable, a directory, or an empty file). The
+  // Open the file as a buffer object. The path was resolved and checked
+  // against base_dir_ by HandleGetObjRequest. The constructor throws if the
+  // object cannot be opened (unreadable, a directory, or an empty file). The
   // requester is blocked waiting for our response on this socket, so we must
   // tell it about the failure.
   std::optional<BufferObject> buffer_obj;
@@ -1134,11 +1182,23 @@ void TransferService::HandleGetObjRequest(int client_fd,
   LOG(INFO) << "Handling get object request for requested_obj_id: "
             << header.source_obj_id;
 
+  // The id comes from the peer. Refuse anything that resolves outside
+  // base_dir_, and hand the resolved path to the task so that what was
+  // checked is what gets opened.
+  absl::StatusOr<std::string> source_path =
+      ResolvePathWithinBaseDir(header.source_obj_id, base_dir_);
+  if (!source_path.ok()) {
+    LOG(ERROR) << "Rejecting get request for '" << header.source_obj_id
+               << "': " << source_path.status().message();
+    SendErrorResponse(client_fd, header.task_id, header.source_obj_id);
+    return;
+  }
+
   // Use the non-throwing overload. A stat error other than "not found" (for
   // example, a name that is too long) must also send kError to the requester
   // instead of throwing and dropping the socket.
   std::error_code ec;
-  if (!std::filesystem::exists(header.source_obj_id, ec)) {
+  if (!std::filesystem::exists(*source_path, ec)) {
     LOG(ERROR) << "Object not found: " << header.source_obj_id
                << (ec ? ": " + ec.message() : "");
     SendErrorResponse(client_fd, header.task_id, header.source_obj_id);
@@ -1149,16 +1209,15 @@ void TransferService::HandleGetObjRequest(int client_fd,
   metric_container->task_id = header.task_id;
   metric_container->task_type = TaskMetricContainer::TaskType::kRespondToGet;
   try {
-    metric_container->data_size =
-        std::filesystem::file_size(header.source_obj_id);
+    metric_container->data_size = std::filesystem::file_size(*source_path);
   } catch (const std::filesystem::filesystem_error& e) {
     LOG(WARNING) << "Failed to get file size for metrics: " << e.what();
     metric_container->data_size = 0;
   }
   metric_container->submit_time = absl::Now();
 
-  RespondToGetTask task(header.task_id, header.source_obj_id,
-                        header.dest_obj_id, client_fd, metric_container);
+  RespondToGetTask task(header.task_id, *source_path, header.dest_obj_id,
+                        client_fd, metric_container);
   // Not registered in pending_tasks_: the task_id is the peer's and could
   // name a local task. ExecuteRespondToGetTask logs its own outcome.
   task.Execute(this);