system-harness 1.0.0

use crate::{Error, ErrorKind, EventPublisher, EventSubscriber, Key, Status, SystemHarness, SystemTerminal};
use crate::SystemConfig;
use cmdstruct::Command;
use serde::{Deserialize, Serialize};
use std::process::Child;
use std::pin::Pin;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::net::UnixStream;

mod args;

mod models;
use models::*;

mod qmp;
use qmp::QmpStream;

fn qemu_system_bin(config: &QemuSystemConfig) -> String {
    format!("qemu-system-{}", config.arch)
}

/// A configuration for running QEMU
///
/// This config can be serialized and deserialized using
/// serde.
#[derive(Clone, Command, Serialize, Deserialize)]
#[command(executable_fn = qemu_system_bin)]
pub struct QemuSystemConfig {
    arch: String,

    #[arg(option = "-boot")]
    boot: Option<Boot>,

    #[arg(option = "-cpu")]
    cpu: Option<String>,

    #[arg(option = "-machine")]
    machine: Option<Machine>,

    #[arg(option = "-smp")]
    smp: Option<Smp>,

    #[arg(option = "-accel")]
    accel: Option<String>,

    #[arg(option = "-bios")]
    bios: Option<String>,

    #[arg(option = "-m")]
    memory: Option<usize>,

    #[arg(option = "-cdrom")]
    cdrom: Option<String>,

    #[arg(option = "-hda")]
    hda: Option<String>,

    #[arg(option = "-hdb")]
    hdb: Option<String>,

    #[arg(option = "-device")]
    device: Option<Vec<Device>>,
    
    #[arg(option = "-fsdev")]
    fsdev: Option<Vec<Backend<FsDev>>>,

    #[arg(option = "-chardev")]
    chardev: Option<Vec<Backend<CharDev>>>,

    #[arg(option = "-netdev")]
    netdev: Option<Vec<Backend<NetDev>>>,

    #[arg(option = "-blockdev")]
    blockdev: Option<Vec<BlockDev>>,

    /// Extra QEMU args
    extra_args: Option<Vec<String>>
}

impl<'sys> SystemConfig<'sys> for QemuSystemConfig {

    type System = QemuSystem;

    async fn spawn(&self) -> Result<QemuSystem, Error> {
        let mut command = self.command();
        command.arg("-nographic");
        command.args(["-qmp", "unix:qmp.sock,server=on,wait=off"]);
        command.args(["-serial", "unix:serial.sock,server=on,wait=off"]);

        if let Some(extra_args) = &self.extra_args {
            command.args(extra_args);
        }

        log::trace!("Starting system...");
        let mut process = command.spawn()?;

        log::trace!("Connecting to QMP socket...");
        let mut qmp_socket = None;
        while process.try_wait()?.is_none() && qmp_socket.is_none() {
            qmp_socket = UnixStream::connect("qmp.sock").await.ok();
        }
        let qmp = QmpStream::new(qmp_socket.unwrap()).await?;
        log::trace!("Connecting to serial socket...");
        let serial = UnixStream::connect("serial.sock").await?;
        log::trace!("System ready.");
        Ok(QemuSystem {
            process,
            serial,
            qmp,
        })
    }

}

/// A running QEMU system
pub struct QemuSystem {
    process: Child,
    serial: UnixStream,
    qmp: QmpStream,
}

pub struct QemuSystemTerminal<'sys> {
    system: &'sys mut QemuSystem 
}

impl AsyncRead for QemuSystemTerminal<'_> {
    fn poll_read(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        Pin::new(&mut self.system.serial).poll_read(cx, buf)
    }
}

impl AsyncWrite for QemuSystemTerminal<'_> {
    fn poll_write(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        Pin::new(&mut self.system.serial).poll_write(cx, buf)
    }

    fn poll_flush(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), std::io::Error>> {
        Pin::new(&mut self.system.serial).poll_flush(cx)
    }

    fn poll_shutdown(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), std::io::Error>> {
        Pin::new(&mut self.system.serial).poll_shutdown(cx)
    }
}

impl SystemTerminal for QemuSystemTerminal<'_> {

    async fn send_key(&mut self, key: Key) -> Result<(), Error> {
        self.system
            .qmp
            .send_command(qmp::QmpCommand::SendKey(qmp::KeyCommand {
                keys: vec![key.into()],
            }))
            .await
            .map(|_| ())
    }

}

impl<'sys> SystemHarness<'sys> for QemuSystem {

    type Terminal = QemuSystemTerminal<'sys>;

    async fn terminal(&'sys mut self) -> Result<Self::Terminal, Error> {
        Ok(QemuSystemTerminal {
            system: self
        })
    }

    async fn running(&mut self) -> Result<bool, Error> {
        self.process
            .try_wait()
            .map(|status| status == None)
            .map_err(|err| err.into())
    }

    async fn pause(&mut self) -> Result<(), Error> {
        self.qmp.send_command(qmp::QmpCommand::Stop).await.map(|_| ())
    }

    async fn resume(&mut self) -> Result<(), Error> {
        self.qmp.send_command(qmp::QmpCommand::Cont).await.map(|_| ())
    }

    async fn shutdown(&mut self) -> Result<(), Error> {
        self.qmp
            .send_command(qmp::QmpCommand::SystemPowerdown)
            .await
            .map(|_| ())
    }

    async fn status(&mut self) -> Result<Status, Error> {
        self.qmp
            .send_command(qmp::QmpCommand::QueryStatus)
            .await
            .and_then(|ret| match ret {
                qmp::QmpReturn::StatusInfo(status) => status.try_into(),
                ret => Err(Error::new(
                    ErrorKind::HarnessError,
                    format!("Unexpected return: {ret:?}"),
                )),
            })
    }
}

impl Drop for QemuSystem {
    fn drop(&mut self) {
        if let Err(err) = self.process.kill() {
            log::error!("{err}")
        }
    }
}

impl EventPublisher for QemuSystem {
    fn subscribe(&mut self, subscriber: impl EventSubscriber) -> Result<(), Error> {
        self.qmp.subscribe(subscriber)
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn json_config() {
        const JSON_CONFIG: &'static str = include_str!("../tests/data/qemu-config.json");
        let config: QemuSystemConfig = serde_json::from_str(JSON_CONFIG).unwrap();
        let command = config.command();
        assert_eq!("qemu-system-i386", command.get_program());
        assert_eq!(
            vec!["-machine", "type=q35", "-m", "512",
                "-device", "driver=virtio-blk,drive=f1",
                "-blockdev", "driver=file,node-name=f1,filename=tests/data/test.raw"],
            command.get_args().collect::<Vec<_>>()
        );
    }
}