From 339fb32e6ee6d8afad04bb8089137006f418b073 Mon Sep 17 00:00:00 2001
From: Changyuan Lyu <changyuanl@google.com>
Date: Sun, 27 Oct 2024 17:24:17 -0700
Subject: [PATCH] refactor(vfio): let VfioPciDev::new take a configured device

Signed-off-by: Changyuan Lyu <changyuanl@google.com>
---
 alioth-cli/src/main.rs |  2 +-
 alioth/src/vfio/pci.rs | 17 +++--------------
 alioth/src/vm.rs       | 16 ++++++++--------
 3 files changed, 12 insertions(+), 23 deletions(-)

diff --git a/alioth-cli/src/main.rs b/alioth-cli/src/main.rs
index ce5130b..2b4e8b0 100644
--- a/alioth-cli/src/main.rs
+++ b/alioth-cli/src/main.rs
@@ -415,7 +415,7 @@ fn main_run(args: RunArgs) -> Result<(), Error> {
     for (index, vfio) in args.vfio.into_iter().enumerate() {
         let param: VfioParam =
             serde_aco::from_args(&vfio, &objects).context(error::ParseArg { arg: vfio })?;
-        vm.add_vfio_dev(format!("vfio-{index}"), param)
+        vm.add_vfio_dev(format!("vfio-{index}").into(), param)
             .context(error::CreateDevice)?;
     }
 
diff --git a/alioth/src/vfio/pci.rs b/alioth/src/vfio/pci.rs
index 7c6ac68..80a5943 100644
--- a/alioth/src/vfio/pci.rs
+++ b/alioth/src/vfio/pci.rs
@@ -44,8 +44,7 @@ use crate::vfio::bindings::{
 };
 use crate::vfio::cdev::Cdev;
 use crate::vfio::device::Device;
-use crate::vfio::iommu::Ioas;
-use crate::vfio::{error, Result, VfioParam};
+use crate::vfio::{error, Result};
 use crate::{align_down, align_up, assign_bits, mask_bits, mem};
 
 fn round_up_range(range: Range<usize>) -> Range<usize> {
@@ -397,18 +396,8 @@ impl<M> VfioPciDev<M>
 where
     M: MsiSender,
 {
-    pub fn new(
-        name: impl Into<Arc<str>>,
-        param: &VfioParam,
-        ioas: Arc<Ioas>,
-        msi_sender: M,
-    ) -> Result<VfioPciDev<M>> {
-        let mut dev = Cdev::new(&param.cdev)?;
-        dev.attach_iommu_ioas(ioas)?;
-        let cdev = Arc::new(VfioCdev {
-            dev,
-            name: name.into(),
-        });
+    pub fn new(name: Arc<str>, dev: Cdev, msi_sender: M) -> Result<VfioPciDev<M>> {
+        let cdev = Arc::new(VfioCdev { dev, name });
 
         let region_config = cdev.dev.get_region_info(VfioPciRegion::CONFIG.raw())?;
 
diff --git a/alioth/src/vm.rs b/alioth/src/vm.rs
index 44adca4..9ef0ccf 100644
--- a/alioth/src/vm.rs
+++ b/alioth/src/vm.rs
@@ -45,6 +45,8 @@ use crate::mem::{MemRegion, MemRegionType};
 use crate::pci::bus::PciBus;
 use crate::pci::{Bdf, PciDevice};
 #[cfg(target_os = "linux")]
+use crate::vfio::cdev::Cdev;
+#[cfg(target_os = "linux")]
 use crate::vfio::iommu::UpdateIommuIoas;
 #[cfg(target_os = "linux")]
 use crate::vfio::iommu::{Ioas, Iommu};
@@ -303,12 +305,7 @@ where
 
 #[cfg(target_os = "linux")]
 impl Machine<Kvm> {
-    pub fn add_vfio_dev(
-        &mut self,
-        name: impl Into<Arc<str>>,
-        param: VfioParam,
-    ) -> Result<(), Error> {
-        let name = name.into();
+    pub fn add_vfio_dev(&mut self, name: Arc<str>, param: VfioParam) -> Result<(), Error> {
         let iommu = if let Some(iommu) = &self.iommu {
             iommu.clone()
         } else {
@@ -340,8 +337,11 @@ impl Machine<Kvm> {
             u32::from(bdf.0),
         )?;
 
-        let dev = Arc::new(VfioPciDev::new(name.clone(), &param, ioas, msi_sender)?);
-        let pci_dev = PciDevice::new(name, dev);
+        let mut cdev = Cdev::new(&param.cdev)?;
+        cdev.attach_iommu_ioas(ioas)?;
+
+        let dev = VfioPciDev::new(name.clone(), cdev, msi_sender)?;
+        let pci_dev = PciDevice::new(name, Arc::new(dev));
         self.add_pci_dev(Some(bdf), pci_dev)?;
         Ok(())
     }