kvm: Set the right bit in kvm_dirty_regs when setting a register (#251)

In #248, `cr2` was updated but `kvm_dirty_regs` wasn't set appropriately so the change doesn't get flushed to the vcpu.

Author: 0vercl0k

scripts/gen_coveragefile_binja.py

@@ -1,32 +1,32 @@
-"""Generate a coverage file for WTF using Binary Ninja.
-
-@australeo - 2023
-"""
-from json import dump
-from pathlib import Path
-from binaryninja import PluginCommand, interaction
-
-def generate_coverage_file(bv):
-    # bv.file.filename: 'C:/path/to/binary.bndb'
-    name = Path(bv.file.filename).stem
-
-    bb_list = []
-
-    for block in bv.basic_blocks:
-        bb_list.append(block.start - bv.start)
-
-    json_object = {
-        "name": name,
-        "addresses": bb_list
-    }
-
-    output_file = interaction.get_save_filename_input("Filename: ", ".cov", name + ".cov")
-
-    with open(output_file, "w", encoding="utf-8") as file:
-        dump(json_object, file)
-
-def register_plugin():
-    PluginCommand.register("Generate WTF coverage file", \
-        "Generate WTF coverage file", generate_coverage_file)
-
-register_plugin()
+"""Generate a coverage file for WTF using Binary Ninja.
+
+@australeo - 2023
+"""
+from json import dump
+from pathlib import Path
+from binaryninja import PluginCommand, interaction
+
+def generate_coverage_file(bv):
+    # bv.file.filename: 'C:/path/to/binary.bndb'
+    name = Path(bv.file.filename).stem
+
+    bb_list = []
+
+    for block in bv.basic_blocks:
+        bb_list.append(block.start - bv.start)
+
+    json_object = {
+        "name": name,
+        "addresses": bb_list
+    }
+
+    output_file = interaction.get_save_filename_input("Filename: ", ".cov", name + ".cov")
+
+    with open(output_file, "w", encoding="utf-8") as file:
+        dump(json_object, file)
+
+def register_plugin():
+    PluginCommand.register("Generate WTF coverage file", \
+        "Generate WTF coverage file", generate_coverage_file)
+
+register_plugin()

src/wtf/kvm_backend.cc

@@ -1968,6 +1968,37 @@ uint64_t KvmBackend_t::GetReg(const Registers_t Reg) const {
   return 0;
 }
 
+uint64_t GetKvmDirtyRegsFromRegister(const Registers_t Reg) {
+  switch (Reg) {
+  case Registers_t::Rax:
+  case Registers_t::Rbx:
+  case Registers_t::Rcx:
+  case Registers_t::Rdx:
+  case Registers_t::Rsi:
+  case Registers_t::Rdi:
+  case Registers_t::Rip:
+  case Registers_t::Rsp:
+  case Registers_t::Rbp:
+  case Registers_t::R8:
+  case Registers_t::R9:
+  case Registers_t::R10:
+  case Registers_t::R11:
+  case Registers_t::R12:
+  case Registers_t::R13:
+  case Registers_t::R14:
+  case Registers_t::R15:
+  case Registers_t::Rflags:
+    return KVM_SYNC_X86_REGS;
+
+  case Registers_t::Cr2:
+  case Registers_t::Cr3:
+    return KVM_SYNC_X86_SREGS;
+  }
+
+  std::abort();
+  return 0;
+}
+
 uint64_t KvmBackend_t::SetReg(const Registers_t Reg, const uint64_t Value) {
   switch (Reg) {
   case Registers_t::Rax: {
@@ -2072,10 +2103,11 @@ uint64_t KvmBackend_t::SetReg(const Registers_t Reg, const uint64_t Value) {
   }
 
   //
-  // Tell KVM to flush the regs into the VCPU next time it runs.
+  // Tell KVM to flush the appropriate set of regs into the VCPU next time it
+  // runs.
   //
 
-  Run_->kvm_dirty_regs |= KVM_SYNC_X86_REGS;
+  Run_->kvm_dirty_regs |= GetKvmDirtyRegsFromRegister(Reg);
   return Value;
 }