From d011d8b8e9d4043b2d7e206de24afc112a67e5da Mon Sep 17 00:00:00 2001
From: Tom Dohrmann <erbse.13@gmx.de>
Date: Wed, 13 Nov 2024 15:41:33 +0100
Subject: [PATCH] unmap KVM run block

It turns out if we don't unmap this, the VM won't get properly shut
down and we eventually run out of HKIDs.
---
 host/mushroom/src/insecure.rs |  1 +
 host/mushroom/src/kvm.rs      | 25 +++++++++++++++++++++----
 host/mushroom/src/snp.rs      |  2 ++
 host/mushroom/src/tdx.rs      |  1 +
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/host/mushroom/src/insecure.rs b/host/mushroom/src/insecure.rs
index 11cdba9f..59b58073 100644
--- 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
index 73889f62..7edc3d21 100644
--- 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
index 7b80f474..d5ac3495 100644
--- 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
index 2e28af4e..5c3c5cca 100644
--- 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();