From d19b575fd1ca3986c4abf669fca6aefa56f461a5 Mon Sep 17 00:00:00 2001
From: Anton Romanov <romanton@google.com>
Date: Tue, 31 May 2022 20:47:30 +0000
Subject: [PATCH] crosvm: argh-ify vfio command

Usage: crosvm vfio <command> [<args>]

add/remove host vfio pci device into guest

Options:
  --help            display usage information

Commands: add               ADD
  remove            REMOVE

Change-Id: I79d0b5205b1ae8d880b26cf12d61a2f70bb73756
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/3680295
Reviewed-by: Dennis Kempin <denniskempin@google.com>
Commit-Queue: Anton Romanov <romanton@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>
---
 src/crosvm.rs | 38 ++++++++++++++++++++++++++++++++--
 src/main.rs   | 57 +++++++++++++++++++--------------------------------
 2 files changed, 57 insertions(+), 38 deletions(-)

diff --git a/src/crosvm.rs b/src/crosvm.rs
index 879e746c88..79f8e07d12 100644
--- a/src/crosvm.rs
+++ b/src/crosvm.rs
@@ -277,10 +277,44 @@ pub struct UsbCommand {
 /// Show package version.
 pub struct VersionCommand {}
 
-#[generate_catchall_args]
+#[derive(FromArgs)]
+#[argh(subcommand, name = "add")]
+/// ADD
+pub struct VfioAddSubCommand {
+    #[argh(positional)]
+    /// path to host's vfio sysfs
+    pub vfio_path: PathBuf,
+    #[argh(positional, arg_name = "VM_SOCKET")]
+    /// VM Socket path
+    pub socket_path: String,
+}
+
+#[derive(FromArgs)]
+#[argh(subcommand, name = "remove")]
+/// REMOVE
+pub struct VfioRemoveSubCommand {
+    #[argh(positional)]
+    /// path to host's vfio sysfs
+    pub vfio_path: PathBuf,
+    #[argh(positional, arg_name = "VM_SOCKET")]
+    /// VM Socket path
+    pub socket_path: String,
+}
+
+#[derive(FromArgs)]
+#[argh(subcommand)]
+pub enum VfioSubCommand {
+    Add(VfioAddSubCommand),
+    Remove(VfioRemoveSubCommand),
+}
+
+#[derive(FromArgs)]
 #[argh(subcommand, name = "vfio")]
 /// add/remove host vfio pci device into guest
-pub struct VfioCrosvmCommand {}
+pub struct VfioCrosvmCommand {
+    #[argh(subcommand)]
+    pub command: VfioSubCommand,
+}
 
 #[derive(FromArgs)]
 #[argh(subcommand, name = "device")]
diff --git a/src/main.rs b/src/main.rs
index 19d2f02299..6811a9e6da 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2779,44 +2779,29 @@ fn modify_battery(cmd: crosvm::BatteryCommand) -> std::result::Result<(), ()> {
 }
 
 fn modify_vfio(cmd: crosvm::VfioCrosvmCommand) -> std::result::Result<(), ()> {
-    if cmd.args.len() < 3 {
-        print_help(
-            "crosvm vfio",
-            "[add | remove host_vfio_sysfs] VM_SOCKET...",
-            &[],
-        );
-        return Err(());
-    }
-
-    let mut args = cmd.args.into_iter();
-    // This unwrap will not panic because of the above length check.
-    let command = args.next().unwrap();
-    let path_str = args.next().unwrap();
-    let vfio_path = PathBuf::from(&path_str);
-    if !vfio_path.exists() || !vfio_path.is_dir() {
-        error!("Invalid host sysfs path: {}", path_str);
-        return Err(());
-    }
-
-    let socket_path = args.next().unwrap();
-    let socket_path = Path::new(&socket_path);
-
-    let add = match command.as_ref() {
-        "add" => true,
-        "remove" => false,
-        other => {
-            error!("Invalid vfio command {}", other);
-            return Err(());
+    let (request, socket_path, vfio_path) = match cmd.command {
+        crosvm::VfioSubCommand::Add(c) => {
+            let request = VmRequest::VfioCommand {
+                vfio_path: c.vfio_path.clone(),
+                add: true,
+                hp_interrupt: true,
+            };
+            (request, c.socket_path, c.vfio_path)
+        }
+        crosvm::VfioSubCommand::Remove(c) => {
+            let request = VmRequest::VfioCommand {
+                vfio_path: c.vfio_path.clone(),
+                add: true,
+                hp_interrupt: true,
+            };
+            (request, c.socket_path, c.vfio_path)
         }
     };
-
-    let hp_interrupt = true;
-    let request = VmRequest::VfioCommand {
-        vfio_path,
-        add,
-        hp_interrupt,
-    };
-    handle_request(&request, socket_path)?;
+    if !vfio_path.exists() || !vfio_path.is_dir() {
+        error!("Invalid host sysfs path: {:?}", vfio_path);
+        return Err(());
+    }
+    handle_request(&request, Path::new(&socket_path))?;
     Ok(())
 }