feat(snapshot): allow block device path override on restore

Allow overriding block device host paths when restoring from a snapshot.
This is useful when disk paths are non-deterministic (e.g. container
runtimes using devmapper) or when restoring on a different host where
the backing file is at a different location.

Adds a drive_overrides parameter to the snapshot load API, following the
same pattern as the existing network_overrides and vsock_override.
Includes Python integration tests that verify the override works with a
different file and that an unknown drive_id is rejected.

Closes #4014

Signed-off-by: Amit Patil <iamitpatil2001@gmail.com>

Author: meAmitPatil

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to
 
 ### Added
 
+- [#5774](https://github.com/firecracker-microvm/firecracker/pull/5774): Add
+  support for block device path overriding on snapshot restore.
 - [#5323](https://github.com/firecracker-microvm/firecracker/pull/5323): Add
   support for Vsock Unix domain socket path overriding on snapshot restore. More
   information can be found in the

docs/snapshotting/network-for-clones.md

@@ -200,6 +200,37 @@ ip addr add 172.16.3.2/30 dev eth0
 ip route add default via 172.16.3.1/30 dev eth0
 ```
 
+### Overriding block device paths
+
+When restoring a VM from a snapshot on a different host or in an environment
+where disk paths are non-deterministic (e.g. container runtimes using
+devmapper), the block device paths baked into the snapshot state may no longer
+be valid. In this case you can use the `drive_overrides` parameter of the
+snapshot restore API to specify the new host path for each block device.
+
+For example, if we have a block device with drive ID `rootfs` in the snapshotted
+microVM, we can override its path during snapshot resume:
+
+```bash
+curl --unix-socket /tmp/firecracker.socket -i \
+    -X PUT 'http://localhost/snapshot/load' \
+    -H  'Accept: application/json' \
+    -H  'Content-Type: application/json' \
+    -d '{
+            "snapshot_path": "./snapshot_file",
+            "mem_backend": {
+                "backend_path": "./mem_file",
+                "backend_type": "File"
+            },
+            "drive_overrides": [
+                 {
+                     "drive_id": "rootfs",
+                     "path_on_host": "/new/path/to/rootfs.ext4"
+                 }
+            ]
+    }'
+```
+
 # Ingress connectivity
 
 The above setup only provides egress connectivity. If in addition we also want

src/firecracker/src/api_server/request/snapshot.rs

@@ -111,6 +111,7 @@ fn parse_put_snapshot_load(body: &Body) -> Result<ParsedRequest, RequestError> {
         resume_vm: snapshot_config.resume_vm,
         network_overrides: snapshot_config.network_overrides,
         vsock_override: snapshot_config.vsock_override,
+        drive_overrides: snapshot_config.drive_overrides,
     };
 
     // Construct the `ParsedRequest` object.
@@ -126,7 +127,9 @@ fn parse_put_snapshot_load(body: &Body) -> Result<ParsedRequest, RequestError> {
 
 #[cfg(test)]
 mod tests {
-    use vmm::vmm_config::snapshot::{MemBackendConfig, MemBackendType, NetworkOverride};
+    use vmm::vmm_config::snapshot::{
+        DriveOverride, MemBackendConfig, MemBackendType, NetworkOverride,
+    };
 
     use super::*;
     use crate::api_server::parsed_request::tests::{depr_action_from_req, vmm_action_from_request};
@@ -189,6 +192,7 @@ mod tests {
             resume_vm: false,
             network_overrides: vec![],
             vsock_override: None,
+            drive_overrides: vec![],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -220,6 +224,7 @@ mod tests {
             resume_vm: false,
             network_overrides: vec![],
             vsock_override: None,
+            drive_overrides: vec![],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -251,6 +256,7 @@ mod tests {
             resume_vm: true,
             network_overrides: vec![],
             vsock_override: None,
+            drive_overrides: vec![],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -291,6 +297,48 @@ mod tests {
                 host_dev_name: String::from("vmtap2"),
             }],
             vsock_override: None,
+            drive_overrides: vec![],
+        };
+        let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
+        assert!(
+            parsed_request
+                .parsing_info()
+                .take_deprecation_message()
+                .is_none()
+        );
+        assert_eq!(
+            vmm_action_from_request(parsed_request),
+            VmmAction::LoadSnapshot(expected_config)
+        );
+
+        let body = r#"{
+            "snapshot_path": "foo",
+            "mem_backend": {
+                "backend_path": "bar",
+                "backend_type": "File"
+            },
+            "resume_vm": true,
+            "drive_overrides": [
+                {
+                    "drive_id": "rootfs",
+                    "path_on_host": "/new/path/rootfs.ext4"
+                }
+            ]
+        }"#;
+        let expected_config = LoadSnapshotParams {
+            snapshot_path: PathBuf::from("foo"),
+            mem_backend: MemBackendConfig {
+                backend_path: PathBuf::from("bar"),
+                backend_type: MemBackendType::File,
+            },
+            track_dirty_pages: false,
+            resume_vm: true,
+            network_overrides: vec![],
+            vsock_override: None,
+            drive_overrides: vec![DriveOverride {
+                drive_id: String::from("rootfs"),
+                path_on_host: String::from("/new/path/rootfs.ext4"),
+            }],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -319,6 +367,7 @@ mod tests {
             resume_vm: true,
             network_overrides: vec![],
             vsock_override: None,
+            drive_overrides: vec![],
         };
         let parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert_eq!(

src/firecracker/swagger/firecracker.yaml

@@ -1603,6 +1603,24 @@ definitions:
         description:
           The new path for the backing Unix Domain Socket.
 
+  DriveOverride:
+    type: object
+    description:
+      Allows for changing the backing host path of a block device
+      during snapshot restore.
+    required:
+      - drive_id
+      - path_on_host
+    properties:
+      drive_id:
+        type: string
+        description:
+          The ID of the drive to modify
+      path_on_host:
+        type: string
+        description:
+          The new path on the host for the block device's backing file
+
   SnapshotLoadParams:
     type: object
     description:
@@ -1650,6 +1668,11 @@ definitions:
           for restoring a snapshot with a different socket path than the one used
           when the snapshot was created. For example, when the original socket path
           is no longer available or when deploying to a different environment.
+      drive_overrides:
+        type: array
+        description: Block device host paths to override
+        items:
+          $ref: "#/definitions/DriveOverride"
 
 
   TokenBucket:

src/vmm/src/devices/virtio/block/persist.rs

@@ -20,8 +20,24 @@ pub enum BlockState {
 impl BlockState {
     pub fn is_activated(&self) -> bool {
         match self {
-            BlockState::Virtio(virtio_block_state) => virtio_block_state.virtio_state.activated,
-            BlockState::VhostUser(vhost_user_block_state) => false,
+            BlockState::Virtio(state) => state.virtio_state.activated,
+            BlockState::VhostUser(_) => false,
+        }
+    }
+
+    /// Returns the drive ID.
+    pub fn id(&self) -> &str {
+        match self {
+            BlockState::Virtio(state) => &state.id,
+            BlockState::VhostUser(state) => &state.id,
+        }
+    }
+
+    /// Overrides the host path (disk path or socket path) of the block device.
+    pub fn set_host_path(&mut self, path: &str) {
+        match self {
+            BlockState::Virtio(state) => state.disk_path = path.to_string(),
+            BlockState::VhostUser(state) => state.socket_path = path.to_string(),
         }
     }
 }
@@ -31,3 +47,90 @@ impl BlockState {
 pub struct BlockConstructorArgs {
     pub mem: GuestMemoryMmap,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::devices::virtio::block::CacheType;
+    use crate::devices::virtio::block::virtio::persist::FileEngineTypeState;
+    use crate::devices::virtio::device::VirtioDeviceType;
+    use crate::devices::virtio::persist::VirtioDeviceState;
+    use crate::rate_limiter::persist::RateLimiterState;
+
+    fn virtio_block_state(id: &str, disk_path: &str, activated: bool) -> VirtioBlockState {
+        VirtioBlockState {
+            id: id.to_string(),
+            partuuid: None,
+            cache_type: CacheType::Unsafe,
+            root_device: false,
+            disk_path: disk_path.to_string(),
+            virtio_state: VirtioDeviceState {
+                device_type: VirtioDeviceType::Block,
+                avail_features: 0,
+                acked_features: 0,
+                queues: vec![],
+                activated,
+            },
+            rate_limiter_state: Default::default(),
+            file_engine_type: FileEngineTypeState::Sync,
+        }
+    }
+
+    fn vhost_user_block_state(id: &str, socket_path: &str) -> VhostUserBlockState {
+        VhostUserBlockState {
+            id: id.to_string(),
+            partuuid: None,
+            cache_type: CacheType::Unsafe,
+            root_device: false,
+            socket_path: socket_path.to_string(),
+            vu_acked_protocol_features: 0,
+            config_space: vec![],
+            virtio_state: VirtioDeviceState {
+                device_type: VirtioDeviceType::Block,
+                avail_features: 0,
+                acked_features: 0,
+                queues: vec![],
+                activated: false,
+            },
+        }
+    }
+
+    #[test]
+    fn test_block_state_id() {
+        let virtio = BlockState::Virtio(virtio_block_state("rootfs", "/path", false));
+        assert_eq!(virtio.id(), "rootfs");
+
+        let vhost = BlockState::VhostUser(vhost_user_block_state("scratch", "/sock"));
+        assert_eq!(vhost.id(), "scratch");
+    }
+
+    #[test]
+    fn test_block_state_is_activated() {
+        let active = BlockState::Virtio(virtio_block_state("rootfs", "/path", true));
+        assert!(active.is_activated());
+
+        let inactive = BlockState::Virtio(virtio_block_state("rootfs", "/path", false));
+        assert!(!inactive.is_activated());
+
+        // vhost-user always returns false
+        let vhost = BlockState::VhostUser(vhost_user_block_state("rootfs", "/sock"));
+        assert!(!vhost.is_activated());
+    }
+
+    #[test]
+    fn test_block_state_set_host_path() {
+        let mut virtio = BlockState::Virtio(virtio_block_state("rootfs", "/old/path", false));
+        virtio.set_host_path("/new/path");
+        match &virtio {
+            BlockState::Virtio(state) => assert_eq!(state.disk_path, "/new/path"),
+            _ => panic!("expected Virtio variant"),
+        }
+
+        let mut vhost = BlockState::VhostUser(vhost_user_block_state("rootfs", "/old/sock"));
+        vhost.set_host_path("/new/sock");
+        match &vhost {
+            BlockState::VhostUser(state) => assert_eq!(state.socket_path, "/new/sock"),
+            _ => panic!("expected VhostUser variant"),
+        }
+    }
+}

src/vmm/src/devices/virtio/block/vhost_user/persist.rs

@@ -15,11 +15,11 @@ use crate::snapshot::Persist;
 /// vhost-user block device state.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct VhostUserBlockState {
-    id: String,
+    pub id: String,
     partuuid: Option<String>,
     cache_type: CacheType,
     root_device: bool,
-    socket_path: String,
+    pub socket_path: String,
     vu_acked_protocol_features: u64,
     config_space: Vec<u8>,
     virtio_state: VirtioDeviceState,

src/vmm/src/devices/virtio/block/virtio/persist.rs

@@ -52,11 +52,11 @@ impl From<FileEngineTypeState> for FileEngineType {
 /// Holds info about the block device. Gets saved in snapshot.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct VirtioBlockState {
-    id: String,
+    pub id: String,
     partuuid: Option<String>,
     cache_type: CacheType,
     root_device: bool,
-    disk_path: String,
+    pub disk_path: String,
     pub virtio_state: VirtioDeviceState,
     rate_limiter_state: RateLimiterState,
     file_engine_type: FileEngineTypeState,

src/vmm/src/persist.rs

@@ -408,6 +408,26 @@ pub fn restore_from_snapshot(
             .clone_from(&vsock_override.uds_path);
     }
 
+    for entry in &params.drive_overrides {
+        microvm_state
+            .device_states
+            .mmio_state
+            .block_devices
+            .iter_mut()
+            .map(|device| &mut device.device_state)
+            .chain(
+                microvm_state
+                    .device_states
+                    .pci_state
+                    .block_devices
+                    .iter_mut()
+                    .map(|device| &mut device.device_state),
+            )
+            .find(|x| x.id() == entry.drive_id)
+            .map(|device_state| device_state.set_host_path(&entry.path_on_host))
+            .ok_or(SnapshotStateFromFileError::UnknownBlockDevice)?;
+    }
+
     let track_dirty_pages = params.track_dirty_pages;
 
     let vcpu_count = microvm_state
@@ -481,6 +501,8 @@ pub enum SnapshotStateFromFileError {
     UnknownNetworkDevice,
     /// Unknown Vsock Device.
     UnknownVsockDevice,
+    /// Unknown Block Device.
+    UnknownBlockDevice,
 }
 
 fn snapshot_state_from_file(

src/vmm/src/rpc_interface.rs

@@ -1286,6 +1286,7 @@ mod tests {
                 resume_vm: false,
                 network_overrides: vec![],
                 vsock_override: None,
+                drive_overrides: vec![],
             },
         )));
         check_unsupported(runtime_request(VmmAction::SetEntropyDevice(