feat(VM): vm utils crate [fixes BRND-19] (#2311)

Cargo.toml

@@ -49,6 +49,7 @@ members = [
     "crates/types",
     "crates/core-distributor",
     "crates/log-format",
+    "crates/vm-utils",
 ]
 exclude = [
     "nox/tests/tetraplets",
@@ -182,6 +183,7 @@ alloy_serde_macro = "0.1.2"
 const-hex = "1.11.3"
 bytesize = "1.3.0"
 cfg-if = "1.0.0"
+nonempty = "0.9.0"
 
 [profile.dev]
 opt-level = 0

crates/core-distributor/Cargo.toml

@@ -18,7 +18,7 @@ toml = { workspace = true }
 multimap = { version = "0.10.0", features = ["serde"] }
 bimap = { version = "0.6.3", features = ["serde"] }
 newtype_derive = "0.1.6"
-nonempty = "0.9.0"
+nonempty.workspace = true
 tokio = { workspace = true, features = ["fs", "rt", "sync", "macros", "tracing"] }
 async-trait.workspace = true
 enum_dispatch.workspace = true

crates/vm-utils/.gitattributes

@@ -0,0 +1 @@
+*.qcow2 filter=lfs diff=lfs merge=lfs -text

crates/vm-utils/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "vm-utils"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+nonempty.workspace = true
+thiserror.workspace = true
+tracing.workspace = true
+rand.workspace = true
+mac_address = "1.1.7"
+virt = "0.3.1"
+
+[dev-dependencies]
+log-utils.workspace = true

crates/vm-utils/src/lib.rs

@@ -0,0 +1,7 @@
+mod vm_utils;
+
+pub use vm_utils::create_vm;
+pub use vm_utils::start_vm;
+pub use vm_utils::stop_vm;
+pub use vm_utils::CreateVMParams;
+pub use vm_utils::VMUtilsError;

crates/vm-utils/src/template.xml

@@ -0,0 +1,32 @@
+<domain type='kvm'>
+    <name>{}</name>
+    <memory unit='KiB'>{}</memory>
+    <vcpu placement='static'>{}</vcpu>
+    <cputune>
+        {}
+    </cputune>
+    <os>
+        <type arch='x86_64' machine='pc-i440fx-2.9'>hvm</type>
+        <boot dev='hd'/>
+    </os>
+    <devices>
+        <disk type='file' device='disk'>
+            <driver name='qemu' type='qcow2'/>
+            <source file='{}'/>
+            <target dev='vda' bus='virtio'/>
+            <address type='pci' domain='0x0000' bus='0x00' slot='0x04' function='0x0'/>
+        </disk>
+        <interface type='bridge'>
+            <mac address='{}'/>
+            <source bridge='br422442'/>
+            <model type='virtio'/>
+            <address type='pci' domain='0x0000' bus='0x00' slot='0x03' function='0x0'/>
+        </interface>
+        <console type='pty'>
+            <target type='serial' port='0'/>
+        </console>
+        <serial type='pty'>
+            <target port='0'/>
+        </serial>
+    </devices>
+</domain>

crates/vm-utils/src/vm_utils.rs

@@ -0,0 +1,196 @@
+use mac_address::MacAddress;
+use nonempty::NonEmpty;
+use rand::Rng;
+use std::path::PathBuf;
+use thiserror::Error;
+use virt::connect::Connect;
+use virt::domain::Domain;
+use virt::sys::VIR_DOMAIN_DEFINE_VALIDATE;
+
+const MAC_PREFIX: [u8; 3] = [0x52, 0x54, 0x00];
+
+#[derive(Debug)]
+pub struct CreateVMParams {
+    name: String,
+    image: PathBuf,
+    cpus: NonEmpty<u32>,
+}
+
+#[derive(Error, Debug)]
+pub enum VMUtilsError {
+    #[error("Failed to connect to the hypervisor")]
+    FailedToConnect {
+        #[source]
+        err: virt::error::Error,
+    },
+    #[error("Failed to create VM")]
+    FailedToCreateVM {
+        #[source]
+        err: virt::error::Error,
+    },
+    #[error("Could not find VM with name {name}")]
+    VmNotFound {
+        name: String,
+        #[source]
+        err: virt::error::Error,
+    },
+    #[error("Failed to shutdown VM")]
+    FailedToShutdownVM {
+        #[source]
+        err: virt::error::Error,
+    },
+    #[error("Failed to get id for VM with name {name}")]
+    FailedToGetVMId { name: String },
+}
+
+pub fn create_vm(uri: &str, params: CreateVMParams) -> Result<(), VMUtilsError> {
+    let conn = Connect::open(uri).map_err(|err| VMUtilsError::FailedToConnect { err })?;
+    let domain = Domain::lookup_by_name(&conn, params.name.as_str()).ok();
+
+    match domain {
+        None => {
+            tracing::info!(target: "vm-utils","Domain with name {} doesn't exists. Creating", params.name);
+            let mac = generate_random_mac();
+            let xml = prepare_xml(&params, mac.to_string().as_str());
+            Domain::define_xml_flags(&conn, xml.as_str(), VIR_DOMAIN_DEFINE_VALIDATE)
+                .map_err(|err| VMUtilsError::FailedToCreateVM { err })?;
+        }
+        Some(_) => {
+            tracing::info!(target: "vm-utils","Domain with name {} already exists. Skipping", params.name);
+        }
+    };
+    Ok(())
+}
+
+fn generate_random_mac() -> MacAddress {
+    let mut rng = rand::thread_rng();
+    let mut result = [0u8; 6];
+    result[..3].copy_from_slice(&MAC_PREFIX);
+    rng.fill(&mut result[3..]);
+    MacAddress::from(result)
+}
+
+pub fn start_vm(uri: &str, name: String) -> Result<u32, VMUtilsError> {
+    tracing::info!(target: "vm-utils","Starting VM with name {name}");
+    let conn = Connect::open(uri).map_err(|err| VMUtilsError::FailedToConnect { err })?;
+    let domain =
+        Domain::lookup_by_name(&conn, name.as_str()).map_err(|err| VMUtilsError::VmNotFound {
+            name: name.clone(),
+            err,
+        })?;
+    domain
+        .create()
+        .map_err(|err| VMUtilsError::FailedToCreateVM { err })?;
+
+    let id = domain
+        .get_id()
+        .ok_or(VMUtilsError::FailedToGetVMId { name })?;
+
+    Ok(id)
+}
+
+pub fn stop_vm(uri: &str, name: String) -> Result<(), VMUtilsError> {
+    tracing::info!(target: "vm-utils","Stopping VM with name {name}");
+    let conn = Connect::open(uri).map_err(|err| VMUtilsError::FailedToConnect { err })?;
+    let domain = Domain::lookup_by_name(&conn, name.as_str())
+        .map_err(|err| VMUtilsError::VmNotFound { name, err })?;
+    domain
+        .shutdown()
+        .map_err(|err| VMUtilsError::FailedToShutdownVM { err })?;
+
+    Ok(())
+}
+
+fn prepare_xml(params: &CreateVMParams, mac_address: &str) -> String {
+    let mut mapping = String::new();
+    for (index, logical_id) in params.cpus.iter().enumerate() {
+        if index > 0 {
+            mapping.push_str("\n        ");
+        }
+        mapping.push_str(format!("<vcpupin vcpu='{index}' cpuset='{logical_id}'/>").as_str());
+    }
+    let memory_in_kb = params.cpus.len() * 4 * 1024 * 1024; // 4Gbs per core
+    format!(
+        include_str!("template.xml"),
+        params.name,
+        memory_in_kb,
+        params.cpus.len(),
+        mapping,
+        params.image.display(),
+        mac_address
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use nonempty::nonempty;
+    use std::fs;
+    const DEFAULT_URI: &str = "test:///default";
+
+    fn list_defined() -> Result<Vec<String>, VMUtilsError> {
+        let conn =
+            Connect::open(DEFAULT_URI).map_err(|err| VMUtilsError::FailedToConnect { err })?;
+        Ok(conn.list_defined_domains().unwrap())
+    }
+
+    fn list() -> Result<Vec<u32>, VMUtilsError> {
+        let conn =
+            Connect::open(DEFAULT_URI).map_err(|err| VMUtilsError::FailedToConnect { err })?;
+        Ok(conn.list_domains().unwrap())
+    }
+
+    #[test]
+    fn test_prepare_xml() {
+        let xml = prepare_xml(
+            &CreateVMParams {
+                name: "test-id".to_string(),
+                image: "test-image".into(),
+                cpus: nonempty![1, 8],
+            },
+            "52:54:00:1e:af:64",
+        );
+        assert_eq!(xml, include_str!("../tests/expected_vm_config.xml"))
+    }
+
+    #[test]
+    fn test_vm_creation() {
+        log_utils::enable_logs();
+
+        let image: PathBuf = "./tests/alpine-virt-3.20.1-x86_64.qcow2".into();
+        let image = fs::canonicalize(image).unwrap();
+
+        let list_before_create = list().unwrap();
+        let list_defined_before_create = list_defined().unwrap();
+        assert!(list_defined_before_create.is_empty());
+
+        create_vm(
+            DEFAULT_URI,
+            CreateVMParams {
+                name: "test-id".to_string(),
+                image: image.clone(),
+                cpus: nonempty![1],
+            },
+        )
+        .unwrap();
+
+        let list_after_create = list().unwrap();
+        let list_defined_after_create = list_defined().unwrap();
+        assert_eq!(list_defined_after_create, vec!["test-id"]);
+        assert_eq!(list_after_create, list_before_create);
+
+        let id = start_vm(DEFAULT_URI, "test-id".to_string()).unwrap();
+
+        let mut list_after_start = list().unwrap();
+        let list_defined_after_start = list_defined().unwrap();
+        let mut expected_list_after_start = Vec::new();
+        expected_list_after_start.push(id);
+        expected_list_after_start.extend(&list_before_create);
+
+        expected_list_after_start.sort();
+        list_after_start.sort();
+
+        assert!(list_defined_after_start.is_empty());
+        assert_eq!(list_after_start, expected_list_after_start);
+    }
+}