From 09a95f5e74c47d43a1f6962fca108466d8cc31e5 Mon Sep 17 00:00:00 2001
From: Zach Reizner <zachr@google.com>
Date: Wed, 2 Aug 2017 14:08:29 -0700
Subject: [PATCH] kvm: add device memory interface to Vm

This interface is to dynamically add and remove memory mappings to guest
physical address space.

TEST=cargo test
BUG=chromium:738638

Change-Id: I695775289d56686ef16a1e3cf7640c97d5da2662
Reviewed-on: https://chromium-review.googlesource.com/599040
Commit-Ready: Zach Reizner <zachr@chromium.org>
Tested-by: Zach Reizner <zachr@chromium.org>
Reviewed-by: Dylan Reid <dgreid@chromium.org>
---
 kvm/src/lib.rs | 72 ++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 64 insertions(+), 8 deletions(-)

diff --git a/kvm/src/lib.rs b/kvm/src/lib.rs
index 1a39cf12a7..084365c664 100644
--- a/kvm/src/lib.rs
+++ b/kvm/src/lib.rs
@@ -12,10 +12,12 @@ mod cap;
 
 use std::fs::File;
 use std::mem;
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
 use std::os::raw::*;
 use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 
-use libc::{open, O_RDWR, EINVAL, ENOSPC};
+use libc::{open, O_RDWR, EINVAL, ENOSPC, ENOENT};
 
 use kvm_sys::*;
 
@@ -150,6 +152,7 @@ impl Into<u64> for NoDatamatch {
 pub struct Vm {
     vm: File,
     guest_mem: GuestMemory,
+    device_memory: HashMap<u32, MemoryMapping>,
     next_mem_slot: usize,
 }
 
@@ -177,6 +180,7 @@ impl Vm {
             Ok(Vm {
                 vm: vm_file,
                 guest_mem: guest_mem,
+                device_memory: HashMap::new(),
                 next_mem_slot: next_mem_slot,
             })
         } else {
@@ -186,28 +190,58 @@ impl Vm {
 
     /// Inserts the given `MemoryMapping` into the VM's address space at `guest_addr`.
     ///
-    /// Note that memory inserted into the VM's address space must not overlap
-    /// with any other memory slot's region.
-    pub fn add_device_memory(&mut self, guest_addr: GuestAddress, mem: MemoryMapping) -> Result<()> {
+    /// The slot that was assigned the device memory mapping is returned on success. The slot can be
+    /// given to `Vm::remove_device_memory` to remove the memory from the VM's address space and
+    /// take back ownership of `mem`.
+    ///
+    /// Note that memory inserted into the VM's address space must not overlap with any other memory
+    /// slot's region.
+    pub fn add_device_memory(&mut self,
+                             guest_addr: GuestAddress,
+                             mem: MemoryMapping)
+                             -> Result<u32> {
         if guest_addr < self.guest_mem.end_addr() {
             return Err(Error::new(ENOSPC));
         }
 
+        let slot = self.next_mem_slot as u32;
+
         // Safe because we check that the given guest address is valid and has no overlaps. We also
         // know that the pointer and size are correct because the MemoryMapping interface ensures
-        // this.
+        // this. We take ownership of the memory mapping so that it won't be unmapped until the slot
+        // is removed.
         unsafe {
-            set_user_memory_region(&self.vm, self.next_mem_slot as u32,
+            set_user_memory_region(&self.vm, slot,
                                         guest_addr.offset() as u64,
                                         mem.size() as u64,
                                         mem.as_ptr() as u64)?;
         };
+        self.device_memory.insert(slot, mem);
         self.next_mem_slot += 1;
 
-        Ok(())
+        Ok(slot)
     }
 
-    /// Gets a reference to the memory at the given address in the VM's address space.
+    /// Removes device memory that was previously added at the given slot.
+    ///
+    /// Ownership of the host memory mapping associated with the given slot is returned on success.
+    pub fn remove_device_memory(&mut self, slot: u32) -> Result<MemoryMapping> {
+        match self.device_memory.entry(slot) {
+            Entry::Occupied(entry) => {
+                // Safe because the slot is checked against the list of device memory slots.
+                unsafe {
+                    set_user_memory_region(&self.vm, slot, 0, 0, 0)?;
+                }
+                Ok(entry.remove())
+            }
+            _ => Err(Error::new(-ENOENT))
+        }
+    }
+
+    /// Gets a reference to the guest memory owned by this VM.
+    ///
+    /// Note that `GuestMemory` does not include any device memory that may have been added after
+    /// this VM was constructed.
     pub fn get_memory(&self) -> &GuestMemory {
         &self.guest_mem
     }
@@ -705,6 +739,28 @@ mod tests {
         vm.add_device_memory(GuestAddress(0x1000), mem).unwrap();
     }
 
+    #[test]
+    fn remove_memory() {
+        let kvm = Kvm::new().unwrap();
+        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
+        let mut vm = Vm::new(&kvm, gm).unwrap();
+        let mem_size = 0x1000;
+        let mem = MemoryMapping::new(mem_size).unwrap();
+        let mem_ptr = mem.as_ptr();
+        let slot = vm.add_device_memory(GuestAddress(0x1000), mem).unwrap();
+        let mem = vm.remove_device_memory(slot).unwrap();
+        assert_eq!(mem.size(), mem_size);
+        assert_eq!(mem.as_ptr(), mem_ptr);
+    }
+
+    #[test]
+    fn remove_invalid_memory() {
+        let kvm = Kvm::new().unwrap();
+        let gm = GuestMemory::new(&vec![(GuestAddress(0), 0x1000)]).unwrap();
+        let mut vm = Vm::new(&kvm, gm).unwrap();
+        assert!(vm.remove_device_memory(0).unwrap());
+    }
+
      #[test]
     fn overlap_memory() {
         let kvm = Kvm::new().unwrap();