add more documentation

src/cli.rs

@@ -1,7 +1,7 @@
 use clap::{Parser, Subcommand};
 use std::{path::PathBuf, str::FromStr};
 
-/// VM controller for AirgapOS
+/// VM controller for `AirgapOS`
 #[derive(Parser, Clone, Debug)]
 pub struct App {
     // global options go here
@@ -16,7 +16,10 @@ pub struct App {
 #[derive(Subcommand, Clone, Debug)]
 pub enum Commands {
     /// Start a headless VM in the background.
-    Start,
+    Start {
+        #[arg(long, default_value = "1G")]
+        memory: String,
+    },
 
     /// Stop a headless VM.
     Stop,
@@ -29,7 +32,7 @@ pub enum Commands {
 
     /// Attach a USB device to a running VM.
     Attach {
-        /// The device to attach.
+        /// The device to attach, in the format of `vendorid:deviceid`.
         device: DeviceIdentifier,
     },
 
@@ -59,19 +62,16 @@ pub enum Commands {
         /// Arguments to pass to the running command.
         args: Vec<String>,
     },
-
-    /// Test synchronization by repeatedly running commands.
-    Test {}
 }
 
 /// An attachable USB device identifier.
 #[derive(Clone, Debug)]
 pub struct DeviceIdentifier {
     /// The Vendor ID.
-    pub vendorid: String,
+    pub vendor_id: u16,
 
     /// The Device ID.
-    pub deviceid: String,
+    pub device_id: u16,
 }
 
 /// An error encountered while parsing a USB device identifier
@@ -80,8 +80,11 @@ pub enum DeviceIdentifierFromStrError {
     #[error("could not split input by colon; expected output similar to `lsusb`")]
     CouldNotSplitByColon,
 
-    #[error("found non-hex {0} at position {1}")]
-    BadChar(char, usize),
+    #[error("could not parse hex from vendor or device ID")]
+    Hex(#[from] hex::FromHexError),
+
+    #[error("could not decode u64 from bytes: {0:?}")]
+    BadBytes(Vec<u8>),
 }
 
 impl FromStr for DeviceIdentifier {
@@ -91,23 +94,22 @@ impl FromStr for DeviceIdentifier {
         let Some((first, last)) = s.split_once(':') else {
             return Err(DeviceIdentifierFromStrError::CouldNotSplitByColon);
         };
-        if let Some((position, ch)) = first
-            .chars()
-            .enumerate()
-            .find(|(_, ch)| !ch.is_ascii_hexdigit())
-        {
-            return Err(DeviceIdentifierFromStrError::BadChar(ch, position));
-        }
-        if let Some((position, ch)) = last
-            .chars()
-            .enumerate()
-            .find(|(_, ch)| !ch.is_ascii_hexdigit())
-        {
-            return Err(DeviceIdentifierFromStrError::BadChar(ch, position));
-        }
+
+        let vendor_id = u16::from_be_bytes(
+            hex::decode(first)?
+                .try_into()
+                .map_err(DeviceIdentifierFromStrError::BadBytes)?,
+        );
+
+        let device_id = u16::from_be_bytes(
+            hex::decode(last)?
+                .try_into()
+                .map_err(DeviceIdentifierFromStrError::BadBytes)?,
+        );
+
         Ok(Self {
-            vendorid: first.to_owned(),
-            deviceid: last.to_owned(),
+            vendor_id,
+            device_id,
         })
     }
 }

src/main.rs

@@ -1,3 +1,5 @@
+#![allow(clippy::redundant_else)]
+
 use clap::Parser;
 use eyre::WrapErr;
 use std::io::Write;
@@ -21,9 +23,12 @@ fn main() -> eyre::Result<()> {
     let opts = cli::App::parse_from(args);
 
     match opts.subcommand {
-        cli::Commands::Start => {
-            let spawn_args = SpawnArguments::default();
-            let mut vm = VirtualMachine::start(spawn_args)?;
+        cli::Commands::Start { memory } => {
+            let spawn_args = SpawnArguments {
+                memory: memory.clone(),
+                ..Default::default()
+            };
+            let vm = VirtualMachine::start(spawn_args)?;
             let pid = vm.pid();
             std::fs::write(&opts.lockfile, pid.to_string()).with_context(|| {
                 format!(
@@ -31,9 +36,6 @@ fn main() -> eyre::Result<()> {
                     lockfile = opts.lockfile.display(),
                 )
             })?;
-
-            // temp
-            vm.run_command("uptime", [])?;
         }
         cli::Commands::Stop => {
             let spawn_arguments = SpawnArguments::default();
@@ -41,8 +43,7 @@ fn main() -> eyre::Result<()> {
             vm.kill()?;
         }
         cli::Commands::Shell => {
-            // TODO: qemu inline, is it possible to pass through stdin/stdout w/o buffering?
-            todo!()
+            todo!("custom args to starting a VM and piping stdin/stdout are not yet implemented");
         }
         cli::Commands::Status => {
             let spawn_arguments = SpawnArguments::default();
@@ -57,7 +58,20 @@ fn main() -> eyre::Result<()> {
             eprintln!("hostname: {hostname}");
             eprint!("{}", String::from_utf8_lossy(&uptime.0));
         }
-        cli::Commands::Attach { device } => todo!(),
+        cli::Commands::Attach { device } => {
+            let spawn_arguments = SpawnArguments::default();
+            let mut vm = VirtualMachine::load(spawn_arguments, None)?;
+            vm.execute_host("qmp_capabilities", serde_json::json!({}))?;
+            vm.execute_host(
+                "device_add",
+                serde_json::json!({
+                    "driver": "usb-host",
+                    "bus": "usb.0",
+                    "vendorid": device.vendor_id,
+                    "productid": device.device_id,
+                }),
+            )?;
+        }
         cli::Commands::Push {
             local_path,
             remote_path,
@@ -79,21 +93,9 @@ fn main() -> eyre::Result<()> {
             let mut vm = VirtualMachine::load(spawn_arguments, None)?;
             let (response, exit_code) = vm.run_command(&command, args)?;
             std::io::stdout().write_all(&response)?;
-            std::process::exit(exit_code as i32);
-        }
-        cli::Commands::Test {} => {
-            let spawn_arguments = SpawnArguments::default();
-            let mut vm = VirtualMachine::load(spawn_arguments, None)?;
-            for i in 0..10 {
-                let sleep_command = format!("sleep 10; echo {i}");
-                let (response, exit_code) =
-                    vm.run_command("sh", [String::from("-c"), sleep_command])?;
-                eprint!(
-                    "exit code {}, output {}",
-                    exit_code,
-                    String::from_utf8_lossy(&response),
-                );
-            }
+            std::process::exit(
+                i32::try_from(exit_code).context(eyre::eyre!("bad PID: pid < i32::MAX << 1"))?,
+            );
         }
     }
 

src/vm.rs

@@ -30,7 +30,7 @@ fn spinner(msg: impl Display) -> ProgressBar {
 }
 
 fn bar(count: u64, msg: impl Display) -> ProgressBar {
-    let template = "[{elapsed_precise}] {wide_bar} {percent}% {msg}";
+    let template = "{elapsed_precise} [{wide_bar}] {percent}% {msg}";
     cfg_if::cfg_if! {
         if #[cfg(feature = "unicode")] {
             let style = ProgressStyle::with_template(template).unwrap();
@@ -66,6 +66,7 @@ fn to_lowercase_hexlike(s: impl AsRef<str>) -> String {
     s.to_ascii_lowercase()
 }
 
+#[allow(clippy::struct_field_names)]
 #[derive(Clone, Debug)]
 struct Device {
     vendor_id: u16,
@@ -100,7 +101,7 @@ fn find_pci_device_by_class(class: u16) -> Result<Vec<Device>> {
             let bus_id = bus_address
                 .into_string()
                 .map_err(|bad| eyre::eyre!("non-utf8 bus address: {bad:?}"))?
-                .split_once(":")
+                .split_once(':')
                 .ok_or(eyre::eyre!("bad path ID"))?
                 .1
                 .to_string();
@@ -116,19 +117,25 @@ fn find_pci_device_by_class(class: u16) -> Result<Vec<Device>> {
 }
 
 // NOTE: Do not implement `clone`, as there is side-effect state involved.
+
+/// A control handle for a virtual machine.
 #[derive(Debug)]
 pub struct VirtualMachine {
     pid: u32,
-    writer: UnixStream,
-    reader: BufReader<UnixStream>,
+    // qemu guest agent (proxied to guest)
+    guest_writer: UnixStream,
+    guest_reader: BufReader<UnixStream>,
+    // qemu machine protocol (host)
+    host_writer: UnixStream,
+    host_reader: BufReader<UnixStream>,
     args: SpawnArguments,
 }
 
 /// The configuration to use when starting a VM.
 #[derive(Clone, Debug)]
 pub struct SpawnArguments {
-    /// The PCI device to use for connecting to a network.
-    pub network_pci_device: Option<String>,
+    /// The amount of memory to allocate to a VM.
+    pub memory: String,
 
     /// The image file to use when booting the machine.
     ///
@@ -148,7 +155,7 @@ pub struct SpawnArguments {
 impl Default for SpawnArguments {
     fn default() -> Self {
         Self {
-            network_pci_device: None,
+            memory: String::from("1G"),
             guest_image: PathBuf::from("/guest.img"),
             guest_agent_socket_path: PathBuf::from("/var/run/netvm_qga.sock"),
             qmp_socket_path: PathBuf::from("/var/run/netvm_qmp.sock"),
@@ -158,6 +165,7 @@ impl Default for SpawnArguments {
 }
 
 impl VirtualMachine {
+    /// Start a virutal machine with the given parameters.
     pub fn start(args: SpawnArguments) -> eyre::Result<Self> {
         let eth_devices = find_pci_device_by_class(0x0200)?;
 
@@ -165,8 +173,9 @@ impl VirtualMachine {
         if std::fs::exists(&args.lockfile_path)? {
             // Check if VM is running
             use nix::unistd::{getpgid, Pid};
-            let pid = get_pid(&args.lockfile_path)?;
-            if getpgid(Some(Pid::from_raw(pid as i32))).is_ok() {
+            let pid = i32::try_from(get_pid(&args.lockfile_path)?)
+                .context(eyre::eyre!("bad PID: pid < i32::MAX << 1"))?;
+            if getpgid(Some(Pid::from_raw(pid))).is_ok() {
                 // process exists, exit
                 return Err(eyre::eyre!(
                     "VM with this configuration exists as PID {pid}"
@@ -213,14 +222,14 @@ impl VirtualMachine {
                 ),
             )?;
             net_args.push("-device".to_string());
-            net_args.push(format!("vfio-pci,host={bus_id}"))
+            net_args.push(format!("vfio-pci,host={bus_id}"));
         }
 
         let mut child = Command::new("qemu-system-x86_64")
             .stdin(Stdio::null())
             .stdout(Stdio::null())
             .stderr(Stdio::null())
-            .args(["-m", "4G"])
+            .args(["-m", &args.memory])
             .args(["-machine", "q35"])
             .arg("-nographic")
             .args(["-serial", "none"])
@@ -255,7 +264,9 @@ impl VirtualMachine {
                 return Err(eyre::eyre!("child exited with code {:?}", status.code()));
             }
 
-            if std::fs::exists(&args.guest_agent_socket_path)? {
+            if std::fs::exists(&args.guest_agent_socket_path)?
+                && std::fs::exists(&args.qmp_socket_path)?
+            {
                 break;
             }
             std::thread::sleep(Duration::from_millis(100));
@@ -266,25 +277,25 @@ impl VirtualMachine {
         Self::load(args, Some(child.id()))
     }
 
+    /// Load a virtual machine with the given parameters and optionally a custom PID.
+    ///
+    /// The custom PID option may be relevant if the virtual machine sockets were loaded but the
+    /// PID of the virtual machine was not properly persisted.
     pub fn load(args: SpawnArguments, pid: Option<u32>) -> Result<Self> {
         let bar = spinner("Connecting to VM");
-        let pid = match pid {
-            Some(pid) => pid,
-            None => {
-                let pid_str = std::fs::read_to_string(&args.lockfile_path)
-                    .context("error reading PID from lockfile")?;
-                pid_str.parse().context("could not parse PID")?
-            }
+        let pid = if let Some(pid) = pid {
+            pid
+        } else {
+            let pid_str = std::fs::read_to_string(&args.lockfile_path)
+                .context("error reading PID from lockfile")?;
+            pid_str.parse().context("could not parse PID")?
         };
 
-        let writer = UnixStream::connect(&args.guest_agent_socket_path)
+        let guest = UnixStream::connect(&args.guest_agent_socket_path)
             .context("could not open socket to QVM guest agent")?;
 
-        let reader = BufReader::new(
-            writer
-                .try_clone()
-                .context("couldn't clone socket to make buffered reader")?,
-        );
+        let host = UnixStream::connect(&args.qmp_socket_path)
+            .context("could not open socket to qemu management socket")?;
 
         bar.println(format!(
             "Connected to VM with PID {} and socket {}",
@@ -293,21 +304,39 @@ impl VirtualMachine {
         ));
         bar.finish_and_clear();
 
-        let vm = Self::from_parts(pid, writer, reader, args)?;
+        let vm = Self::from_parts(pid, guest, host, args)?;
 
         Ok(vm)
     }
 
     fn from_parts(
         pid: u32,
-        writer: UnixStream,
-        reader: BufReader<UnixStream>,
+        guest_socket: UnixStream,
+        host_socket: UnixStream,
         args: SpawnArguments,
     ) -> Result<Self> {
+        let guest_reader = BufReader::new(
+            guest_socket
+                .try_clone()
+                .context("couldn't clone socket to make buffered reader")?,
+        );
+        let mut host_reader = BufReader::new(
+            host_socket
+                .try_clone()
+                .context("couldn't clone socket to make buffered reader")?,
+        );
+
+        let mut server_hello = String::new();
+        host_reader
+            .read_line(&mut server_hello)
+            .context("can't read line from socket (pre-load)")?;
+
         let mut vm = Self {
             pid,
-            writer,
-            reader,
+            guest_writer: guest_socket,
+            guest_reader,
+            host_writer: host_socket,
+            host_reader,
             args,
         };
 
@@ -319,7 +348,7 @@ impl VirtualMachine {
         // crashing if those circumstances happen to be met.
         let time = SystemTime::now().duration_since(UNIX_EPOCH)?;
 
-        let identifier = time.as_secs() % (u32::MAX as u64);
+        let identifier = time.as_secs() % u64::from(u32::MAX);
 
         let ping_response = vm
             .execute_internal("guest-sync", serde_json::json!({"id": identifier}))
@@ -334,6 +363,7 @@ impl VirtualMachine {
         Ok(vm)
     }
 
+    /// The PID of the virtual machine.
     pub fn pid(&self) -> u32 {
         self.pid
     }
@@ -347,21 +377,21 @@ impl VirtualMachine {
         // * read a line from the parser to reset the input
 
         let bar = spinner("Re-establishing connection...");
-        self.writer
+        self.guest_writer
             .set_nonblocking(true)
             .context("flush: can't set nonblocking")?;
-        if let Err(e) = self.reader.read_to_end(&mut vec![]) {
+        if let Err(e) = self.guest_reader.read_to_end(&mut vec![]) {
             if e.kind() != std::io::ErrorKind::WouldBlock {
                 return Err(e).context("flush: can't read nonblocked data");
             }
         }
-        self.writer
+        self.guest_writer
             .set_nonblocking(false)
             .context("flush: can't set blocking")?;
-        self.writer
+        self.guest_writer
             .write_all(&[0x1b])
             .context("flush: can't send reset byte")?;
-        self.reader
+        self.guest_reader
             .read_line(&mut String::new())
             .context("flush: can't read error")?;
         bar.finish_and_clear();
@@ -369,6 +399,7 @@ impl VirtualMachine {
         Ok(())
     }
 
+    /// Push a single file from the local machine to the VM.
     pub fn push(
         &mut self,
         local_path: impl AsRef<Path>,
@@ -421,7 +452,8 @@ impl VirtualMachine {
                     .ok_or(eyre::eyre!("not given 'count' of bytes written"))?
                     .as_u64()
                     .ok_or(eyre::eyre!("'count' not u64"))?;
-                written += response_written as usize;
+                written +=
+                    usize::try_from(response_written).expect("wrote more than u46::MAX bytes");
                 if written == size {
                     break;
                 }
@@ -437,6 +469,8 @@ impl VirtualMachine {
         Ok(())
     }
 
+    /// Pull a single file from the VM to the local machine. This operation is destructive and will
+    /// overwrite existing files.
     pub fn pull(
         &mut self,
         remote_path: impl AsRef<Path>,
@@ -534,6 +568,8 @@ impl VirtualMachine {
 
     // TODO: make this return status, stdout, stderr
     // TODO: accept optional: env, input-data, disable capture-output
+    /// Run a command on the virtual machine. Standard input is not sent to the process, and only
+    /// standard output is received from the process.
     pub fn run_command(
         &mut self,
         command: &str,
@@ -599,12 +635,12 @@ impl VirtualMachine {
             "arguments": args,
         });
 
-        serde_json::to_writer(&mut self.writer, &message)
+        serde_json::to_writer(&mut self.guest_writer, &message)
             .context("could not send message over socket")?;
-        writeln!(&mut self.writer).context("could not send newline over socket")?;
+        writeln!(&mut self.guest_writer).context("could not send newline over socket")?;
 
         let mut line = String::new();
-        self.reader
+        self.guest_reader
             .read_line(&mut line)
             .context("can't read line from socket")?;
 
@@ -620,6 +656,7 @@ impl VirtualMachine {
         }
     }
 
+    /// Execute an operation via QEMU Guest Agent. This modifies state inside the VM.
     pub fn execute<S: serde::Serialize + Debug>(
         &mut self,
         command: &'static str,
@@ -631,16 +668,56 @@ impl VirtualMachine {
         result
     }
 
+    /// Execute an operation via QEMU Machine Protocol. This modifies state on the host machine and
+    /// the VM.
+    pub fn execute_host<S: serde::Serialize + Debug>(
+        &mut self,
+        command: &'static str,
+        args: S,
+    ) -> Result<serde_json::Value> {
+        let bar = spinner(format!("Executing: {command:?} with {args:?}"));
+
+        let message = serde_json::json!({
+            "execute": command,
+            "arguments": args,
+        });
+
+        serde_json::to_writer(&mut self.host_writer, &message)
+            .context("could not send message over socket")?;
+        writeln!(&mut self.host_writer).context("could not send newline over socket")?;
+
+        let mut line = String::new();
+        self.host_reader
+            .read_line(&mut line)
+            .context("can't read line from socket")?;
+
+        let response: serde_json::Value =
+            serde_json::from_str(&line).context("response from qemu is not json")?;
+
+        bar.finish_and_clear();
+
+        if let Some(response) = response.get("return") {
+            Ok(response.clone())
+        } else if let Some(error) = response.get("error") {
+            Err(eyre::eyre!("error response from qemu: {error:?}"))
+        } else {
+            Err(eyre::eyre!("invalid response from qemu: {response:?}"))
+        }
+    }
+
     // NOTE: u32 is returned from Process::id(), i32 is the Linux internal version
     // This should be safe; the kernel wouldn't give a value that, when converted
     // to a u32, can't be made back into an i32
+    /// Terminate the VM and remove any stateful files.
     pub fn kill(self) -> Result<()> {
         use nix::{
             errno::Errno,
             sys::signal::{kill, SIGKILL},
             unistd::{getpgid, Pid},
         };
-        let pid = Pid::from_raw(self.pid as i32);
+        let pid = Pid::from_raw(
+            i32::try_from(self.pid).context(eyre::eyre!("bad PID: pid < i32::MAX << 1"))?,
+        );
         if getpgid(Some(pid)).is_err() {
             eprintln!("Process not found");
             return Ok(());