From fdecba695814affa573d8eef2d2cf85f697b2292 Mon Sep 17 00:00:00 2001
From: Muminul Islam <muislam@microsoft.com>
Date: Thu, 8 Jul 2021 10:14:24 -0700
Subject: [PATCH] hypervisor: MSHV needs gpa to retrieve dirty logs

Right now, get_dirty_log API has two parameters,
slot and memory_size.
MSHV needs gpa to retrieve the page states. GPA is
needed as MSHV returns the state base on PFN.

Signed-off-by: Muminul Islam <muislam@microsoft.com>
---
 hypervisor/src/kvm/mod.rs  |  2 +-
 hypervisor/src/mshv/mod.rs | 13 +++++++++----
 hypervisor/src/vm.rs       |  2 +-
 vmm/src/memory_manager.rs  |  2 +-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs
index 9b865de2b..ae205bfc5 100644
--- a/hypervisor/src/kvm/mod.rs
+++ b/hypervisor/src/kvm/mod.rs
@@ -472,7 +472,7 @@ impl vm::Vm for KvmVm {
     ///
     /// Get dirty pages bitmap (one bit per page)
     ///
-    fn get_dirty_log(&self, slot: u32, memory_size: u64) -> vm::Result<Vec<u64>> {
+    fn get_dirty_log(&self, slot: u32, _base_gpa: u64, memory_size: u64) -> vm::Result<Vec<u64>> {
         self.fd
             .get_dirty_log(slot, memory_size as usize)
             .map_err(|e| vm::HypervisorVmError::GetDirtyLog(e.into()))
diff --git a/hypervisor/src/mshv/mod.rs b/hypervisor/src/mshv/mod.rs
index a9033e894..3410899c5 100644
--- a/hypervisor/src/mshv/mod.rs
+++ b/hypervisor/src/mshv/mod.rs
@@ -33,6 +33,7 @@ use std::fs::File;
 use std::os::unix::io::AsRawFd;
 use std::sync::RwLock;
 
+const DIRTY_BITMAP_CLEAR_DIRTY: u64 = 0x4;
 pub const PAGE_SHIFT: usize = 12;
 
 #[derive(Debug, Default, Copy, Clone, Serialize, Deserialize)]
@@ -895,10 +896,14 @@ impl vm::Vm for MshvVm {
     ///
     /// Get dirty pages bitmap (one bit per page)
     ///
-    fn get_dirty_log(&self, _slot: u32, _memory_size: u64) -> vm::Result<Vec<u64>> {
-        Err(vm::HypervisorVmError::GetDirtyLog(anyhow!(
-            "get_dirty_log not implemented"
-        )))
+    fn get_dirty_log(&self, _slot: u32, base_gpa: u64, memory_size: u64) -> vm::Result<Vec<u64>> {
+        self.fd
+            .get_dirty_log(
+                base_gpa >> PAGE_SHIFT,
+                memory_size as usize,
+                DIRTY_BITMAP_CLEAR_DIRTY,
+            )
+            .map_err(|e| vm::HypervisorVmError::GetDirtyLog(e.into()))
     }
 }
 pub use hv_cpuid_entry as CpuIdEntry;
diff --git a/hypervisor/src/vm.rs b/hypervisor/src/vm.rs
index 11ab20944..112fc0014 100644
--- a/hypervisor/src/vm.rs
+++ b/hypervisor/src/vm.rs
@@ -285,7 +285,7 @@ pub trait Vm: Send + Sync {
     /// Stop logging dirty pages
     fn stop_dirty_log(&self) -> Result<()>;
     /// Get dirty pages bitmap
-    fn get_dirty_log(&self, slot: u32, memory_size: u64) -> Result<Vec<u64>>;
+    fn get_dirty_log(&self, slot: u32, base_gpa: u64, memory_size: u64) -> Result<Vec<u64>>;
     #[cfg(feature = "tdx")]
     /// Initalize TDX on this VM
     fn tdx_init(&self, cpuid: &CpuId, max_vcpus: u32) -> Result<()>;
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index c6f6a5ef3..939dff707 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -1502,7 +1502,7 @@ impl MemoryManager {
         let page_size = 4096; // TODO: Does this need to vary?
         let mut table = MemoryRangeTable::default();
         for r in &self.guest_ram_mappings {
-            let vm_dirty_bitmap = self.vm.get_dirty_log(r.slot, r.size).map_err(|e| {
+            let vm_dirty_bitmap = self.vm.get_dirty_log(r.slot, r.gpa, r.size).map_err(|e| {
                 MigratableError::MigrateSend(anyhow!("Error getting VM dirty log {}", e))
             })?;
             let vmm_dirty_bitmap = match self.guest_memory.memory().find_region(GuestAddress(r.gpa))