diff --git a/crates/libs/node/mgr/server/src/service/snapshots.rs b/crates/libs/node/mgr/server/src/service/snapshots.rs index 55a28e5a56..18f8f0cd09 100644 --- a/crates/libs/node/mgr/server/src/service/snapshots.rs +++ b/crates/libs/node/mgr/server/src/service/snapshots.rs @@ -572,8 +572,15 @@ impl NodeManagerService { ) -> Result, 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,7 +714,7 @@ 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 @@ -722,7 +722,7 @@ impl NodeManagerService { // 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!( diff --git a/crates/libs/node/mgr/server/src/service/snapshots/tests.rs b/crates/libs/node/mgr/server/src/service/snapshots/tests.rs index 1c4a192e18..f59a063a04 100644 --- a/crates/libs/node/mgr/server/src/service/snapshots/tests.rs +++ b/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,7 +645,7 @@ 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", @@ -645,6 +653,32 @@ async fn load_snapshot_refuses_invalid_requests_before_staging_the_jail() { ( "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 {