1 files changed, 17 insertions, 59 deletions

diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index a995f3a..c63fbfe 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -305,7 +305,7 @@ impl VmMemoryRequest {
                     Err(e) => VmMemoryResponse::Err(e),
                 }
             }
-            UnregisterMemory(slot) => match vm.remove_mmio_memory(slot) {
+            UnregisterMemory(slot) => match vm.remove_memory_region(slot) {
                 Ok(_) => VmMemoryResponse::Ok,
                 Err(e) => VmMemoryResponse::Err(e),
             },
@@ -347,7 +347,7 @@ impl VmMemoryRequest {
                     Ok(v) => v,
                     Err(_e) => return VmMemoryResponse::Err(SysError::new(EINVAL)),
                 };
-                match vm.add_mmio_memory(GuestAddress(gpa), mmap, false, false) {
+                match vm.add_memory_region(GuestAddress(gpa), Box::new(mmap), false, false) {
                     Ok(_) => VmMemoryResponse::Ok,
                     Err(e) => VmMemoryResponse::Err(e),
                 }
@@ -479,19 +479,10 @@ impl VmMsyncRequest {
     pub fn execute(&self, vm: &mut Vm) -> VmMsyncResponse {
         use self::VmMsyncRequest::*;
         match *self {
-            MsyncArena { slot, offset, size } => {
-                if let Some(arena) = vm.get_mmap_arena(slot) {
-                    match arena.msync(offset, size) {
-                        Ok(()) => VmMsyncResponse::Ok,
-                        Err(e) => match e {
-                            MmapError::SystemCallFailed(errno) => VmMsyncResponse::Err(errno),
-                            _ => VmMsyncResponse::Err(SysError::new(EINVAL)),
-                        },
-                    }
-                } else {
-                    VmMsyncResponse::Err(SysError::new(EINVAL))
-                }
-            }
+            MsyncArena { slot, offset, size } => match vm.mysnc_memory_region(slot, offset, size) {
+                Ok(()) => VmMsyncResponse::Ok,
+                Err(e) => VmMsyncResponse::Err(e),
+            },
         }
     }
 }
@@ -544,7 +535,7 @@ fn register_memory(
     allocator: &mut SystemAllocator,
     fd: &dyn AsRawFd,
     size: usize,
-    allocation: Option<(Alloc, u64)>,
+    pci_allocation: Option<(Alloc, u64)>,
 ) -> Result<(u64, u32)> {
     let mmap = match MemoryMapping::from_fd(fd, size) {
         Ok(v) => v,
@@ -552,54 +543,21 @@ fn register_memory(
         _ => return Err(SysError::new(EINVAL)),
     };
 
-    let addr = match allocation {
-        Some((
-            Alloc::PciBar {
-                bus,
-                dev,
-                func,
-                bar,
-            },
-            offset,
-        )) => {
-            match allocator
-                .mmio_allocator(MmioType::High)
-                .get(&Alloc::PciBar {
-                    bus,
-                    dev,
-                    func,
-                    bar,
-                }) {
-                Some((start_addr, length, _)) => {
-                    let address = *start_addr + offset;
-                    let range = *start_addr..*start_addr + *length;
-                    let end = address + (size as u64);
-                    match (range.contains(&address), range.contains(&end)) {
-                        (true, true) => address,
-                        _ => return Err(SysError::new(EINVAL)),
-                    }
-                }
-                None => return Err(SysError::new(EINVAL)),
-            }
-        }
+    let addr = match pci_allocation {
+        Some(pci_allocation) => allocator
+            .mmio_allocator(MmioType::High)
+            .address_from_pci_offset(pci_allocation.0, pci_allocation.1, size as u64)
+            .map_err(|_e| SysError::new(EINVAL))?,
         None => {
             let alloc = allocator.get_anon_alloc();
-            match allocator.mmio_allocator(MmioType::High).allocate(
-                size as u64,
-                alloc,
-                "vmcontrol_register_memory".to_string(),
-            ) {
-                Ok(a) => a,
-                _ => return Err(SysError::new(EINVAL)),
-            }
+            allocator
+                .mmio_allocator(MmioType::High)
+                .allocate(size as u64, alloc, "vmcontrol_register_memory".to_string())
+                .map_err(|_e| SysError::new(EINVAL))?
         }
-        _ => return Err(SysError::new(EINVAL)),
     };
 
-    let slot = match vm.add_mmio_memory(GuestAddress(addr), mmap, false, false) {
-        Ok(v) => v,
-        Err(e) => return Err(e),
-    };
+    let slot = vm.add_memory_region(GuestAddress(addr), Box::new(mmap), false, false)?;
 
     Ok((addr >> 12, slot))
 }