From daa000e659ba9df6da32bb054950424b8e7b7e6c Mon Sep 17 00:00:00 2001
From: Elie Kheirallah <khei@google.com>
Date: Wed, 28 Sep 2022 19:01:02 +0000
Subject: [PATCH] crosvm: cmdline: Add Snapshot command to crosvm

Bug=b:232437513, b:253937826
Test=cargo build

Change-Id: I3e90e5dc86e53a7b2fa0eae6be2f63fcb91a345e
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/3924744
Commit-Queue: Elie Kheirallah <khei@google.com>
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>
Auto-Submit: Elie Kheirallah <khei@google.com>
---
 arch/src/lib.rs       |  3 ++-
 src/crosvm/cmdline.rs | 25 +++++++++++++++++++++++++
 src/main.rs           | 16 ++++++++++++++++
 vm_control/src/lib.rs | 21 +++++++++++++++++++++
 4 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/arch/src/lib.rs b/arch/src/lib.rs
index 5394b01b2a..0fb2d6f9cf 100644
--- a/arch/src/lib.rs
+++ b/arch/src/lib.rs
@@ -562,7 +562,7 @@ pub fn configure_pci_device<V: VmArch, Vcpu: VcpuArch>(
     Ok(pci_address)
 }
 
-/// Creates a Virtio MMIO devices for use by this Vm.
+/// Creates Virtio MMIO devices for use by this Vm.
 pub fn generate_virtio_mmio_bus(
     devices: Vec<(VirtioMmioDevice, Option<Minijail>)>,
     irq_chip: &mut dyn IrqChip,
@@ -885,6 +885,7 @@ pub fn generate_pci_root(
             .partition(|(_, (_, jail))| jail.is_some());
         sandboxed.into_iter().chain(non_sandboxed.into_iter())
     };
+
     let mut amls = BTreeMap::new();
     for (dev_idx, dev_value) in devices {
         #[cfg(unix)]
diff --git a/src/crosvm/cmdline.rs b/src/crosvm/cmdline.rs
index d7800b78e6..8bef99b2a3 100644
--- a/src/crosvm/cmdline.rs
+++ b/src/crosvm/cmdline.rs
@@ -152,6 +152,7 @@ pub enum CrossPlatformCommands {
     Usb(UsbCommand),
     Version(VersionCommand),
     Vfio(VfioCrosvmCommand),
+    Snapshot(SnapshotCommand),
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -597,6 +598,30 @@ impl From<DiskOption> for DiskOptionWithId {
     }
 }
 
+#[derive(FromArgs)]
+#[argh(subcommand, name = "snapshot", description = "Snapshot commands")]
+/// Snapshot commands
+pub struct SnapshotCommand {
+    #[argh(subcommand)]
+    pub snapshot_command: SnapshotSubCommands,
+}
+
+#[derive(FromArgs)]
+#[argh(subcommand, name = "take")]
+/// Take a snapshot of the VM
+pub struct SnapshotTakeCommand {
+    #[argh(positional, arg_name = "VM_SOCKET")]
+    /// VM Socket path
+    pub socket_path: String,
+}
+
+#[derive(FromArgs)]
+#[argh(subcommand)]
+/// Snapshot commands
+pub enum SnapshotSubCommands {
+    Take(SnapshotTakeCommand),
+}
+
 /// Container for GpuParameters that have been fixed after parsing using serde.
 ///
 /// This deserializes as a regular `GpuParameters` and applies validation.
diff --git a/src/main.rs b/src/main.rs
index a9aa195c77..73fa4c371d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -70,6 +70,7 @@ use vm_control::BalloonControlCommand;
 use vm_control::DiskControlCommand;
 use vm_control::HotPlugDeviceInfo;
 use vm_control::HotPlugDeviceType;
+use vm_control::SnapshotCommand;
 use vm_control::SwapCommand;
 use vm_control::UsbControlResult;
 use vm_control::VmRequest;
@@ -525,6 +526,18 @@ fn modify_usb(cmd: cmdline::UsbCommand) -> std::result::Result<(), ()> {
     }
 }
 
+fn snapshot_vm(cmd: cmdline::SnapshotCommand) -> std::result::Result<(), ()> {
+    use cmdline::SnapshotSubCommands::*;
+    let (socket_path, request) = match cmd.snapshot_command {
+        Take(path) => {
+            let req = VmRequest::Snapshot(SnapshotCommand::Take);
+            (path.socket_path, req)
+        }
+    };
+    let socket_path = Path::new(&socket_path);
+    vms_request(&request, socket_path)
+}
+
 #[allow(clippy::unnecessary_wraps)]
 fn pkg_version() -> std::result::Result<(), ()> {
     const VERSION: Option<&'static str> = option_env!("CARGO_PKG_VERSION");
@@ -700,6 +713,9 @@ fn crosvm_main<I: IntoIterator<Item = String>>(args: I) -> Result<CommandStatus>
                     CrossPlatformCommands::Vfio(cmd) => {
                         modify_vfio(cmd).map_err(|_| anyhow!("vfio subcommand failed"))
                     }
+                    CrossPlatformCommands::Snapshot(cmd) => {
+                        snapshot_vm(cmd).map_err(|_| anyhow!("snapshot subcommand failed"))
+                    }
                 }
                 .map(|_| CommandStatus::SuccessOrVmStop)
             }
diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index fd6fbc763a..922a5d47d1 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -265,6 +265,21 @@ impl Display for UsbControlResult {
     }
 }
 
+/// Commands for snapshot feature
+#[derive(Serialize, Deserialize, Debug)]
+pub enum SnapshotCommand {
+    Take,
+}
+
+/// Response for [SnapshotCommand]
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum SnapshotControlResult {
+    /// The request is accepted successfully.
+    Ok,
+    /// The command fails.
+    Failed,
+}
+
 /// Source of a `VmMemoryRequest::RegisterMemory` mapping.
 #[derive(Serialize, Deserialize)]
 pub enum VmMemorySource {
@@ -966,6 +981,8 @@ pub enum VmRequest {
         device: HotPlugDeviceInfo,
         add: bool,
     },
+    /// Command to Snapshot devices
+    Snapshot(SnapshotCommand),
 }
 
 pub fn handle_disk_command(command: &DiskControlCommand, disk_host_tube: &Tube) -> VmResponse {
@@ -1275,6 +1292,7 @@ impl VmRequest {
                 }
             }
             VmRequest::HotPlugCommand { device: _, add: _ } => VmResponse::Ok,
+            VmRequest::Snapshot(SnapshotCommand::Take) => VmResponse::Ok,
         }
     }
 }
@@ -1305,6 +1323,8 @@ pub enum VmResponse {
     BatResponse(BatControlResult),
     /// Results of swap status command.
     SwapStatus(SwapStatus),
+    /// Results of snapshot commands.
+    SnapshotResponse(SnapshotControlResult),
 }
 
 impl Display for VmResponse {
@@ -1343,6 +1363,7 @@ impl Display for VmResponse {
                         .unwrap_or_else(|_| "invalid_response".to_string()),
                 )
             }
+            SnapshotResponse(result) => write!(f, "snapshot control request result {:?}", result),
         }
     }
 }