Jitter: add remove_memory_page api

5 files changed, 74 insertions, 6 deletions

diff --git a/miasm/jitter/vm_mngr.c b/miasm/jitter/vm_mngr.c
index 53ec9065..0c8a0586 100644
--- a/miasm/jitter/vm_mngr.c
+++ b/miasm/jitter/vm_mngr.c
@@ -820,7 +820,32 @@ void add_memory_page(vm_mngr_t* vm_mngr, struct memory_page_node* mpn_a)
 
 	vm_mngr->memory_pages_array[i] = *mpn_a;
 	vm_mngr->memory_pages_number ++;
+}
+
+void remove_memory_page(vm_mngr_t* vm_mngr, uint64_t ad)
+{
+  struct memory_page_node * mpn;
+  int i;
 
+  i = find_page_node(vm_mngr->memory_pages_array,
+		     ad,
+		     0,
+		     vm_mngr->memory_pages_number - 1);
+  if (i < 0) {
+    return;
+  }
+
+  mpn = &vm_mngr->memory_pages_array[i];
+  free(mpn->name);
+  free(mpn->ad_hp);
+  memmove(&vm_mngr->memory_pages_array[i],
+  	      &vm_mngr->memory_pages_array[i+1],
+  	      sizeof(struct memory_page_node) * (vm_mngr->memory_pages_number - i - 1)
+  	      );
+  vm_mngr->memory_pages_number --;
+  vm_mngr->memory_pages_array = realloc(vm_mngr->memory_pages_array,
+					sizeof(struct memory_page_node) *
+					(vm_mngr->memory_pages_number));
 }
 
 /* Return a char* representing the repr of vm_mngr_t object */
@@ -961,5 +986,3 @@ _MIASM_EXPORT uint64_t get_exception_flag(vm_mngr_t* vm_mngr)
 {
 	return vm_mngr->exception_flags;
 }
-
-
diff --git a/miasm/jitter/vm_mngr.h b/miasm/jitter/vm_mngr.h
index 7ae44d99..946d3b48 100644
--- a/miasm/jitter/vm_mngr.h
+++ b/miasm/jitter/vm_mngr.h
@@ -237,6 +237,7 @@ void reset_memory_page_pool(vm_mngr_t* vm_mngr);
 void reset_code_bloc_pool(vm_mngr_t* vm_mngr);
 void dump_code_bloc_pool(vm_mngr_t* vm_mngr);
 void add_memory_page(vm_mngr_t* vm_mngr, struct memory_page_node* mpn_a);
+void remove_memory_page(vm_mngr_t* vm_mngr, uint64_t ad);
 
 
 void init_memory_breakpoint(vm_mngr_t* vm_mngr);
diff --git a/miasm/jitter/vm_mngr_py.c b/miasm/jitter/vm_mngr_py.c
index 69e62fef..8bee0586 100644
--- a/miasm/jitter/vm_mngr_py.c
+++ b/miasm/jitter/vm_mngr_py.c
@@ -129,6 +129,23 @@ PyObject* vm_add_memory_page(VmMngr* self, PyObject* args)
 }
 
 
+PyObject* vm_remove_memory_page(VmMngr* self, PyObject* args)
+{
+  PyObject *addr;
+  uint64_t page_addr;
+
+  if (!PyArg_ParseTuple(args, "O", &addr))
+    RAISE(PyExc_TypeError,"Cannot parse arguments");
+
+  PyGetInt_uint64_t(addr, page_addr);
+
+  remove_memory_page(&self->vm_mngr, page_addr);
+
+  Py_INCREF(Py_None);
+  return Py_None;
+
+
+}
 
 PyObject* vm_set_mem_access(VmMngr* self, PyObject* args)
 {
@@ -877,6 +894,8 @@ static PyMethodDef VmMngr_methods[] = {
 	{"add_memory_page",(PyCFunction)vm_add_memory_page, METH_VARARGS,
 	 "add_memory_page(address, access, content [, cmt]) -> Maps a memory page at @address of len(@content) bytes containing @content with protection @access\n"
 	"@cmt is a comment linked to the memory page"},
+	{"remove_memory_page",(PyCFunction)vm_remove_memory_page, METH_VARARGS,
+	 "remove_memory_page(address) -> removes a previously allocated memory page at @address"},
 	{"add_memory_breakpoint",(PyCFunction)vm_add_memory_breakpoint, METH_VARARGS,
 	 "add_memory_breakpoint(address, size, access) -> Add a memory breakpoint at @address of @size bytes with @access type"},
 	{"remove_memory_breakpoint",(PyCFunction)vm_remove_memory_breakpoint, METH_VARARGS,
diff --git a/miasm/os_dep/win_api_x86_32.py b/miasm/os_dep/win_api_x86_32.py
index 606c8792..5e2fb680 100644
--- a/miasm/os_dep/win_api_x86_32.py
+++ b/miasm/os_dep/win_api_x86_32.py
@@ -757,14 +757,31 @@ def kernel32_VirtualProtect(jitter):
         old = jitter.vm.get_mem_access(args.lpvoid)
         jitter.vm.set_u32(args.lpfloldprotect, ACCESS_DICT_INV[old])
 
-    log.warn("set page %x %x", args.lpvoid, args.dwsize)
+    protect_addr = args.lpvoid - (args.lpvoid % winobjs.alloc_align)
     for addr, data in jitter.vm.get_all_memory().items():
         size = data["size"]
-        # Multi-page
-        if addr <= args.lpvoid < addr + size:
+        # Exact match
+        if addr == protect_addr and args.dwsize == size:
             log.warn("set page %x %x", addr, ACCESS_DICT[flnewprotect])
             jitter.vm.set_mem_access(addr, ACCESS_DICT[flnewprotect])
-
+            continue
+        # Split memory area
+        if addr <= protect_addr < protect_addr + args.dwsize <= addr + size:
+            part1 = jitter.vm.get_mem(addr, protect_addr + args.dwsize - addr)
+            part2 = jitter.vm.get_mem(protect_addr + args.dwsize,  addr + size - protect_addr - args.dwsize)
+            old_access = jitter.vm.get_mem_access(addr)
+            jitter.vm.remove_memory_page(addr)
+            log.warn("create page %x %x", addr, old_access)
+            log.warn("create page %x %x", protect_addr + args.dwsize, ACCESS_DICT[flnewprotect])
+            jitter.vm.add_memory_page(addr, old_access, part1, "VirtualProtect split")
+            jitter.vm.add_memory_page(protect_addr + args.dwsize, ACCESS_DICT[flnewprotect], part2, "VirtualProtect split")
+            continue
+        # Multi page
+        if (addr <= protect_addr < addr + size or
+            addr < protect_addr + args.dwsize <= addr + size):
+            log.warn("set page %x %x", addr, ACCESS_DICT[flnewprotect])
+            jitter.vm.set_mem_access(addr, ACCESS_DICT[flnewprotect])
+            continue
     jitter.func_ret_stdcall(ret_ad, 1)
 
 
diff --git a/test/jitter/vm_mngr.py b/test/jitter/vm_mngr.py
index 468fb347..0fec1734 100644
--- a/test/jitter/vm_mngr.py
+++ b/test/jitter/vm_mngr.py
@@ -25,3 +25,11 @@ for i, access_right in enumerate(shuffled_rights):
 # Check for modification
 for i, access_right in enumerate(shuffled_rights):
     assert myjit.vm.get_mem_access(base_addr + i * page_size) == access_right
+
+# Remove pages
+for i in range(len(rights)):
+    myjit.vm.remove_memory_page(base_addr + i * page_size)
+
+# Add pages again
+for i, access_right in enumerate(rights):
+    myjit.vm.add_memory_page(base_addr + i * page_size, access_right, data)