VirtWeb/virtweb_backend/src/actors/libvirt_actor.rs

use crate::app_config::AppConfig;
use crate::libvirt_lib_structures::{DomainState, DomainXML, DomainXMLUuid};
use crate::libvirt_rest_structures::*;
use actix::{Actor, Context, Handler, Message};
use image::ImageOutputFormat;
use std::io::Cursor;
use virt::connect::Connect;
use virt::domain::Domain;
use virt::stream::Stream;
use virt::sys;
use virt::sys::VIR_DOMAIN_XML_SECURE;

pub struct LibVirtActor {
    m: Connect,
}

impl LibVirtActor {
    /// Connect to hypervisor
    pub async fn connect() -> anyhow::Result<Self> {
        let hypervisor_uri = AppConfig::get().hypervisor_uri.as_deref().unwrap_or("");
        log::info!(
            "Will connect to hypvervisor at address '{}'",
            hypervisor_uri
        );
        let conn = Connect::open(hypervisor_uri)?;

        Ok(Self { m: conn })
    }
}

impl Actor for LibVirtActor {
    type Context = Context<Self>;
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<HypervisorInfo>")]
pub struct GetHypervisorInfo;

impl Handler<GetHypervisorInfo> for LibVirtActor {
    type Result = anyhow::Result<HypervisorInfo>;

    fn handle(&mut self, _msg: GetHypervisorInfo, _ctx: &mut Self::Context) -> Self::Result {
        let node = self.m.get_node_info()?;
        Ok(HypervisorInfo {
            r#type: self.m.get_type()?,
            hyp_version: self.m.get_hyp_version()?,
            lib_version: self.m.get_lib_version()?,
            capabilities: self.m.get_capabilities()?,
            free_memory: self.m.get_free_memory()?,
            hostname: self.m.get_hostname()?,
            node: HypervisorNodeInfo {
                cpu_model: node.model,
                memory_size: node.memory,
                number_of_active_cpus: node.cpus,
                cpu_frequency_mhz: node.mhz,
                number_of_numa_cell: node.nodes,
                number_of_cpu_socket_per_node: node.sockets,
                number_of_core_per_sockets: node.cores,
                number_of_threads_per_core: node.threads,
            },
        })
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<Vec<DomainXMLUuid>>")]
pub struct GetDomainsListReq;

impl Handler<GetDomainsListReq> for LibVirtActor {
    type Result = anyhow::Result<Vec<DomainXMLUuid>>;

    fn handle(&mut self, _msg: GetDomainsListReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Get full list of domains");
        let domains = self.m.list_all_domains(0)?;
        let mut ids = Vec::with_capacity(domains.len());

        for d in domains {
            ids.push(DomainXMLUuid::parse_from_str(&d.get_uuid_string()?)?);
        }

        Ok(ids)
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<DomainXML>")]
pub struct GetDomainXMLReq(pub DomainXMLUuid);

impl Handler<GetDomainXMLReq> for LibVirtActor {
    type Result = anyhow::Result<DomainXML>;

    fn handle(&mut self, msg: GetDomainXMLReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Get domain XML:\n{}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        let xml = domain.get_xml_desc(VIR_DOMAIN_XML_SECURE)?;
        log::debug!("XML = {}", xml);
        Ok(serde_xml_rs::from_str(&xml)?)
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<DomainXMLUuid>")]
pub struct DefineDomainReq(pub DomainXML);

impl Handler<DefineDomainReq> for LibVirtActor {
    type Result = anyhow::Result<DomainXMLUuid>;

    fn handle(&mut self, mut msg: DefineDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        // A issue with the disks definition serialization needs them to be serialized aside
        let mut disks_xml = Vec::with_capacity(msg.0.devices.disks.len());
        for disk in msg.0.devices.disks {
            let disk_xml = serde_xml_rs::to_string(&disk)?;
            let start_offset = disk_xml.find("<disk").unwrap();
            disks_xml.push(disk_xml[start_offset..].to_string());
        }
        msg.0.devices.disks = vec![];

        let mut xml = serde_xml_rs::to_string(&msg.0)?;
        let disks_xml = disks_xml.join("\n");
        xml = xml.replacen("<devices>", &format!("<devices>{disks_xml}"), 1);

        log::debug!("Define domain:\n{}", xml);
        let domain = Domain::define_xml(&self.m, &xml)?;
        DomainXMLUuid::parse_from_str(&domain.get_uuid_string()?)
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct DeleteDomainReq {
    pub id: DomainXMLUuid,
    pub keep_files: bool,
}

impl Handler<DeleteDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: DeleteDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!(
            "Delete domain: {:?} (keep files: {})",
            msg.id,
            msg.keep_files
        );

        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.id.as_string())?;
        domain.undefine_flags(match msg.keep_files {
            true => sys::VIR_DOMAIN_UNDEFINE_KEEP_NVRAM,
            false => sys::VIR_DOMAIN_UNDEFINE_NVRAM,
        })?;

        if !msg.keep_files {
            log::info!("Delete storage associated with the domain");
            let path = AppConfig::get().vm_storage_path(msg.id);
            if path.exists() {
                std::fs::remove_dir_all(path)?;
            }
        }

        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<DomainState>")]
pub struct GetDomainStateReq(pub DomainXMLUuid);

impl Handler<GetDomainStateReq> for LibVirtActor {
    type Result = anyhow::Result<DomainState>;

    fn handle(&mut self, msg: GetDomainStateReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Get domain state: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        let (state, _) = domain.get_state()?;
        Ok(match state {
            sys::VIR_DOMAIN_NOSTATE => DomainState::NoState,
            sys::VIR_DOMAIN_RUNNING => DomainState::Running,
            sys::VIR_DOMAIN_BLOCKED => DomainState::Blocked,
            sys::VIR_DOMAIN_PAUSED => DomainState::Paused,
            sys::VIR_DOMAIN_SHUTDOWN => DomainState::Shutdown,
            sys::VIR_DOMAIN_SHUTOFF => DomainState::Shutoff,
            sys::VIR_DOMAIN_CRASHED => DomainState::Crashed,
            sys::VIR_DOMAIN_PMSUSPENDED => DomainState::PowerManagementSuspended,

            _ => DomainState::Other,
        })
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct StartDomainReq(pub DomainXMLUuid);

impl Handler<StartDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: StartDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Start domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        domain.create()?;
        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct ShutdownDomainReq(pub DomainXMLUuid);

impl Handler<ShutdownDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: ShutdownDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Shutdown domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        domain.shutdown()?;
        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct KillDomainReq(pub DomainXMLUuid);

impl Handler<KillDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: KillDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Kill domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        domain.destroy()?;
        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct ResetDomainReq(pub DomainXMLUuid);

impl Handler<ResetDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: ResetDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Reset domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        domain.reset()?;
        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct SuspendDomainReq(pub DomainXMLUuid);

impl Handler<SuspendDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: SuspendDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Suspend domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        domain.suspend()?;
        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<()>")]
pub struct ResumeDomainReq(pub DomainXMLUuid);

impl Handler<ResumeDomainReq> for LibVirtActor {
    type Result = anyhow::Result<()>;

    fn handle(&mut self, msg: ResumeDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Resume domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;
        domain.resume()?;
        Ok(())
    }
}

#[derive(Message)]
#[rtype(result = "anyhow::Result<Vec<u8>>")]
pub struct ScreenshotDomainReq(pub DomainXMLUuid);

impl Handler<ScreenshotDomainReq> for LibVirtActor {
    type Result = anyhow::Result<Vec<u8>>;

    fn handle(&mut self, msg: ScreenshotDomainReq, _ctx: &mut Self::Context) -> Self::Result {
        log::debug!("Take screenshot of domain: {}", msg.0.as_string());
        let domain = Domain::lookup_by_uuid_string(&self.m, &msg.0.as_string())?;

        let stream = Stream::new(&self.m, 0)?;
        domain.screenshot(&stream, 0, 0)?;

        let mut screen_out = Vec::with_capacity(1000000);
        let mut buff = [0u8; 1000];
        loop {
            let size = stream.recv(&mut buff)?;

            if size == 0 {
                break;
            }
            screen_out.extend_from_slice(&buff[0..size]);
        }

        let image = image::load_from_memory(&screen_out)?;
        let mut png_out = Cursor::new(Vec::new());
        image.write_to(&mut png_out, ImageOutputFormat::Png)?;

        Ok(png_out.into_inner())
    }
}