[RFC] feat(snapshot): allow vsock uds path override on restore

docs/snapshotting/snapshot-support.md

@@ -24,7 +24,7 @@
 - [Snapshot security and uniqueness](#snapshot-security-and-uniqueness)
   - [Secure and insecure usage examples](#usage-examples)
   - [Reusing snapshotted states securely](#reusing-snapshotted-states-securely)
-- [Vsock device limitation](#vsock-device-limitation)
+- [Vsock device reset](#vsock-device-reset)
 - [VMGenID device limitation](#vmgenid-device-limitation)
 - [Where can I resume my snapshots?](#where-can-i-resume-my-snapshots)
 

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

@@ -110,6 +110,7 @@ fn parse_put_snapshot_load(body: &Body) -> Result<ParsedRequest, RequestError> {
             || snapshot_config.track_dirty_pages,
         resume_vm: snapshot_config.resume_vm,
         network_overrides: snapshot_config.network_overrides,
+        vsock_override: snapshot_config.vsock_override,
     };
 
     // Construct the `ParsedRequest` object.
@@ -187,6 +188,7 @@ mod tests {
             track_dirty_pages: false,
             resume_vm: false,
             network_overrides: vec![],
+            vsock_override: None,
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -217,6 +219,7 @@ mod tests {
             track_dirty_pages: true,
             resume_vm: false,
             network_overrides: vec![],
+            vsock_override: None,
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -247,6 +250,7 @@ mod tests {
             track_dirty_pages: false,
             resume_vm: true,
             network_overrides: vec![],
+            vsock_override: None,
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -286,6 +290,7 @@ mod tests {
                 iface_id: String::from("eth0"),
                 host_dev_name: String::from("vmtap2"),
             }],
+            vsock_override: None,
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -313,6 +318,7 @@ mod tests {
             track_dirty_pages: false,
             resume_vm: true,
             network_overrides: vec![],
+            vsock_override: None,
         };
         let parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert_eq!(

src/firecracker/swagger/firecracker.yaml

@@ -1239,7 +1239,19 @@ definitions:
         type: string
         description:
           The new host device of the interface
-
+  VsockOverride:
+    type: object
+    description:
+      Allows for changing the backing Unix Domain Socket of a vsock device
+      during snapshot restore.
+    required:
+      - uds_path
+    properties:
+      uds_path:
+        type: string
+        description:
+          The new path for the backing Unix Domain Socket.
+
   SnapshotLoadParams:
     type: object
     description:
@@ -1280,6 +1292,13 @@ definitions:
         description: Network host device names to override
         items:
           $ref: "#/definitions/NetworkOverride"
+      vsock_override:
+        $ref: "#/definitions/VsockOverride"
+        description:
+          Overrides the vsock device's UDS path on snapshot restore. This is useful
+          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.
 
 
   TokenBucket:

src/vmm/src/persist.rs

@@ -26,6 +26,7 @@ use crate::cpu_config::x86_64::cpuid::CpuidTrait;
 #[cfg(target_arch = "x86_64")]
 use crate::cpu_config::x86_64::cpuid::common::get_vendor_id_from_host;
 use crate::device_manager::persist::{ACPIDeviceManagerState, DevicePersistError, DeviceStates};
+use crate::devices::virtio::vsock::persist::{VsockBackendState, VsockUdsState};
 use crate::logger::{info, warn};
 use crate::resources::VmResources;
 use crate::seccomp::BpfThreadMap;
@@ -347,6 +348,17 @@ pub fn restore_from_snapshot(
             return Err(SnapshotStateFromFileError::UnknownNetworkDevice.into());
         }
     }
+
+    if let Some(vsock_override) = &params.vsock_override {
+        if let Some(vsock_device) = microvm_state.device_states.vsock_device.as_mut() {
+            vsock_device.device_state.backend = VsockBackendState::Uds(VsockUdsState {
+                path: vsock_override.uds_path.clone(),
+            });
+        } else {
+            return Err(SnapshotStateFromFileError::UnknownVsockDevice.into());
+        }
+    }
+
     let track_dirty_pages = params.track_dirty_pages;
 
     let vcpu_count = microvm_state
@@ -420,6 +432,8 @@ pub enum SnapshotStateFromFileError {
     Load(#[from] crate::snapshot::SnapshotError),
     /// Unknown Network Device.
     UnknownNetworkDevice,
+    /// Unknown Vsock Device.
+    UnknownVsockDevice,
 }
 
 fn snapshot_state_from_file(

src/vmm/src/rpc_interface.rs

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

src/vmm/src/vmm_config/snapshot.rs

@@ -57,6 +57,13 @@ pub struct NetworkOverride {
     pub host_dev_name: String,
 }
 
+/// Allows for changing the host UDS of the vsock backend during snapshot restore
+#[derive(Debug, PartialEq, Eq, Deserialize)]
+pub struct VsockOverride {
+    /// The path to the UDS that will be used for the vsock interface
+    pub uds_path: String,
+}
+
 /// Stores the configuration that will be used for loading a snapshot.
 #[derive(Debug, PartialEq, Eq)]
 pub struct LoadSnapshotParams {
@@ -72,6 +79,8 @@ pub struct LoadSnapshotParams {
     pub resume_vm: bool,
     /// The network devices to override on load.
     pub network_overrides: Vec<NetworkOverride>,
+    /// When set, the vsock backend UDS path will be overridden
+    pub vsock_override: Option<VsockOverride>,
 }
 
 /// Stores the configuration for loading a snapshot that is provided by the user.
@@ -101,6 +110,9 @@ pub struct LoadSnapshotConfig {
     /// The network devices to override on load.
     #[serde(default)]
     pub network_overrides: Vec<NetworkOverride>,
+    /// Whether or not to override the vsock backend UDS path.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub vsock_override: Option<VsockOverride>,
 }
 
 /// Stores the configuration used for managing snapshot memory.