feat: pci passthrough (#114)

* feat: pci passthrough

* feat: guest device management

* feat: addons mounting and kernel modules support

* feat: more pci work

* fix: kernel build squashfs fixes

* fix: e820entry should be available on all platforms

crates/daemon/src/config.rs

@@ -0,0 +1,55 @@
+use std::{collections::HashMap, path::Path};
+
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use tokio::fs;
+
+#[derive(Serialize, Deserialize, Clone, Debug, Default)]
+pub struct DaemonConfig {
+    #[serde(default)]
+    pub pci: DaemonPciConfig,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, Default)]
+pub struct DaemonPciConfig {
+    #[serde(default)]
+    pub devices: HashMap<String, DaemonPciDeviceConfig>,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct DaemonPciDeviceConfig {
+    pub locations: Vec<String>,
+    #[serde(default)]
+    pub permissive: bool,
+    #[serde(default)]
+    #[serde(rename = "msi-translate")]
+    pub msi_translate: bool,
+    #[serde(default)]
+    #[serde(rename = "power-management")]
+    pub power_management: bool,
+    #[serde(default)]
+    #[serde(rename = "rdm-reserve-policy")]
+    pub rdm_reserve_policy: DaemonPciDeviceRdmReservePolicy,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, Default)]
+pub enum DaemonPciDeviceRdmReservePolicy {
+    #[default]
+    #[serde(rename = "strict")]
+    Strict,
+    #[serde(rename = "relaxed")]
+    Relaxed,
+}
+
+impl DaemonConfig {
+    pub async fn load(path: &Path) -> Result<DaemonConfig> {
+        if path.exists() {
+            let content = fs::read_to_string(path).await?;
+            let config: DaemonConfig = toml::from_str(&content)?;
+            Ok(config)
+        } else {
+            fs::write(&path, "").await?;
+            Ok(DaemonConfig::default())
+        }
+    }
+}

crates/daemon/src/control.rs

@@ -11,10 +11,11 @@ use krata::{
         control::{
             control_service_server::ControlService, ConsoleDataReply, ConsoleDataRequest,
             CreateGuestReply, CreateGuestRequest, DestroyGuestReply, DestroyGuestRequest,
-            ExecGuestReply, ExecGuestRequest, IdentifyHostReply, IdentifyHostRequest,
-            ListGuestsReply, ListGuestsRequest, PullImageReply, PullImageRequest,
-            ReadGuestMetricsReply, ReadGuestMetricsRequest, ResolveGuestReply, ResolveGuestRequest,
-            SnoopIdmReply, SnoopIdmRequest, WatchEventsReply, WatchEventsRequest,
+            DeviceInfo, ExecGuestReply, ExecGuestRequest, IdentifyHostReply, IdentifyHostRequest,
+            ListDevicesReply, ListDevicesRequest, ListGuestsReply, ListGuestsRequest,
+            PullImageReply, PullImageRequest, ReadGuestMetricsReply, ReadGuestMetricsRequest,
+            ResolveGuestReply, ResolveGuestRequest, SnoopIdmReply, SnoopIdmRequest,
+            WatchEventsReply, WatchEventsRequest,
         },
     },
 };
@@ -35,8 +36,8 @@ use uuid::Uuid;
 
 use crate::{
     command::DaemonCommand, console::DaemonConsoleHandle, db::GuestStore,
-    event::DaemonEventContext, glt::GuestLookupTable, idm::DaemonIdmHandle,
-    metrics::idm_metric_to_api, oci::convert_oci_progress,
+    devices::DaemonDeviceManager, event::DaemonEventContext, glt::GuestLookupTable,
+    idm::DaemonIdmHandle, metrics::idm_metric_to_api, oci::convert_oci_progress,
 };
 
 pub struct ApiError {
@@ -60,6 +61,7 @@ impl From<ApiError> for Status {
 #[derive(Clone)]
 pub struct DaemonControlService {
     glt: GuestLookupTable,
+    devices: DaemonDeviceManager,
     events: DaemonEventContext,
     console: DaemonConsoleHandle,
     idm: DaemonIdmHandle,
@@ -69,8 +71,10 @@ pub struct DaemonControlService {
 }
 
 impl DaemonControlService {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         glt: GuestLookupTable,
+        devices: DaemonDeviceManager,
         events: DaemonEventContext,
         console: DaemonConsoleHandle,
         idm: DaemonIdmHandle,
@@ -80,6 +84,7 @@ impl DaemonControlService {
     ) -> Self {
         Self {
             glt,
+            devices,
             events,
             console,
             idm,
@@ -524,4 +529,23 @@ impl ControlService for DaemonControlService {
         };
         Ok(Response::new(Box::pin(output) as Self::SnoopIdmStream))
     }
+
+    async fn list_devices(
+        &self,
+        request: Request<ListDevicesRequest>,
+    ) -> Result<Response<ListDevicesReply>, Status> {
+        let _ = request.into_inner();
+        let mut devices = Vec::new();
+        let state = self.devices.copy().await.map_err(|error| ApiError {
+            message: error.to_string(),
+        })?;
+        for (name, state) in state {
+            devices.push(DeviceInfo {
+                name,
+                claimed: state.owner.is_some(),
+                owner: state.owner.map(|x| x.to_string()).unwrap_or_default(),
+            });
+        }
+        Ok(Response::new(ListDevicesReply { devices }))
+    }
 }

crates/daemon/src/devices.rs

@@ -0,0 +1,106 @@
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::{anyhow, Result};
+use log::warn;
+use tokio::sync::RwLock;
+use uuid::Uuid;
+
+use crate::config::{DaemonConfig, DaemonPciDeviceConfig};
+
+#[derive(Clone)]
+pub struct DaemonDeviceState {
+    pub pci: Option<DaemonPciDeviceConfig>,
+    pub owner: Option<Uuid>,
+}
+
+#[derive(Clone)]
+pub struct DaemonDeviceManager {
+    config: Arc<DaemonConfig>,
+    devices: Arc<RwLock<HashMap<String, DaemonDeviceState>>>,
+}
+
+impl DaemonDeviceManager {
+    pub fn new(config: Arc<DaemonConfig>) -> Self {
+        Self {
+            config,
+            devices: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    pub async fn claim(&self, device: &str, uuid: Uuid) -> Result<DaemonDeviceState> {
+        let mut devices = self.devices.write().await;
+        let Some(state) = devices.get_mut(device) else {
+            return Err(anyhow!(
+                "unable to claim unknown device '{}' for guest {}",
+                device,
+                uuid
+            ));
+        };
+
+        if let Some(owner) = state.owner {
+            return Err(anyhow!(
+                "unable to claim device '{}' for guest {}: already claimed by {}",
+                device,
+                uuid,
+                owner
+            ));
+        }
+
+        state.owner = Some(uuid);
+        Ok(state.clone())
+    }
+
+    pub async fn release_all(&self, uuid: Uuid) -> Result<()> {
+        let mut devices = self.devices.write().await;
+        for state in (*devices).values_mut() {
+            if state.owner == Some(uuid) {
+                state.owner = None;
+            }
+        }
+        Ok(())
+    }
+
+    pub async fn release(&self, device: &str, uuid: Uuid) -> Result<()> {
+        let mut devices = self.devices.write().await;
+        let Some(state) = devices.get_mut(device) else {
+            return Ok(());
+        };
+
+        if let Some(owner) = state.owner {
+            if owner != uuid {
+                return Ok(());
+            }
+        }
+
+        state.owner = None;
+        Ok(())
+    }
+
+    pub async fn update_claims(&self, claims: HashMap<String, Uuid>) -> Result<()> {
+        let mut devices = self.devices.write().await;
+        devices.clear();
+        for (name, pci) in &self.config.pci.devices {
+            let owner = claims.get(name).cloned();
+            devices.insert(
+                name.clone(),
+                DaemonDeviceState {
+                    owner,
+                    pci: Some(pci.clone()),
+                },
+            );
+        }
+
+        for (name, uuid) in &claims {
+            if !devices.contains_key(name) {
+                warn!("unknown device '{}' assigned to guest {}", name, uuid);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn copy(&self) -> Result<HashMap<String, DaemonDeviceState>> {
+        let devices = self.devices.read().await;
+        Ok(devices.clone())
+    }
+}

crates/daemon/src/lib.rs

@@ -1,9 +1,11 @@
-use std::{net::SocketAddr, path::PathBuf, str::FromStr};
+use std::{net::SocketAddr, path::PathBuf, str::FromStr, sync::Arc};
 
 use anyhow::{anyhow, Result};
+use config::DaemonConfig;
 use console::{DaemonConsole, DaemonConsoleHandle};
 use control::DaemonControlService;
 use db::GuestStore;
+use devices::DaemonDeviceManager;
 use event::{DaemonEventContext, DaemonEventGenerator};
 use glt::GuestLookupTable;
 use idm::{DaemonIdm, DaemonIdmHandle};
@@ -23,9 +25,11 @@ use tonic::transport::{Identity, Server, ServerTlsConfig};
 use uuid::Uuid;
 
 pub mod command;
+pub mod config;
 pub mod console;
 pub mod control;
 pub mod db;
+pub mod devices;
 pub mod event;
 pub mod glt;
 pub mod idm;
@@ -35,7 +39,9 @@ pub mod reconcile;
 
 pub struct Daemon {
     store: String,
+    _config: Arc<DaemonConfig>,
     glt: GuestLookupTable,
+    devices: DaemonDeviceManager,
     guests: GuestStore,
     events: DaemonEventContext,
     guest_reconciler_task: JoinHandle<()>,
@@ -50,12 +56,20 @@ const GUEST_RECONCILER_QUEUE_LEN: usize = 1000;
 
 impl Daemon {
     pub async fn new(store: String) -> Result<Self> {
-        let mut image_cache_dir = PathBuf::from(store.clone());
+        let store_dir = PathBuf::from(store.clone());
+        let mut config_path = store_dir.clone();
+        config_path.push("config.toml");
+
+        let config = DaemonConfig::load(&config_path).await?;
+        let config = Arc::new(config);
+        let devices = DaemonDeviceManager::new(config.clone());
+
+        let mut image_cache_dir = store_dir.clone();
         image_cache_dir.push("cache");
         image_cache_dir.push("image");
         fs::create_dir_all(&image_cache_dir).await?;
 
-        let mut host_uuid_path = PathBuf::from(store.clone());
+        let mut host_uuid_path = store_dir.clone();
         host_uuid_path.push("host.uuid");
         let host_uuid = if host_uuid_path.is_file() {
             let content = fs::read_to_string(&host_uuid_path).await?;
@@ -74,8 +88,9 @@ impl Daemon {
             generated
         };
 
-        let initrd_path = detect_guest_file(&store, "initrd")?;
-        let kernel_path = detect_guest_file(&store, "kernel")?;
+        let initrd_path = detect_guest_path(&store, "initrd")?;
+        let kernel_path = detect_guest_path(&store, "kernel")?;
+        let addons_path = detect_guest_path(&store, "addons.squashfs")?;
 
         let packer = OciPackerService::new(None, &image_cache_dir, OciPlatform::current()).await?;
         let runtime = Runtime::new().await?;
@@ -93,6 +108,7 @@ impl Daemon {
                 .await?;
         let runtime_for_reconciler = runtime.dupe().await?;
         let guest_reconciler = GuestReconciler::new(
+            devices.clone(),
             glt.clone(),
             guests.clone(),
             events.clone(),
@@ -101,6 +117,7 @@ impl Daemon {
             guest_reconciler_notify.clone(),
             kernel_path,
             initrd_path,
+            addons_path,
         )?;
 
         let guest_reconciler_task = guest_reconciler.launch(guest_reconciler_receiver).await?;
@@ -108,7 +125,9 @@ impl Daemon {
 
         Ok(Self {
             store,
+            _config: config,
             glt,
+            devices,
             guests,
             events,
             guest_reconciler_task,
@@ -123,6 +142,7 @@ impl Daemon {
     pub async fn listen(&mut self, addr: ControlDialAddress) -> Result<()> {
         let control_service = DaemonControlService::new(
             self.glt.clone(),
+            self.devices.clone(),
             self.events.clone(),
             self.console.clone(),
             self.idm.clone(),
@@ -186,7 +206,7 @@ impl Drop for Daemon {
     }
 }
 
-fn detect_guest_file(store: &str, name: &str) -> Result<PathBuf> {
+fn detect_guest_path(store: &str, name: &str) -> Result<PathBuf> {
     let mut path = PathBuf::from(format!("{}/guest/{}", store, name));
     if path.is_file() {
         return Ok(path);

crates/daemon/src/reconcile/guest/mod.rs

@@ -26,6 +26,7 @@ use uuid::Uuid;
 
 use crate::{
     db::GuestStore,
+    devices::DaemonDeviceManager,
     event::{DaemonEvent, DaemonEventContext},
     glt::GuestLookupTable,
 };
@@ -55,6 +56,7 @@ impl Drop for GuestReconcilerEntry {
 
 #[derive(Clone)]
 pub struct GuestReconciler {
+    devices: DaemonDeviceManager,
     glt: GuestLookupTable,
     guests: GuestStore,
     events: DaemonEventContext,
@@ -62,6 +64,7 @@ pub struct GuestReconciler {
     packer: OciPackerService,
     kernel_path: PathBuf,
     initrd_path: PathBuf,
+    addons_path: PathBuf,
     tasks: Arc<Mutex<HashMap<Uuid, GuestReconcilerEntry>>>,
     guest_reconciler_notify: Sender<Uuid>,
     reconcile_lock: Arc<RwLock<()>>,
@@ -70,6 +73,7 @@ pub struct GuestReconciler {
 impl GuestReconciler {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
+        devices: DaemonDeviceManager,
         glt: GuestLookupTable,
         guests: GuestStore,
         events: DaemonEventContext,
@@ -78,8 +82,10 @@ impl GuestReconciler {
         guest_reconciler_notify: Sender<Uuid>,
         kernel_path: PathBuf,
         initrd_path: PathBuf,
+        modules_path: PathBuf,
     ) -> Result<Self> {
         Ok(Self {
+            devices,
             glt,
             guests,
             events,
@@ -87,6 +93,7 @@ impl GuestReconciler {
             packer,
             kernel_path,
             initrd_path,
+            addons_path: modules_path,
             tasks: Arc::new(Mutex::new(HashMap::new())),
             guest_reconciler_notify,
             reconcile_lock: Arc::new(RwLock::with_max_readers((), PARALLEL_LIMIT)),
@@ -152,6 +159,8 @@ impl GuestReconciler {
             self.guests.remove(guest.uuid).await?;
         }
 
+        let mut device_claims = HashMap::new();
+
         for (uuid, mut stored_guest) in stored_guests {
             let previous_guest = stored_guest.clone();
             let runtime_guest = runtime_guests.iter().find(|x| x.uuid == uuid);
@@ -173,6 +182,17 @@ impl GuestReconciler {
                     } else {
                         state.status = GuestStatus::Started.into();
                     }
+
+                    for device in &stored_guest
+                        .spec
+                        .as_ref()
+                        .cloned()
+                        .unwrap_or_default()
+                        .devices
+                    {
+                        device_claims.insert(device.name.clone(), uuid);
+                    }
+
                     state.network = Some(guestinfo_to_networkstate(runtime));
                     stored_guest.state = Some(state);
                 }
@@ -185,6 +205,9 @@ impl GuestReconciler {
                 let _ = self.guest_reconciler_notify.try_send(uuid);
             }
         }
+
+        self.devices.update_claims(device_claims).await?;
+
         Ok(())
     }
 
@@ -255,8 +278,10 @@ impl GuestReconciler {
 
     async fn start(&self, uuid: Uuid, guest: &mut Guest) -> Result<GuestReconcilerResult> {
         let starter = GuestStarter {
+            devices: &self.devices,
             kernel_path: &self.kernel_path,
             initrd_path: &self.initrd_path,
+            addons_path: &self.addons_path,
             packer: &self.packer,
             glt: &self.glt,
             runtime: &self.runtime,
@@ -293,6 +318,7 @@ impl GuestReconciler {
             host: self.glt.host_uuid().to_string(),
             domid: domid.unwrap_or(u32::MAX),
         });
+        self.devices.release_all(uuid).await?;
         Ok(GuestReconcilerResult::Changed { rerun: false })
     }
 

crates/daemon/src/reconcile/guest/start.rs

@@ -1,5 +1,7 @@
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
+use std::str::FromStr;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use anyhow::{anyhow, Result};
 use futures::StreamExt;
@@ -7,6 +9,7 @@ use krata::launchcfg::LaunchPackedFormat;
 use krata::v1::common::GuestOciImageSpec;
 use krata::v1::common::{guest_image_spec::Image, Guest, GuestState, GuestStatus, OciImageFormat};
 use krataoci::packer::{service::OciPackerService, OciPackedFormat};
+use kratart::launch::{PciBdf, PciDevice, PciRdmReservePolicy};
 use kratart::{launch::GuestLaunchRequest, Runtime};
 use log::info;
 
@@ -15,6 +18,8 @@ use tokio::io::AsyncReadExt;
 use tokio_tar::Archive;
 use uuid::Uuid;
 
+use crate::config::DaemonPciDeviceRdmReservePolicy;
+use crate::devices::DaemonDeviceManager;
 use crate::{
     glt::GuestLookupTable,
     reconcile::guest::{guestinfo_to_networkstate, GuestReconcilerResult},
@@ -24,8 +29,10 @@ use crate::{
 const OCI_SPEC_TAR_FILE_MAX_SIZE: usize = 100 * 1024 * 1024;
 
 pub struct GuestStarter<'a> {
+    pub devices: &'a DaemonDeviceManager,
     pub kernel_path: &'a Path,
     pub initrd_path: &'a Path,
+    pub addons_path: &'a Path,
     pub packer: &'a OciPackerService,
     pub glt: &'a GuestLookupTable,
     pub runtime: &'a Runtime,
@@ -135,6 +142,48 @@ impl GuestStarter<'_> {
             fs::read(&self.initrd_path).await?
         };
 
+        let success = AtomicBool::new(false);
+
+        let _device_release_guard = scopeguard::guard(
+            (spec.devices.clone(), self.devices.clone()),
+            |(devices, manager)| {
+                if !success.load(Ordering::Acquire) {
+                    tokio::task::spawn(async move {
+                        for device in devices {
+                            let _ = manager.release(&device.name, uuid).await;
+                        }
+                    });
+                }
+            },
+        );
+
+        let mut pcis = Vec::new();
+        for device in &spec.devices {
+            let state = self.devices.claim(&device.name, uuid).await?;
+            if let Some(cfg) = state.pci {
+                for location in cfg.locations {
+                    let pci = PciDevice {
+                        bdf: PciBdf::from_str(&location)?.with_domain(0),
+                        permissive: cfg.permissive,
+                        msi_translate: cfg.msi_translate,
+                        power_management: cfg.power_management,
+                        rdm_reserve_policy: match cfg.rdm_reserve_policy {
+                            DaemonPciDeviceRdmReservePolicy::Strict => PciRdmReservePolicy::Strict,
+                            DaemonPciDeviceRdmReservePolicy::Relaxed => {
+                                PciRdmReservePolicy::Relaxed
+                            }
+                        },
+                    };
+                    pcis.push(pci);
+                }
+            } else {
+                return Err(anyhow!(
+                    "device '{}' isn't a known device type",
+                    device.name
+                ));
+            }
+        }
+
         let info = self
             .runtime
             .launch(GuestLaunchRequest {
@@ -150,6 +199,7 @@ impl GuestStarter<'_> {
                 initrd,
                 vcpus: spec.vcpus,
                 mem: spec.mem,
+                pcis,
                 env: task
                     .environment
                     .iter()
@@ -157,6 +207,7 @@ impl GuestStarter<'_> {
                     .collect::<HashMap<_, _>>(),
                 run: empty_vec_optional(task.command.clone()),
                 debug: false,
+                addons_image: Some(self.addons_path.to_path_buf()),
             })
             .await?;
         self.glt.associate(uuid, info.domid).await;
@@ -169,6 +220,7 @@ impl GuestStarter<'_> {
             host: self.glt.host_uuid().to_string(),
             domid: info.domid,
         });
+        success.store(true, Ordering::Release);
         Ok(GuestReconcilerResult::Changed { rerun: false })
     }
 }