Add support for KVM_EXIT_MEMORY_FAULT

The KVM_EXIT_MEMORY_FAULT vmexit is triggered when guest wants to switch
a region of memory from private to shared and viceversa. To support this
when tee is enabled, add an extra thread named sender_io that gets the
parameters from the vcpu thread and triggers the
set_memory_properties(). The vcpu fd is owned only by this thread.

Signed-off-by: Matias Ezequiel Vara Larsen <[email protected]>

Author: MatiasVara

src/libkrun/src/lib.rs

@@ -1,6 +1,12 @@
 #[macro_use]
 extern crate log;
 
+#[cfg(feature = "tee")]
+use crossbeam_channel::unbounded;
+#[cfg(feature = "tee")]
+use kvm_bindings::kvm_memory_attributes;
+#[cfg(feature = "tee")]
+use libc::{fallocate, MADV_DONTNEED, madvise, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::convert::TryInto;
@@ -11,10 +17,14 @@ use std::ffi::CString;
 #[cfg(target_os = "linux")]
 use std::os::fd::AsRawFd;
 use std::os::fd::RawFd;
+#[cfg(feature = "tee")]
+use std::os::raw::c_void;
 use std::path::PathBuf;
 use std::slice;
 use std::sync::atomic::{AtomicI32, Ordering};
 use std::sync::Mutex;
+#[cfg(feature = "tee")]
+use vm_memory::{Address, GuestAddress, GuestMemory, GuestMemoryRegion, MemoryRegionAddress};
 
 #[cfg(target_os = "macos")]
 use crossbeam_channel::unbounded;
@@ -1225,12 +1235,17 @@ pub extern "C" fn krun_start_enter(ctx_id: u32) -> i32 {
     #[cfg(target_os = "macos")]
     let (sender, receiver) = unbounded();
 
+    #[cfg(feature = "tee")]
+    let (io_sender, receiver) = unbounded();
+
     let _vmm = match vmm::builder::build_microvm(
         &ctx_cfg.vmr,
         &mut event_manager,
         ctx_cfg.shutdown_efd,
         #[cfg(target_os = "macos")]
         sender,
+        #[cfg(feature = "tee")]
+        io_sender,
     ) {
         Ok(vmm) => vmm,
         Err(e) => {
@@ -1242,6 +1257,83 @@ pub extern "C" fn krun_start_enter(ctx_id: u32) -> i32 {
     #[cfg(target_os = "macos")]
     let mapper_vmm = _vmm.clone();
 
+    #[cfg(feature = "tee")]
+    let vm = _vmm.lock().unwrap().kvm_vm().fd.clone();
+    #[cfg(feature = "tee")]
+    let guest_mem = _vmm.lock().unwrap().guest_memory().clone();
+    #[cfg(feature = "tee")]
+    let guest_memfd = _vmm.lock().unwrap().guest_memfd_vec.clone();
+
+    #[cfg(feature = "tee")]
+    std::thread::spawn(move || loop {
+        match receiver.recv() {
+            Err(e) => error!("Error in receiver: {:?}", e),
+            Ok(m) => {
+                let _ret = vm
+                    .lock()
+                    .unwrap()
+                    .set_memory_attributes(kvm_memory_attributes {
+                        address: m.addr,
+                        size: m.size,
+                        attributes: m.attributes as u64,
+                        flags: 0,
+                    });
+
+                let region = guest_mem.find_region(GuestAddress(m.addr));
+
+                if let None = region {
+                    error!("Region not found");
+                    continue;
+                }
+
+                let offset = m.addr - region.unwrap().start_addr().raw_value();
+
+                let mut guest_memfd_index = 0;
+
+                // loop through the regions to get the index of the guestmemfd fd
+                // in the future, the fd may be stored in the GuestRegionMmap
+                for (index, region) in guest_mem.iter().enumerate() {
+                    if (region.start_addr().raw_value() + region.size() as u64) > m.addr {
+                        break;
+                    }
+                    guest_memfd_index = index;
+                }
+
+                // from private to shared
+                if m.attributes == 0 {
+                    let ret = unsafe {
+                        fallocate(
+                            *guest_memfd.get(guest_memfd_index).unwrap(),
+                            FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE,
+                            offset as i64,
+                            m.size as i64,
+                        )
+                    };
+                    if ret < 0 {
+                        error!("fallocate has failed");
+                    }
+                // from shared to private
+                } else {
+                    let host_startaddr = region
+                        .unwrap()
+                        .get_host_address(MemoryRegionAddress(m.addr + offset));
+
+                    let ret = unsafe {
+                        madvise(
+                            host_startaddr.unwrap() as *mut c_void,
+                            m.size.try_into().unwrap(),
+                            MADV_DONTNEED,
+                        )
+                    };
+
+                    if ret < 0 {
+                        error!("madvise has failed");
+                    }
+                }
+            }
+        }
+    });
+
     #[cfg(target_os = "macos")]
     std::thread::Builder::new()
         .name("mapping worker".into())

src/vmm/src/builder.rs

@@ -3,8 +3,12 @@
 
 //! Enables pre-boot setup, instantiation and booting of a Firecracker VMM.
 
+#[cfg(feature = "tee")]
+use crate::vstate::MemProperties;
 #[cfg(target_os = "macos")]
-use crossbeam_channel::{unbounded, Sender};
+use crossbeam_channel::unbounded;
+#[cfg(any(target_os = "macos", feature = "tee"))]
+use crossbeam_channel::Sender;
 use std::fmt::{Display, Formatter};
 use std::fs::File;
 use std::io;
@@ -354,6 +358,7 @@ pub fn build_microvm(
     event_manager: &mut EventManager,
     _shutdown_efd: Option<EventFd>,
     #[cfg(target_os = "macos")] _map_sender: Sender<MemoryMapping>,
+    #[cfg(feature = "tee")] io_sender: Sender<MemProperties>,
 ) -> std::result::Result<Arc<Mutex<Vmm>>, StartMicrovmError> {
     #[cfg(not(feature = "efi"))]
     let kernel_bundle = vm_resources
@@ -556,6 +561,7 @@ pub fn build_microvm(
             boot_ip,
             &pio_device_manager.io_bus,
             &exit_evt,
+            // TODO: missing the io_sender
         )
         .map_err(StartMicrovmError::Internal)?;
     }
@@ -572,6 +578,8 @@ pub fn build_microvm(
             &guest_memory,
             GuestAddress(kernel_bundle.guest_addr),
             &exit_evt,
+            #[cfg(feature = "tee")]
+            io_sender,
         )
         .map_err(StartMicrovmError::Internal)?;
 
@@ -1105,13 +1113,16 @@ fn create_vcpus_aarch64(
     guest_mem: &GuestMemoryMmap,
     entry_addr: GuestAddress,
     exit_evt: &EventFd,
+    #[cfg(feature = "tee")] sender_io: Sender<MemProperties>,
 ) -> 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 {
         let mut vcpu = Vcpu::new_aarch64(
             cpu_index,
             &vm.fd.lock().unwrap(),
             exit_evt.try_clone().map_err(Error::EventFd)?,
+            #[cfg(feature = "tee")]
+            sender_io.clone(),
         )
         .map_err(Error::Vcpu)?;
 

src/vmm/src/linux/vstate.rs

@@ -14,6 +14,9 @@ use std::io;
 #[cfg(feature = "tee")]
 use std::os::unix::io::RawFd;
 
+#[cfg(feature = "tee")]
+use kvm_ioctls::VcpuExit::Unsupported;
+
 use std::sync::Arc;
 use std::sync::Mutex;
 
@@ -51,7 +54,7 @@ use kvm_bindings::{
 #[cfg(feature = "tee")]
 use kvm_bindings::{
     kvm_create_guest_memfd, kvm_memory_attributes, kvm_userspace_memory_region2, KVM_API_VERSION,
-    KVM_MEMORY_ATTRIBUTE_PRIVATE, KVM_MEM_GUEST_MEMFD,
+    KVM_MEMORY_ATTRIBUTE_PRIVATE, KVM_MEMORY_EXIT_FLAG_PRIVATE, KVM_MEM_GUEST_MEMFD, KVM_EXIT_MEMORY_FAULT
 };
 #[cfg(not(feature = "tee"))]
 use kvm_bindings::{kvm_userspace_memory_region, KVM_API_VERSION};
@@ -835,6 +838,13 @@ pub struct VcpuConfig {
     pub cpu_template: Option<CpuFeaturesTemplate>,
 }
 
+#[cfg(feature = "tee")]
+pub struct MemProperties {
+    pub addr: u64,
+    pub size: u64,
+    pub attributes: u32,
+}
+
 // Using this for easier explicit type-casting to help IDEs interpret the code.
 type VcpuCell = Cell<Option<*mut Vcpu>>;
 
@@ -857,6 +867,10 @@ pub struct Vcpu {
     #[cfg(target_arch = "aarch64")]
     mpidr: u64,
 
+    // The transmitting end of the events channel which will be given to the vcpu side
+    #[cfg(feature = "tee")]
+    sender_io: Sender<MemProperties>,
+
     // The receiving end of events channel owned by the vcpu side.
     event_receiver: Receiver<VcpuEvent>,
     // The transmitting end of the events channel which will be given to the handler.
@@ -998,7 +1012,12 @@ impl Vcpu {
     /// * `exit_evt` - An `EventFd` that will be written into when this vcpu exits.
     /// * `create_ts` - A timestamp used by the vcpu to calculate its lifetime.
     #[cfg(target_arch = "aarch64")]
-    pub fn new_aarch64(id: u8, vm_fd: &VmFd, exit_evt: EventFd) -> Result<Self> {
+    pub fn new_aarch64(
+        id: u8,
+        vm_fd: &VmFd,
+        exit_evt: EventFd,
+        #[cfg(feature = "tee")] sender_io: Sender<MemProperties>,
+    ) -> Result<Self> {
         let kvm_vcpu = vm_fd.create_vcpu(id as u64).map_err(Error::VcpuFd)?;
         let (event_sender, event_receiver) = unbounded();
         let (response_sender, response_receiver) = unbounded();
@@ -1013,6 +1032,8 @@ impl Vcpu {
             event_sender: Some(event_sender),
             response_receiver: Some(response_receiver),
             response_sender,
+            #[cfg(feature = "tee")]
+            sender_io,
         })
     }
 
@@ -1299,16 +1320,44 @@ impl Vcpu {
                     info!("Received KVM_EXIT_SHUTDOWN signal");
                     Ok(VcpuEmulation::Stopped)
                 }
+                #[cfg(feature = "tee")]
+                VcpuExit::MemoryFault { flags, gpa, size } => {
+                    if flags & !KVM_MEMORY_EXIT_FLAG_PRIVATE as u64 != 0 {
+                        error!("KVM_EXIT_MEMORY_FAULT: Unknown flag {}", flags);
+                        Err(Error::VcpuUnhandledKvmExit)
+                    } else {
+                        // from private to shared
+                        let mut attr = 0;
+                        // from shared to private
+                        if flags & KVM_MEMORY_EXIT_FLAG_PRIVATE as u64
+                            == KVM_MEMORY_EXIT_FLAG_PRIVATE as u64
+                        {
+                            attr = KVM_MEMORY_ATTRIBUTE_PRIVATE;
+                        };
+
+                        let _ = self.sender_io.try_send(MemProperties {
+                            addr: gpa,
+                            size,
+                            attributes: attr,
+                        });
+                        Ok(VcpuEmulation::Handled)
+                    }
+                }
+                // Documentation specifices that when KVM exists with KVM_EXIT_MEMORY_FAULT,
+                // userspace should assume kvm_run.exit_reason is stale/undefined for error numbers
+                // different than EFAULT or EHWPOISON
+                #[cfg(feature = "tee")]
+                Unsupported(KVM_EXIT_MEMORY_FAULT) => Ok(VcpuEmulation::Handled),
+                VcpuExit::InternalError => {
+                    error!("Received KVM_EXIT_INTERNAL_ERROR signal");
+                    Err(Error::VcpuUnhandledKvmExit)
+                }
                 // Documentation specifies that below kvm exits are considered
                 // errors.
                 VcpuExit::FailEntry(reason, vcpu) => {
                     error!("Received KVM_EXIT_FAIL_ENTRY signal: reason={reason}, vcpu={vcpu}");
                     Err(Error::VcpuUnhandledKvmExit)
                 }
-                VcpuExit::InternalError => {
-                    error!("Received KVM_EXIT_INTERNAL_ERROR signal");
-                    Err(Error::VcpuUnhandledKvmExit)
-                }
                 r => {
                     // TODO: Are we sure we want to finish running a vcpu upon
                     // receiving a vm exit that is not necessarily an error?