photoncloud-monorepo/plasmavmc/crates/plasmavmc-kvm/src/qmp.rs

use std::path::Path;

use serde::Serialize;
use serde_json::Value;
use tokio::{
    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
    net::UnixStream,
};

use plasmavmc_types::{Error, Result, VmState, VmStatus};

/// Minimal async QMP client for lifecycle operations.
pub struct QmpClient {
    reader: BufReader<tokio::net::unix::OwnedReadHalf>,
    writer: tokio::net::unix::OwnedWriteHalf,
}

impl QmpClient {
    /// Connect to a QMP Unix socket and negotiate capabilities.
    pub async fn connect(path: impl AsRef<Path>) -> Result<Self> {
        let stream = UnixStream::connect(path.as_ref())
            .await
            .map_err(|e| Error::HypervisorError(format!("Failed to connect QMP: {e}")))?;
        Self::from_stream(stream).await
    }

    async fn from_stream(stream: UnixStream) -> Result<Self> {
        let (read, write) = stream.into_split();
        let mut client = QmpClient {
            reader: BufReader::new(read),
            writer: write,
        };

        client.read_greeting().await?;
        // Negotiate capabilities per QMP handshake.
        client.command::<Value>("qmp_capabilities", None::<Value>).await?;

        Ok(client)
    }

    /// Send an arbitrary QMP command with optional arguments.
    pub async fn command<A: Serialize>(
        &mut self,
        name: &str,
        args: Option<A>,
    ) -> Result<Value> {
        let mut payload = serde_json::json!({ "execute": name });
        if let Some(arguments) = args {
            payload["arguments"] = serde_json::to_value(arguments).map_err(|e| {
                Error::HypervisorError(format!("Failed to serialize QMP args: {e}"))
            })?;
        }

        let mut buf = serde_json::to_vec(&payload)
            .map_err(|e| Error::HypervisorError(format!("Failed to encode QMP command: {e}")))?;
        buf.push(b'\n');
        self.writer
            .write_all(&buf)
            .await
            .map_err(|e| Error::HypervisorError(format!("Failed to send QMP command: {e}")))?;
        self.writer
            .flush()
            .await
            .map_err(|e| Error::HypervisorError(format!("Failed to flush QMP command: {e}")))?;

        let response = self.read_message().await?;
        if let Some(error) = response.get("error") {
            let desc = error
                .get("desc")
                .and_then(Value::as_str)
                .unwrap_or("unknown error");
            return Err(Error::HypervisorError(format!(
                "QMP command {name} failed: {desc}"
            )));
        }

        response.get("return").cloned().ok_or_else(|| {
            Error::HypervisorError(format!(
                "Unexpected QMP response for {name}: {}",
                response
            ))
        })
    }

    /// Query VM status and map to VmStatus.
    pub async fn query_status(&mut self) -> Result<VmStatus> {
        let resp = self
            .command::<Value>("query-status", None::<Value>)
            .await?;
        let status = resp
            .get("status")
            .and_then(Value::as_str)
            .unwrap_or("unknown");

        let mapped_state = match status {
            "running" => VmState::Running,
            "paused" => VmState::Stopped,
            "shutdown" | "quit" => VmState::Stopped,
            "inmigrate" | "postmigrate" => VmState::Migrating,
            "watchdog" | "guest-panicked" | "internal-error" | "io-error" => VmState::Error,
            _ => VmState::Error,
        };

        Ok(VmStatus {
            actual_state: mapped_state,
            ..VmStatus::default()
        })
    }

    async fn read_greeting(&mut self) -> Result<()> {
        let greeting = self.read_message().await?;
        if greeting.get("QMP").is_some() {
            Ok(())
        } else {
            Err(Error::HypervisorError(format!(
                "Invalid QMP greeting: {greeting}"
            )))
        }
    }

    async fn read_message(&mut self) -> Result<Value> {
        loop {
            let mut line = String::new();
            let read = self
                .reader
                .read_line(&mut line)
                .await
                .map_err(|e| Error::HypervisorError(format!("Failed to read QMP: {e}")))?;
            if read == 0 {
                return Err(Error::HypervisorError(
                    "QMP connection closed".to_string(),
                ));
            }

            if line.trim().is_empty() {
                continue;
            }

            let value: Value = serde_json::from_str(&line).map_err(|e| {
                Error::HypervisorError(format!("Failed to parse QMP message: {e}: {line}"))
            })?;

            // Skip async events; return first reply.
            if value.get("event").is_some() {
                continue;
            }

            return Ok(value);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::{
        io::BufReader,
        net::{UnixListener, UnixStream},
    };

    async fn spawn_qmp_server(socket_path: &Path) -> tokio::task::JoinHandle<()> {
        let listener = UnixListener::bind(socket_path).expect("bind qmp socket");
        let socket_path = socket_path.to_owned();
        tokio::spawn(async move {
            let (stream, _) = listener.accept().await.expect("accept connection");
            handle_qmp_session(stream).await;
            // Clean up socket file to avoid leaking between tests.
            let _ = std::fs::remove_file(&socket_path);
        })
    }

    async fn handle_qmp_session(stream: UnixStream) {
        let (read, mut write) = stream.into_split();
        // Send greeting first.
        let greeting = r#"{"QMP":{"version":{"qemu":{"major":8,"minor":0,"micro":0}},"capabilities":[]}}"#;
        write.write_all(greeting.as_bytes()).await.unwrap();
        write.write_all(b"\n").await.unwrap();
        write.flush().await.unwrap();

        let mut reader = BufReader::new(read);
        let mut line = String::new();

        // Expect qmp_capabilities.
        line.clear();
        reader.read_line(&mut line).await.unwrap();
        assert!(
            line.contains("qmp_capabilities"),
            "expected qmp_capabilities, got {line}"
        );
        write.write_all(br#"{"return":{}}"#).await.unwrap();
        write.write_all(b"\n").await.unwrap();
        write.flush().await.unwrap();

        // Next command (query-status) will be handled in tests.
        line.clear();
        reader.read_line(&mut line).await.unwrap();
        if line.contains("query-status") {
            write
                .write_all(br#"{"return":{"status":"running","running":true}}"#)
                .await
                .unwrap();
            write.write_all(b"\n").await.unwrap();
            write.flush().await.unwrap();
        } else {
            write
                .write_all(br#"{"error":{"class":"CommandNotFound","desc":"unexpected"}}"#)
                .await
                .unwrap();
            write.write_all(b"\n").await.unwrap();
            write.flush().await.unwrap();
        }
    }

    #[tokio::test]
    async fn qmp_client_performs_handshake_and_status() {
        let dir = tempfile::tempdir().unwrap();
        let socket_path = dir.path().join("qmp.sock");
        let server_handle = spawn_qmp_server(&socket_path).await;

        let mut client = QmpClient::connect(&socket_path).await.unwrap();
        let status = client.query_status().await.unwrap();
        assert_eq!(status.actual_state, VmState::Running);

        server_handle.await.unwrap();
    }

    #[tokio::test]
    async fn qmp_client_reports_command_error() {
        let dir = tempfile::tempdir().unwrap();
        let socket_path = dir.path().join("qmp.sock");

        // Server will only understand qmp_capabilities and then respond error for others.
        let listener = UnixListener::bind(&socket_path).unwrap();
        let server = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let (read, mut write) = stream.into_split();
            let greeting = r#"{"QMP":{"version":{"qemu":{"major":8,"minor":0,"micro":0}},"capabilities":[]}}"#;
            write.write_all(greeting.as_bytes()).await.unwrap();
            write.write_all(b"\n").await.unwrap();
            write.flush().await.unwrap();

            let mut reader = BufReader::new(read);
            let mut line = String::new();
            reader.read_line(&mut line).await.unwrap(); // qmp_capabilities
            write.write_all(br#"{"return":{}}"#).await.unwrap();
            write.write_all(b"\n").await.unwrap();
            write.flush().await.unwrap();

            line.clear();
            reader.read_line(&mut line).await.unwrap(); // query-status
            write
                .write_all(br#"{"error":{"class":"GenericError","desc":"bad command"}}"#)
                .await
                .unwrap();
            write.write_all(b"\n").await.unwrap();
            write.flush().await.unwrap();
        });

        let mut client = QmpClient::connect(&socket_path).await.unwrap();
        let err = client.query_status().await.unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("failed"), "unexpected error: {msg}");
        server.await.unwrap();
    }
}