node-mgr: validate LoadSnapshot requests before staging host state

crates/libs/node/mgr/server/src/service/snapshots.rs

@@ -572,8 +572,15 @@ impl NodeManagerService {
     ) -> Result<Response<LoadSnapshotResponse>, Status> {
         let req = request.get_ref();
         validate_sandbox_id(&req.sandbox_id)?;
+        // Validate before either path stages the jail or re-owns the snapshot
+        // files, and so before the load, which cannot be undone short of
+        // killing the VM.
+        let rootfs_limiter = req
+            .rootfs_rate_limit
+            .as_ref()
+            .map(drive_rate_limiter)
+            .transpose()?;
         let jail_root = format!("{JAIL_DIR}/firecracker/{}/root", req.sandbox_id);
-        let api_socket = format!("{jail_root}/firecracker.sock");
         let _lifecycle_permit = self
             .lifecycle_limiter
             .acquire(crate::lifecycle_limiter::LifecycleMethod::LoadSnapshot)
@@ -614,9 +621,9 @@ impl NodeManagerService {
                 format!("/{ENC_SNAPSHOT_JAIL_SUBDIR}/{ENC_SNAPSHOT_MEMORY_NAME}");
             self.finish_load_snapshot(
                 &lease,
-                req,
+                &req.sandbox_id,
+                rootfs_limiter.as_ref(),
                 &jail_root,
-                &api_socket,
                 &fc_snapshot_path,
                 &fc_mem_file_path,
             )
@@ -663,9 +670,9 @@ impl NodeManagerService {
 
         self.finish_load_snapshot(
             &lease,
-            req,
+            &req.sandbox_id,
+            rootfs_limiter.as_ref(),
             &jail_root,
-            &api_socket,
             &fc_snapshot_path,
             &fc_mem_file_path,
         )
@@ -675,29 +682,22 @@ impl NodeManagerService {
 
     /// Shared tail of the load-snapshot paths (plaintext source_dir and BYOK
     /// encrypted volume): wait for the FC API socket, load the snapshot with the
-    /// given memory file, apply the sandbox's rootfs rate limit, then fix
-    /// v.sock permissions.
+    /// given memory file, apply the sandbox's validated rootfs rate limit, then
+    /// fix v.sock permissions.
     async fn finish_load_snapshot(
         &self,
         lease: &SlotInfo,
-        request: &LoadSnapshotRequest,
+        sandbox_id: &str,
+        rootfs_limiter: Option<&pplx_space_common_core::types::RateLimiter>,
         jail_root: &str,
-        api_socket: &str,
         fc_snapshot_path: &str,
         fc_mem_file_path: &str,
     ) -> Result<(), Status> {
-        let sandbox_id = request.sandbox_id.as_str();
-        // Validate before the load, which cannot be undone short of killing
-        // the VM.
-        let rootfs_limiter = request
-            .rootfs_rate_limit
-            .as_ref()
-            .map(drive_rate_limiter)
-            .transpose()?;
+        let api_socket = format!("{jail_root}/firecracker.sock");
         // Wait for the API socket to appear (FC needs time to create it after launch)
         let api_socket_wait_started = Instant::now();
         loop {
-            if std::path::Path::new(api_socket).exists() {
+            if std::path::Path::new(&api_socket).exists() {
                 break;
             }
             if api_socket_wait_started.elapsed() >= FIRECRACKER_API_SOCKET_TIMEOUT {
@@ -714,15 +714,15 @@ impl NodeManagerService {
             "firecracker api socket observed"
         );
 
-        let client = crate::fc_api::FcApiClient::new(api_socket.to_string());
+        let client = crate::fc_api::FcApiClient::new(api_socket);
         client
             .load_snapshot(fc_snapshot_path, fc_mem_file_path)
             .await
             .map_err(|e| Status::internal(format!("load snapshot failed: {e}")))?;
         // The loaded drive limiter is the one of the VM the snapshot was taken
         // from. Replace it with the sandbox's own while the VM is still
         // paused; a sandbox must not run under another VM's limits.
-        if let Some(limiter) = &rootfs_limiter
+        if let Some(limiter) = rootfs_limiter
             && let Err(error) = client.update_drive_rate_limiter("rootfs", limiter).await
         {
             let status = Status::internal(format!(

crates/libs/node/mgr/server/src/service/snapshots/tests.rs

@@ -547,8 +547,16 @@ async fn create_snapshot_refuses_invalid_requests_before_dialing_firecracker() {
 
 // No invalid case may pass validation: the next step stages files or binds the
 // volume into the node's fixed jail directory, or waits for the Firecracker
-// socket. The valid encrypted load shows where validation ends: its volume bind
-// is the first blocking host step, which the fixture refuses.
+// socket. The valid encrypted loads show where validation ends: their volume
+// bind is the first blocking host step, which the fixture refuses. Rows whose
+// guard precedes the lease lookup and that stage nothing use the unallocated
+// sandbox, so dropping the guard stops at the lease check instead of polling
+// the jail's Firecracker socket. The encrypted rate-limit row shows the limit
+// is refused before the refused bind. The checkpoint row only shows it is
+// refused before the path checks: the source check first reads the sandbox's
+// checkpoint files under the fixed cache root, which do not exist, so no
+// plaintext row can reach the staging step, and a valid checkpoint load has no
+// row because that read's outcome depends on the host.
 #[tokio::test]
 async fn load_snapshot_refuses_invalid_requests_before_staging_the_jail() {
     use crate::lifecycle_limiter::LifecycleMethod;
@@ -637,14 +645,40 @@ async fn load_snapshot_refuses_invalid_requests_before_staging_the_jail() {
                     bandwidth: None,
                     operations: bucket,
                 }),
-                ..valid.clone()
+                ..unallocated.clone()
             },
             Code::InvalidArgument,
             "rootfs_rate_limit.bandwidth is required",
         ),
         (
             "rate limit without operations",
             LoadSnapshotRequest {
+                rootfs_rate_limit: Some(DriveRateLimit {
+                    bandwidth: bucket,
+                    operations: None,
+                }),
+                ..unallocated.clone()
+            },
+            Code::InvalidArgument,
+            "rootfs_rate_limit.operations is required",
+        ),
+        (
+            "checkpoint load with a rate limit without bandwidth",
+            LoadSnapshotRequest {
+                source_dir: Some(format!("{}/{id}/snapshot", super::CACHE_DIR)),
+                rootfs_rate_limit: Some(DriveRateLimit {
+                    bandwidth: None,
+                    operations: bucket,
+                }),
+                ..valid.clone()
+            },
+            Code::InvalidArgument,
+            "rootfs_rate_limit.bandwidth is required",
+        ),
+        (
+            "encrypted load with a rate limit without operations",
+            LoadSnapshotRequest {
+                snapshot_volume_mount_dir: Some(snapshot_volume(&id)),
                 rootfs_rate_limit: Some(DriveRateLimit {
                     bandwidth: bucket,
                     operations: None,
@@ -663,6 +697,19 @@ async fn load_snapshot_refuses_invalid_requests_before_staging_the_jail() {
             Code::Cancelled,
             "artifact work cancelled",
         ),
+        (
+            "valid encrypted load with a rate limit",
+            LoadSnapshotRequest {
+                snapshot_volume_mount_dir: Some(snapshot_volume(&id)),
+                rootfs_rate_limit: Some(DriveRateLimit {
+                    bandwidth: bucket,
+                    operations: bucket,
+                }),
+                ..valid.clone()
+            },
+            Code::Cancelled,
+            "artifact work cancelled",
+        ),
     ];
     let mut mismatches = Vec::new();
     for (case, request, code, message) in cases {