Update worker thread for TEE

jakecorrenti · jakecorrenti · commit ab523d211964 · 2025-05-01T12:09:29.000-04:00
For TEE VMs, use the same sender as we would for macOS or an x86 VM with
a split IRQCHIP. ADditionally, use a channel for inter-process
communication instead of an EventFd.

Signed-off-by: Jake Correnti &lt;jakecorrenti+github@proton.me&gt;
diff --git a/src/libkrun/src/lib.rs b/src/libkrun/src/lib.rs
@@ -5,8 +5,6 @@ use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::convert::TryInto;
 use std::env;
-#[cfg(feature = "tee")]
-use std::ffi::c_void;
 use std::ffi::CStr;
 #[cfg(target_os = "linux")]
 use std::ffi::CString;
@@ -53,17 +51,6 @@ use vmm::vmm_config::machine_config::VmConfig;
 use vmm::vmm_config::net::NetworkInterfaceConfig;
 use vmm::vmm_config::vsock::VsockDeviceConfig;
 
-#[cfg(feature = "tee")]
-use kvm_bindings::{kvm_memory_attributes, KVM_MEMORY_ATTRIBUTE_PRIVATE};
-
-#[cfg(feature = "tee")]
-use vm_memory::{guest_memory::GuestMemory, GuestAddress, GuestMemoryRegion, MemoryRegionAddress};
-
-#[cfg(feature = "tee")]
-use libc::{
-    fallocate, madvise, EFD_SEMAPHORE, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE, MADV_DONTNEED,
-};
-
 // Value returned on success. We use libc's errors otherwise.
 const KRUN_SUCCESS: i32 = 0;
 // Maximum number of arguments/environment variables we allow
@@ -1504,24 +1491,12 @@ pub extern "C" fn krun_start_enter(ctx_id: u32) -> i32 {
         }
     }
 
-    #[cfg(feature = "tee")]
-    let (pm_sender, pm_receiver) = unbounded();
-    #[cfg(feature = "tee")]
-    let pm_efd =
-        EventFd::new(EFD_SEMAPHORE).expect("unable to create TEE memory properties eventfd");
-    let (sender, receiver) = unbounded();
+    let (sender, _receiver) = unbounded();
 
     let _vmm = match vmm::builder::build_microvm(
         &ctx_cfg.vmr,
         &mut event_manager,
         ctx_cfg.shutdown_efd,
-        #[cfg(feature = "tee")]
-        (
-            pm_sender,
-            pm_efd
-                .try_clone()
-                .expect("unable to clone TEE memory properties eventfd"),
-        ),
         sender,
     ) {
         Ok(vmm) => vmm,
@@ -1531,105 +1506,18 @@ pub extern "C" fn krun_start_enter(ctx_id: u32) -> i32 {
         }
     };
 
-    #[cfg(feature = "tee")]
-    let mapper_vmm = _vmm.clone();
-
     #[cfg(target_os = "macos")]
     if ctx_cfg.gpu_virgl_flags.is_some() {
-        vmm::worker::start_worker_thread(_vmm.clone(), receiver).unwrap();
+        vmm::worker::start_worker_thread(_vmm.clone(), _receiver).unwrap();
     }
 
     #[cfg(target_arch = "x86_64")]
     if ctx_cfg.vmr.split_irqchip {
-        vmm::worker::start_worker_thread(_vmm.clone(), receiver.clone()).unwrap();
+        vmm::worker::start_worker_thread(_vmm.clone(), _receiver.clone()).unwrap();
     }
 
-    #[cfg(feature = "tee")]
-    let guest_mem = _vmm.lock().unwrap().guest_memory().clone();
-
-    #[cfg(feature = "tee")]
-    std::thread::Builder::new()
-        .name("TEE memory properties worker".into())
-        .spawn(move || loop {
-            match pm_receiver.recv() {
-                Err(e) => error!("Error in pm receiver: {:?}", e),
-                Ok(m) => {
-                    let (guest_memfd, region_start) = mapper_vmm
-                        .lock()
-                        .unwrap()
-                        .kvm_vm()
-                        .guest_memfd_get(m.gpa)
-                        .unwrap_or_else(|| panic!("unable to find KVM guest_memfd for memory region corresponding to GPA 0x{:x}", m.gpa));
-
-                    let attributes: u64 = if m.private {
-                        KVM_MEMORY_ATTRIBUTE_PRIVATE as u64
-                    } else {
-                        0
-                    };
-
-                    let attr = kvm_memory_attributes {
-                        address: m.gpa,
-                        size: m.size,
-                        attributes,
-                        flags: 0,
-                    };
-
-                    mapper_vmm
-                        .lock()
-                        .unwrap()
-                        .kvm_vm()
-                        .fd()
-                        .set_memory_attributes(attr)
-                        .unwrap_or_else(|_| panic!("unable to set memory attributes for memory region corresponding to guest address 0x{:x}", m.gpa));
-
-                    let region = guest_mem.find_region(GuestAddress(m.gpa));
-                    if region.is_none() {
-                        error!("guest memory region corresponding to GPA 0x{:x} not found", m.gpa);
-                        pm_efd.write(1).unwrap();
-                        continue;
-                    }
-
-                    let offset = m.gpa - region_start;
-
-                    if m.private {
-                        let region_addr = MemoryRegionAddress(offset);
-
-                        let host_startaddr = region
-                            .unwrap()
-                            .get_host_address(region_addr)
-                            .expect("host address corresponding to memory region address 0x{:x} not found");
-
-                        let ret = unsafe {
-                            madvise(
-                                host_startaddr as *mut c_void,
-                                m.size.try_into().unwrap(),
-                                MADV_DONTNEED,
-                            )
-                        };
-
-                        if ret < 0 {
-                            error!("unable to advise kernel that memory region corresponding to GPA 0x{:x} will likely not be needed (madvise)", m.gpa);
-                        }
-                    } else {
-                        let ret = unsafe {
-                            fallocate(
-                                guest_memfd,
-                                FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE,
-                                offset as i64,
-                                m.size as i64,
-                            )
-                        };
-
-                        if ret < 0 {
-                            error!("unable to allocate space in guest_memfd for shared memory (fallocate)");
-                        }
-                    }
-
-                    pm_efd.write(1).unwrap();
-                }
-            }
-        })
-        .unwrap();
+    #[cfg(feature = "amd-sev")]
+    vmm::worker::start_worker_thread(_vmm.clone(), _receiver.clone()).unwrap();
 
     loop {
         match event_manager.run() {
diff --git a/src/utils/src/worker_message.rs b/src/utils/src/worker_message.rs
@@ -1,3 +1,10 @@
+#[derive(Debug)]
+pub struct MemoryProperties {
+    pub gpa: u64,
+    pub size: u64,
+    pub private: bool,
+}
+
 #[derive(Debug)]
 pub enum WorkerMessage {
     #[cfg(target_arch = "x86_64")]
@@ -11,4 +18,5 @@ pub enum WorkerMessage {
     GpuAddMapping(crossbeam_channel::Sender<bool>, u64, u64, u64),
     #[cfg(target_os = "macos")]
     GpuRemoveMapping(crossbeam_channel::Sender<bool>, u64, u64),
+    ConvertMemory(crossbeam_channel::Sender<bool>, MemoryProperties),
 }
diff --git a/src/vmm/src/builder.rs b/src/vmm/src/builder.rs
@@ -23,8 +23,6 @@ use super::{Error, Vmm};
 #[cfg(target_arch = "x86_64")]
 use crate::device_manager::legacy::PortIODeviceManager;
 use crate::device_manager::mmio::MMIODeviceManager;
-#[cfg(feature = "tee")]
-use crate::linux::vstate::MemoryProperties;
 use crate::resources::VmResources;
 use crate::vmm_config::external_kernel::{ExternalKernel, KernelFormat};
 #[cfg(all(target_os = "linux", target_arch = "aarch64"))]
@@ -511,7 +509,6 @@ pub fn build_microvm(
     vm_resources: &super::resources::VmResources,
     event_manager: &mut EventManager,
     _shutdown_efd: Option<EventFd>,
-    #[cfg(feature = "tee")] pm_sender: (Sender<MemoryProperties>, EventFd),
     _sender: Sender<WorkerMessage>,
 ) -> std::result::Result<Arc<Mutex<Vmm>>, StartMicrovmError> {
     let payload = choose_payload(vm_resources)?;
@@ -667,7 +664,10 @@ pub fn build_microvm(
     #[cfg(target_arch = "x86_64")]
     {
         let ioapic: Box<dyn IrqChipT> = if vm_resources.split_irqchip {
-            Box::new(IoApic::new(vm.fd(), _sender).map_err(StartMicrovmError::CreateKvmIrqChip)?)
+            Box::new(
+                IoApic::new(vm.fd(), _sender.clone())
+                    .map_err(StartMicrovmError::CreateKvmIrqChip)?,
+            )
         } else {
             Box::new(KvmIoapic::new(vm.fd()).map_err(StartMicrovmError::CreateKvmIrqChip)?)
         };
@@ -689,7 +689,7 @@ pub fn build_microvm(
             &pio_device_manager.io_bus,
             &exit_evt,
             #[cfg(feature = "tee")]
-            pm_sender,
+            _sender,
         )
         .map_err(StartMicrovmError::Internal)?;
     }
@@ -1456,7 +1456,7 @@ fn create_vcpus_x86_64(
     entry_addr: GuestAddress,
     io_bus: &devices::Bus,
     exit_evt: &EventFd,
-    #[cfg(feature = "tee")] pm_sender: (Sender<MemoryProperties>, EventFd),
+    #[cfg(feature = "tee")] pm_sender: Sender<WorkerMessage>,
 ) -> super::Result<Vec<Vcpu>> {
     let mut vcpus = Vec::with_capacity(vcpu_config.vcpu_count as usize);
     for cpu_index in 0..vcpu_config.vcpu_count {
@@ -1468,7 +1468,7 @@ fn create_vcpus_x86_64(
             io_bus.clone(),
             exit_evt.try_clone().map_err(Error::EventFd)?,
             #[cfg(feature = "tee")]
-            (pm_sender.0.clone(), pm_sender.1.try_clone().unwrap()),
+            pm_sender.clone(),
         )
         .map_err(Error::Vcpu)?;
 
diff --git a/src/vmm/src/linux/vstate.rs b/src/vmm/src/linux/vstate.rs
@@ -50,6 +50,8 @@ use kvm_ioctls::{Cap::*, *};
 use utils::eventfd::EventFd;
 use utils::signal::{register_signal_handler, sigrtmin, Killable};
 use utils::sm::StateMachine;
+#[cfg(feature = "tee")]
+use utils::worker_message::{MemoryProperties, WorkerMessage};
 use vm_memory::{
     Address, GuestAddress, GuestMemory, GuestMemoryError, GuestMemoryMmap, GuestMemoryRegion,
     GuestRegionMmap,
@@ -746,13 +748,6 @@ pub struct VmState {
     ioapic: kvm_irqchip,
 }
 
-#[cfg(feature = "tee")]
-pub struct MemoryProperties {
-    pub gpa: u64,
-    pub size: u64,
-    pub private: bool,
-}
-
 /// Encapsulates configuration parameters for the guest vCPUS.
 #[derive(Debug, Eq, PartialEq)]
 pub struct VcpuConfig {
@@ -796,7 +791,7 @@ pub struct Vcpu {
     response_sender: Sender<VcpuResponse>,
 
     #[cfg(feature = "tee")]
-    pm_sender: (Sender<MemoryProperties>, EventFd),
+    pm_sender: Sender<WorkerMessage>,
 }
 
 impl Vcpu {
@@ -900,7 +895,7 @@ impl Vcpu {
         msr_list: MsrList,
         io_bus: devices::Bus,
         exit_evt: EventFd,
-        #[cfg(feature = "tee")] pm_sender: (Sender<MemoryProperties>, EventFd),
+        #[cfg(feature = "tee")] pm_sender: Sender<WorkerMessage>,
     ) -> Result<Self> {
         let kvm_vcpu = vm_fd.create_vcpu(id as u64).map_err(Error::VcpuFd)?;
         let (event_sender, event_receiver) = unbounded();
@@ -1222,8 +1217,17 @@ impl Vcpu {
 
                     let mem_properties = MemoryProperties { gpa, size, private };
 
-                    self.pm_sender.0.send(mem_properties).unwrap();
-                    let _ = self.pm_sender.1.read().unwrap();
+                    let (response_sender, response_receiver) = unbounded();
+                    self.pm_sender
+                        .send(WorkerMessage::ConvertMemory(
+                            response_sender.clone(),
+                            mem_properties,
+                        ))
+                        .unwrap();
+                    if !response_receiver.recv().unwrap() {
+                        error!("Unable to convert memory with properties: gpa: 0x{:x} size: 0x{:x} to_private: {}", gpa, size, private);
+                        return Err(Error::VcpuUnhandledKvmExit);
+                    }
                     Ok(VcpuEmulation::Handled)
                 }
                 #[cfg(target_arch = "x86_64")]
@@ -1242,8 +1246,17 @@ impl Vcpu {
 
                     let mem_properties = MemoryProperties { gpa, size, private };
 
-                    self.pm_sender.0.send(mem_properties).unwrap();
-                    let _ = self.pm_sender.1.read().unwrap();
+                    let (response_sender, response_receiver) = unbounded();
+                    self.pm_sender
+                        .send(WorkerMessage::ConvertMemory(
+                            response_sender.clone(),
+                            mem_properties,
+                        ))
+                        .unwrap();
+                    if !response_receiver.recv().unwrap() {
+                        error!("Unable to convert memory with properties: gpa: 0x{:x} size: 0x{:x} to_private: {}", gpa, size, private);
+                        return Err(Error::VcpuUnhandledKvmExit);
+                    }
 
                     Ok(VcpuEmulation::Handled)
                 }
diff --git a/src/vmm/src/worker.rs b/src/vmm/src/worker.rs
