refine: Refine code

crates/vm-core/src/arch/aarch64/vm_exit.rs

@@ -1,23 +1,10 @@
 use tracing::trace;
 
-use crate::arch::aarch64::firmware::psci::error::PsciError;
 use crate::arch::aarch64::vcpu::reg::CoreRegister;
 use crate::arch::aarch64::vcpu::reg::SysRegister;
-use crate::device_manager::vm_exit::DeviceError;
 use crate::vcpu::error::VcpuError;
 use crate::vcpu::vcpu::Vcpu;
-
-#[derive(Debug, thiserror::Error)]
-pub enum Error {
-    #[error("Failed to handle mmio, err: {0}")]
-    DeviceError(#[from] DeviceError),
-
-    #[error("Failed to rw vcpu, err: {0}")]
-    VcpuError(String),
-
-    #[error("Failed to handle smc, err: {0}")]
-    PsciError(#[from] PsciError),
-}
+use crate::vcpu::vm_exit::VmExit;
 
 #[derive(Debug)]
 pub enum VmExitReason {
@@ -52,24 +39,23 @@ pub enum HandleVmExitResult {
 pub fn handle_vm_exit(
     vcpu: &mut Vcpu,
     exit_reason: VmExitReason,
+    vm_exit_handler: &dyn VmExit,
 ) -> Result<HandleVmExitResult, VcpuError> {
-    let device_vm_exit_handler = vcpu.device_vm_exit_handler.clone();
-
     trace!(?exit_reason);
 
     match exit_reason {
         VmExitReason::Unknown => Ok(HandleVmExitResult::Continue),
         VmExitReason::Wf => Ok(HandleVmExitResult::NextInstruction),
-        VmExitReason::MMRead { gpa, srt, len } if device_vm_exit_handler.in_mmio_region(gpa) => {
+        VmExitReason::MMRead { gpa, srt, len } if vm_exit_handler.in_mmio_region(gpa) => {
             let mut buf = [0; 8];
-            device_vm_exit_handler.mmio_read(gpa, len, &mut buf[0..len])?;
+            vm_exit_handler.mmio_read(gpa, len, &mut buf[0..len])?;
             vcpu.vcpu_instance
                 .set_core_reg(srt, u64::from_le_bytes(buf))?;
             Ok(HandleVmExitResult::NextInstruction)
         }
         VmExitReason::MMRead { .. } => todo!(),
-        VmExitReason::MMWrite { gpa, buf, len } if device_vm_exit_handler.in_mmio_region(gpa) => {
-            device_vm_exit_handler.mmio_write(gpa, len, &buf)?;
+        VmExitReason::MMWrite { gpa, buf, len } if vm_exit_handler.in_mmio_region(gpa) => {
+            vm_exit_handler.mmio_write(gpa, len, &buf)?;
             Ok(HandleVmExitResult::NextInstruction)
         }
         VmExitReason::MMWrite { .. } => todo!(),
@@ -83,10 +69,8 @@ pub fn handle_vm_exit(
             Ok(HandleVmExitResult::NextInstruction)
         }
         VmExitReason::Smc => {
-            let psci = vcpu.psci.clone();
-
             // We only support psci for smc now
-            psci.call(vcpu.vcpu_instance.as_mut())?;
+            vm_exit_handler.call_smc(vcpu.vcpu_instance.as_mut())?;
 
             Ok(HandleVmExitResult::NextInstruction)
         }

crates/vm-core/src/arch/x86_64.rs

@@ -1,5 +1,3 @@
 pub mod layout;
 pub mod vcpu;
 pub mod vm_exit;
-
-pub const BASE_ADDRESS: u64 = 0x0;

crates/vm-core/src/device/pio/pio_as_manager.rs

@@ -30,9 +30,9 @@ impl PioAddressSpaceManager {
         Ok(())
     }
 
-    pub fn get_device_by_port(&mut self, port: u16) -> Option<&mut dyn PioDevice> {
+    pub fn get_device_by_port(&self, port: u16) -> Option<&dyn PioDevice> {
         let (_, &idx) = self.address_space.try_get_value_by_key(port)?;
 
-        Some(self.device[idx].as_mut())
+        Some(self.device[idx].as_ref())
     }
 }

crates/vm-core/src/device/pio/pio_device.rs

@@ -6,7 +6,7 @@ pub type PortRange = Range<u16>;
 pub trait PioDevice: Device {
     fn ports(&self) -> Vec<PortRange>;
 
-    fn io_in(&mut self, port: u16, data: &mut [u8]);
+    fn io_in(&self, port: u16, data: &mut [u8]);
 
-    fn io_out(&mut self, port: u16, data: &[u8]);
+    fn io_out(&self, port: u16, data: &[u8]);
 }

crates/vm-core/src/device_manager.rs

@@ -1,2 +1 @@
 pub mod manager;
-pub mod vm_exit;

crates/vm-core/src/device_manager/manager.rs

@@ -5,87 +5,11 @@ use crate::device::mmio::mmio_as_manager::MmioAddressSpaceManager;
 use crate::device::mmio::mmio_device::MmioDevice;
 use crate::device::pio::pio_as_manager::PioAddressSpaceManager;
 use crate::device::pio::pio_device::PioDevice;
-use crate::device_manager::vm_exit::DeviceError;
-use crate::device_manager::vm_exit::DeviceVmExitHandler;
 use crate::utils::address_space::AddressSpaceError;
 
 pub struct DeviceManager {
-    pio_manager: PioAddressSpaceManager,
-    mmio_manager: MmioAddressSpaceManager,
-}
-
-impl DeviceVmExitHandler for DeviceManager {
-    fn io_in(&mut self, port: u16, data: &mut [u8]) -> Result<(), DeviceError> {
-        let device = self
-            .pio_manager
-            .get_device_by_port(port)
-            .ok_or(DeviceError::NoDeviceForPort(port))?;
-
-        device.io_in(port, data);
-
-        Ok(())
-    }
-
-    fn io_out(&mut self, port: u16, data: &[u8]) -> Result<(), DeviceError> {
-        let device = self
-            .pio_manager
-            .get_device_by_port(port)
-            .ok_or(DeviceError::NoDeviceForPort(port))?;
-
-        device.io_out(port, data);
-
-        Ok(())
-    }
-
-    fn mmio_read(&self, addr: u64, len: usize, data: &mut [u8]) -> Result<(), DeviceError> {
-        let (range, handler) = self
-            .mmio_manager
-            .get_handler_by_addr(addr)
-            .ok_or(DeviceError::NoDeviceForAddr(addr))?;
-
-        let err = || DeviceError::MmioOutOfMemory {
-            mmio_start: range.start,
-            mmio_len: range.len,
-            addr,
-        };
-
-        if addr.checked_add(len as u64).ok_or_else(err)?
-            > range.start.checked_add(range.len as u64).ok_or_else(err)?
-        {
-            return Err(err());
-        }
-
-        handler.mmio_read(addr - range.start, len, data);
-
-        Ok(())
-    }
-
-    fn mmio_write(&self, addr: u64, len: usize, data: &[u8]) -> Result<(), DeviceError> {
-        let (range, handler) = self
-            .mmio_manager
-            .get_handler_by_addr(addr)
-            .ok_or(DeviceError::NoDeviceForAddr(addr))?;
-
-        let err = || DeviceError::MmioOutOfMemory {
-            mmio_start: range.start,
-            mmio_len: range.len,
-            addr,
-        };
-
-        if addr.checked_add(len as u64).ok_or_else(err)?
-            > range.start.checked_add(range.len as u64).ok_or_else(err)?
-        {
-            return Err(err());
-        }
-
-        handler.mmio_write(addr - range.start, len, data);
-
-        Ok(())
-    }
-
-    fn in_mmio_region(&self, addr: u64) -> bool {
-        self.mmio_manager.mmio_layout().in_mmio_region(addr)
-    }
+    pub pio_manager: PioAddressSpaceManager,
+    pub mmio_manager: MmioAddressSpaceManager,
 }
 
 impl DeviceManager {

crates/vm-core/src/vcpu.rs

@@ -1,3 +1,4 @@
 pub mod error;
 pub mod vcpu;
 pub mod vcpu_manager;
+pub mod vm_exit;

crates/vm-core/src/vcpu/error.rs

@@ -1,6 +1,4 @@
-#[cfg(target_arch = "aarch64")]
-use crate::arch::aarch64::firmware::psci::error::PsciError;
-use crate::device_manager::vm_exit::DeviceError;
+use crate::vcpu::vm_exit::VmExitHandlerError;
 
 #[derive(Debug, thiserror::Error)]
 pub enum VcpuError {
@@ -15,13 +13,9 @@ pub enum VcpuError {
     #[error("{0}")]
     KvmError(#[from] kvm_ioctls::Error),
 
-    #[cfg(target_arch = "aarch64")]
     #[error("{0}")]
-    PsciError(#[from] PsciError),
-
-    #[error("{0}")]
-    DeviceError(#[from] DeviceError),
+    GuestError(String),
 
     #[error("{0}")]
-    GuestError(String),
+    VmExitHandlerErr(#[from] VmExitHandlerError),
 }

crates/vm-core/src/vcpu/vcpu.rs

@@ -1,12 +1,8 @@
 use std::sync::Arc;
 
-#[cfg(target_arch = "aarch64")]
-use crate::arch::aarch64::firmware::psci::Psci;
-use crate::device_manager::vm_exit::DeviceVmExitHandler;
+use crate::vcpu::vm_exit::VmExit;
 
 pub struct Vcpu {
     pub vcpu_instance: Box<dyn crate::virt::vcpu::Vcpu>,
-    pub device_vm_exit_handler: Arc<dyn DeviceVmExitHandler>,
-    #[cfg(target_arch = "aarch64")]
-    pub psci: Arc<dyn Psci>,
+    pub vm_exit_handler: Arc<dyn VmExit>,
 }

crates/vm-core/src/vcpu/vcpu_manager.rs

@@ -1,58 +1,56 @@
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::thread;
+use std::thread::JoinHandle;
 
-#[cfg(target_arch = "aarch64")]
-use crate::arch::aarch64::firmware::psci::Psci;
 #[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::vcpu::setup_cpu;
-use crate::device_manager::vm_exit::DeviceVmExitHandler;
 use crate::error::Error as VmError;
 use crate::vcpu::error::VcpuError;
 use crate::vcpu::vcpu::Vcpu;
-use crate::virt::Vm;
+use crate::vcpu::vm_exit::VmExit;
+use crate::virt::vm::Vm;
 
 pub struct VcpuManager {
     vm_instance: Arc<dyn Vm>,
     vcpus: Vec<Arc<Mutex<Vcpu>>>,
+    handlers: Vec<JoinHandle<Result<(), VcpuError>>>,
 }
 
 impl VcpuManager {
     pub fn new(vm_instance: Arc<dyn Vm>) -> Self {
         VcpuManager {
             vm_instance,
             vcpus: Default::default(),
+            handlers: Default::default(),
         }
     }
 
     pub fn create_vcpu(
         &mut self,
         vcpu_id: usize,
-        device_vm_exit_handler: Arc<dyn DeviceVmExitHandler>,
-        #[cfg(target_arch = "aarch64")] psci: Arc<dyn Psci>,
+        vm_exit_handler: Arc<dyn VmExit>,
     ) -> Result<(), VmError> {
         let vcpu_instance = self.vm_instance.create_vcpu(vcpu_id)?;
 
         let vcpu = Vcpu {
             vcpu_instance,
-            device_vm_exit_handler,
-            #[cfg(target_arch = "aarch64")]
-            psci,
+            vm_exit_handler,
         };
 
         self.vcpus.push(Arc::new(Mutex::new(vcpu)));
 
         Ok(())
     }
 
-    pub fn start_vcpu(&self, vcpu_id: usize, start_pc: u64, x0: u64) -> Result<(), VcpuError> {
+    pub fn start_vcpu(&mut self, vcpu_id: usize, start_pc: u64, x0: u64) -> Result<(), VcpuError> {
         let vcpu = self
             .vcpus
             .get(vcpu_id)
             .ok_or(VcpuError::VcpuNotCreated(vcpu_id))?
             .clone();
 
-        thread::spawn(move || -> Result<(), VcpuError> {
+        let handle = thread::spawn(move || -> Result<(), VcpuError> {
             let mut vcpu = vcpu.lock().unwrap();
 
             vcpu.vcpu_instance.post_init_within_thread()?;
@@ -62,6 +60,8 @@ impl VcpuManager {
                 setup_cpu(x0, start_pc, vcpu_id, &mut *vcpu.vcpu_instance)?;
             }
 
+            let vm_exit_handler = vcpu.vm_exit_handler.clone();
+
             loop {
                 let vm_exit_reason = vcpu.vcpu_instance.run()?;
 
@@ -70,8 +70,11 @@ impl VcpuManager {
                     use crate::arch::aarch64::vcpu::reg::CoreRegister;
                     use crate::arch::aarch64::vm_exit::HandleVmExitResult;
 
-                    match crate::arch::aarch64::vm_exit::handle_vm_exit(&mut *vcpu, vm_exit_reason)?
-                    {
+                    match crate::arch::aarch64::vm_exit::handle_vm_exit(
+                        &mut *vcpu,
+                        vm_exit_reason,
+                        vm_exit_handler.as_ref(),
+                    )? {
                         HandleVmExitResult::Continue => (),
                         HandleVmExitResult::NextInstruction => {
                             let pc = vcpu.vcpu_instance.get_core_reg(CoreRegister::PC)?;
@@ -82,6 +85,8 @@ impl VcpuManager {
             }
         });
 
+        self.handlers.push(handle);
+
         Ok(())
     }
 }

crates/vm-core/src/device_manager/vm_exit.rs → crates/vm-core/src/vcpu/vm_exit.rs

@@ -1,5 +1,7 @@
+use crate::virt::vcpu::Vcpu;
+
 #[derive(Debug, thiserror::Error)]
-pub enum DeviceError {
+pub enum VmExitHandlerError {
     #[error("no device found for port 0x{0:#x}")]
     NoDeviceForPort(u16),
 
@@ -14,12 +16,23 @@ pub enum DeviceError {
         mmio_len: usize,
         addr: u64,
     },
+
+    #[cfg(target_arch = "aarch64")]
+    #[error("{0}")]
+    SmcError(#[from] crate::arch::aarch64::firmware::psci::error::PsciError),
 }
 
-pub trait DeviceVmExitHandler: Send + Sync {
-    fn io_in(&mut self, port: u16, data: &mut [u8]) -> Result<(), DeviceError>;
-    fn io_out(&mut self, port: u16, data: &[u8]) -> Result<(), DeviceError>;
-    fn mmio_read(&self, addr: u64, len: usize, data: &mut [u8]) -> Result<(), DeviceError>;
-    fn mmio_write(&self, addr: u64, len: usize, data: &[u8]) -> Result<(), DeviceError>;
+pub trait VmExit: Send + Sync {
+    fn io_in(&mut self, port: u16, data: &mut [u8]) -> Result<(), VmExitHandlerError>;
+
+    fn io_out(&mut self, port: u16, data: &[u8]) -> Result<(), VmExitHandlerError>;
+
+    fn mmio_read(&self, addr: u64, len: usize, data: &mut [u8]) -> Result<(), VmExitHandlerError>;
+
+    fn mmio_write(&self, addr: u64, len: usize, data: &[u8]) -> Result<(), VmExitHandlerError>;
+
     fn in_mmio_region(&self, addr: u64) -> bool;
+
+    #[cfg(target_arch = "aarch64")]
+    fn call_smc(&self, vcpu: &mut dyn Vcpu) -> Result<(), VmExitHandlerError>;
 }