Merge pull request #99 from Freax13/fix/unmap-kvm-run

unmap KVM run block
Freax13 · Nov 13, 2024 · ca57941 · ca57941
2 parents c1766a1 + d011d8b
commit ca57941
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 4 deletions.
diff --git a/host/mushroom/src/insecure.rs b/host/mushroom/src/insecure.rs
@@ -236,6 +236,7 @@ fn run_kernel_vcpu(
         ap.set_cpuid(&cpuid_entries).unwrap();
 
         let kvm_run = ap.get_kvm_run_block().unwrap();
+        let kvm_run = kvm_run.as_ptr();
 
         let mut sregs = ap.get_sregs().unwrap();
         sregs.es = KvmSegment::DATA64;

diff --git a/host/mushroom/src/kvm.rs b/host/mushroom/src/kvm.rs
@@ -2,11 +2,13 @@
 
 use std::{
     array::{self, from_fn},
+    ffi::c_void,
     fmt,
     fs::OpenOptions,
     mem::{size_of, size_of_val},
     num::{NonZeroU32, NonZeroUsize},
     os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd},
+    ptr::NonNull,
 };
 
 use anyhow::{ensure, Context, Result};
@@ -827,8 +829,7 @@ impl VcpuHandle {
         Ok(())
     }
 
-    pub fn get_kvm_run_block(&self) -> Result<VolatilePtr<KvmRun>> {
-        // FIXME: unmap the memory
+    pub fn get_kvm_run_block(&self) -> Result<KvmRunBox> {
         let res = unsafe {
             nix::sys::mman::mmap(
                 None,
@@ -840,8 +841,7 @@ impl VcpuHandle {
             )
         };
         let ptr = res.context("failed to map vcpu kvm_run block")?;
-        let ptr = unsafe { VolatilePtr::new(ptr.cast()) };
-        Ok(ptr)
+        Ok(KvmRunBox { ptr })
     }
 
     /// Returns `true` if the cpu ran uninterrupted or returns `false` if the
@@ -863,6 +863,23 @@ impl VcpuHandle {
     }
 }
 
+pub struct KvmRunBox {
+    ptr: NonNull<c_void>,
+}
+
+impl KvmRunBox {
+    pub fn as_ptr(&self) -> VolatilePtr<'_, KvmRun> {
+        unsafe { VolatilePtr::new(self.ptr.cast()) }
+    }
+}
+
+impl Drop for KvmRunBox {
+    fn drop(&mut self) {
+        let res = unsafe { nix::sys::mman::munmap(self.ptr, size_of::<KvmRun>()) };
+        res.unwrap();
+    }
+}
+
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
 #[repr(C)]
 pub struct KvmRegs {

diff --git a/host/mushroom/src/snp.rs b/host/mushroom/src/snp.rs
@@ -248,6 +248,7 @@ impl VmContext {
     pub fn run_supervisor(&mut self) -> Result<MushroomResult> {
         let mut output = Vec::new();
         let kvm_run = self.bsp.get_kvm_run_block()?;
+        let kvm_run = kvm_run.as_ptr();
 
         loop {
             let exit = kvm_run.read().exit();
@@ -414,6 +415,7 @@ impl VmContext {
             ap.set_msr(DEBUG_CTL, 1).unwrap();
 
             let kvm_run = ap.get_kvm_run_block().unwrap();
+            let kvm_run = kvm_run.as_ptr();
 
             std::thread::park();
 

diff --git a/host/mushroom/src/tdx.rs b/host/mushroom/src/tdx.rs
@@ -281,6 +281,7 @@ impl VmContext {
         sender: &Sender<OutputEvent>,
     ) -> Result<()> {
         let kvm_run = bsp.get_kvm_run_block()?;
+        let kvm_run = kvm_run.as_ptr();
 
         while !done.load(Ordering::Relaxed) {
             let exit = kvm_run.read().exit();