Jitter: Use mem access for automod

4 files changed, 106 insertions, 114 deletions

diff --git a/miasm2/jitter/jitcore.py b/miasm2/jitter/jitcore.py
index 42e6d7a9..04bd707a 100644
--- a/miasm2/jitter/jitcore.py
+++ b/miasm2/jitter/jitcore.py
@@ -51,8 +51,6 @@ class JitCore(object):
         self.blocs_mem_interval = interval()
         self.disasm_cb = None
         self.split_dis = set()
-        self.addr_mod = interval()
-
         self.options = {"jit_maxline": 50  # Maximum number of line jitted
                         }
 
@@ -261,12 +259,21 @@ class JitCore(object):
 
         return modified_blocs
 
-    def updt_automod_code(self, vm):
-        """Remove code jitted in range self.addr_mod
+    def updt_automod_code_range(self, vm, mem_range):
+        """Remove jitted code in range @mem_range
         @vm: VmMngr instance
+        @mem_range: list of start/stop addresses
         """
-        for addr_start, addr_stop in vm.get_memory_write():
+        for addr_start, addr_stop in mem_range:
             self.del_bloc_in_range(addr_start, addr_stop)
         self.__updt_jitcode_mem_range(vm)
-        self.addr_mod = interval()
         vm.reset_memory_access()
+
+    def updt_automod_code(self, vm):
+        """Remove jitted code updated by memory write
+        @vm: VmMngr instance
+        """
+        mem_range = []
+        for addr_start, addr_stop in vm.get_memory_write():
+            mem_range.append((addr_start, addr_stop))
+        self.updt_automod_code_range(vm, mem_range)
diff --git a/miasm2/jitter/jitload.py b/miasm2/jitter/jitload.py
index 1447c899..d8393230 100644
--- a/miasm2/jitter/jitload.py
+++ b/miasm2/jitter/jitload.py
@@ -267,8 +267,7 @@ class jitter:
         self.breakpoints_handler.add_callback(addr, callback)
         self.jit.add_disassembly_splits(addr)
         # De-jit previously jitted blocks
-        self.jit.addr_mod = interval([(addr, addr)])
-        self.jit.updt_automod_code(self.vm)
+        self.jit.updt_automod_code_range(self.vm, [(addr, addr)])
 
     def set_breakpoint(self, addr, *args):
         """Set callbacks associated with addr.
diff --git a/miasm2/jitter/vm_mngr.c b/miasm2/jitter/vm_mngr.c
index 8022b32c..df8326d9 100644
--- a/miasm2/jitter/vm_mngr.c
+++ b/miasm2/jitter/vm_mngr.c
@@ -78,6 +78,37 @@ const uint8_t parity_table[256] = {
 
 // #define DEBUG_MIASM_AUTOMOD_CODE
 
+void memory_access_list_init(struct memory_access_list * access)
+{
+	access->array = NULL;
+	access->allocated = 0;
+	access->num = 0;
+}
+
+void memory_access_list_reset(struct memory_access_list * access)
+{
+	if (access->array) {
+		free(access->array);
+		access->array = NULL;
+	}
+	access->allocated = 0;
+	access->num = 0;
+}
+
+void memory_access_list_add(struct memory_access_list * access, uint64_t start, uint64_t stop)
+{
+	if (access->num >= access->allocated) {
+		if (access->allocated == 0)
+			access->allocated = 1;
+		else
+			access->allocated *= 2;
+		access->array = realloc(access->array, access->allocated * sizeof(struct memory_access));
+	}
+	access->array[access->num].start = start;
+	access->array[access->num].stop = stop;
+	access->num += 1;
+}
+
 
 
 uint16_t set_endian16(vm_mngr_t* vm_mngr, uint16_t val)
@@ -393,92 +424,56 @@ void dump_code_bloc(vm_mngr_t* vm_mngr)
 
 }
 
-void add_range_to_pylist(PyObject* pylist, uint64_t addr1, uint64_t addr2)
+void add_range_to_list(struct memory_access_list * access, uint64_t addr1, uint64_t addr2)
 {
-	PyObject* range;
-	PyObject* element;
-	int list_size;
-	uint64_t addr_start, addr_stop;
-
-	list_size = PyList_Size(pylist);
-
-	if (list_size > 0) {
+	if (access->num > 0) {
 		/* Check match on upper bound */
-		 element = PyList_GetItem(pylist, list_size - 1);
-
-		 addr_start = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 0));
-		 addr_stop = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 1));
-
-		 if (addr_stop == addr1) {
-			 range = PyTuple_New(2);
-			 PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr_start));
-			 PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr2));
-			 PyList_SetItem(pylist, list_size - 1, range);
+		 if (access->array[access->num-1].stop == addr1) {
+			 access->array[access->num-1].stop = addr2;
 			 return;
-
 		 }
 
 		/* Check match on lower bound */
-		 element = PyList_GetItem(pylist, 0);
-		 addr_start = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 0));
-
-		 if (addr_start == addr2) {
-			 range = PyTuple_New(2);
-			 PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr1));
-			 PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr_start));
-			 PyList_SetItem(pylist, 0, range);
+		 if (access->array[0].start == addr2) {
+			 access->array[0].start = addr1;
 			 return;
 		 }
-
 	}
 
 	/* No merge, add to the list */
-	range = PyTuple_New(2);
-	PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr1));
-	PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr2));
-
-	PyList_Append(pylist, range);
+	memory_access_list_add(access, addr1, addr2);
 }
 
 
 void add_mem_read(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t size)
 {
-	add_range_to_pylist(vm_mngr->memory_r, addr, addr + size);
+	add_range_to_list(&(vm_mngr->memory_r), addr, addr + size);
 }
 
 void add_mem_write(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t size)
 {
-	add_range_to_pylist(vm_mngr->memory_w, addr, addr + size);
-	vm_mngr->write_num++;
+	add_range_to_list(&(vm_mngr->memory_w), addr, addr + size);
 }
 
 void check_invalid_code_blocs(vm_mngr_t* vm_mngr)
 {
 	int i;
-	int list_size;
-	PyObject* element;
 	struct code_bloc_node * cbp;
-	uint64_t addr_start, addr_stop;
-
-	if (vm_mngr->write_num == 0)
-		return;
-
-	list_size = PyList_Size(vm_mngr->memory_w);
-
-	LIST_FOREACH(cbp, &vm_mngr->code_bloc_pool, next){
+	for (i=0;i<vm_mngr->memory_w.num; i++) {
 		if (vm_mngr->exception_flags & EXCEPT_CODE_AUTOMOD)
 			break;
-		for (i=0;i<list_size; i++) {
-			element = PyList_GetItem(vm_mngr->memory_w, i);
-			addr_start = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 0));
-			addr_stop = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 1));
+		if (vm_mngr->memory_w.array[i].stop <= vm_mngr->code_bloc_pool_ad_min ||
+		    vm_mngr->memory_w.array[i].start >=vm_mngr->code_bloc_pool_ad_max)
+			continue;
 
-			if ((cbp->ad_start < addr_stop) &&
-			    (addr_start < cbp->ad_stop)){
+		LIST_FOREACH(cbp, &vm_mngr->code_bloc_pool, next){
+			if ((cbp->ad_start < vm_mngr->memory_w.array[i].stop) &&
+			    (vm_mngr->memory_w.array[i].start < cbp->ad_stop)){
 #ifdef DEBUG_MIASM_AUTOMOD_CODE
 				fprintf(stderr, "**********************************\n");
 				fprintf(stderr, "self modifying code %"PRIX64" %"PRIX64"\n",
-					addr, my_size);
+					vm_mngr->memory_w.array[i].start,
+					vm_mngr->memory_w.array[i].stop);
 				fprintf(stderr, "**********************************\n");
 #endif
 				vm_mngr->exception_flags |= EXCEPT_CODE_AUTOMOD;
@@ -492,41 +487,25 @@ void check_invalid_code_blocs(vm_mngr_t* vm_mngr)
 void check_memory_breakpoint(vm_mngr_t* vm_mngr)
 {
 	int i;
-	int list_size;
-	PyObject* element;
-	uint64_t addr_start, addr_stop;
 	struct memory_breakpoint_info * memory_bp;
 
-	if (vm_mngr->write_num == 0)
-		return;
-
-	/* Check Write memory breakpoint */
-	list_size = PyList_Size(vm_mngr->memory_w);
+	/* Check memory breakpoints */
 	LIST_FOREACH(memory_bp, &vm_mngr->memory_breakpoint_pool, next) {
 		if (vm_mngr->exception_flags & EXCEPT_BREAKPOINT_INTERN)
 			break;
 		if (memory_bp->access & BREAKPOINT_READ) {
-			for (i=0;i<list_size; i++) {
-				element = PyList_GetItem(vm_mngr->memory_w, i);
-				addr_start = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 0));
-				addr_stop = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 1));
-
-
-				if ((memory_bp->ad < addr_stop) &&
-				    (addr_start < memory_bp->ad + memory_bp->size)) {
+			for (i=0;i<vm_mngr->memory_r.num; i++) {
+				if ((memory_bp->ad < vm_mngr->memory_r.array[i].stop) &&
+				    (vm_mngr->memory_r.array[i].start < memory_bp->ad + memory_bp->size)) {
 					vm_mngr->exception_flags |= EXCEPT_BREAKPOINT_INTERN;
 					break;
 				}
 			}
-		} else if (memory_bp->access & BREAKPOINT_READ) {
-			for (i=0;i<list_size; i++) {
-				element = PyList_GetItem(vm_mngr->memory_w, i);
-				addr_start = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 0));
-				addr_stop = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 1));
-
-
-				if ((memory_bp->ad < addr_stop) &&
-				    (addr_start < memory_bp->ad + memory_bp->size)) {
+		}
+		if (memory_bp->access & BREAKPOINT_WRITE) {
+			for (i=0;i<vm_mngr->memory_w.num; i++) {
+				if ((memory_bp->ad < vm_mngr->memory_w.array[i].stop) &&
+				    (vm_mngr->memory_w.array[i].start < memory_bp->ad + memory_bp->size)) {
 					vm_mngr->exception_flags |= EXCEPT_BREAKPOINT_INTERN;
 					break;
 				}
@@ -536,33 +515,30 @@ void check_memory_breakpoint(vm_mngr_t* vm_mngr)
 }
 
 
-
-void reset_pylist(PyObject* pylist)
+PyObject* get_memory_pylist(vm_mngr_t* vm_mngr, struct memory_access_list* memory_list)
 {
 	int i;
-	int list_size;
-	PyObject* element;
-
-	list_size = PyList_Size(pylist);
-
-	for (i=0;i<list_size; i++) {
-		element = PyList_GetItem(pylist, i);
-		Py_DECREF(element);
+	PyObject *pylist;
+	PyObject *range;
+	pylist = PyList_New(memory_list->num);
+	for (i=0;i<memory_list->num;i++) {
+		range = PyTuple_New(2);
+		PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)memory_list->array[i].start));
+		PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)memory_list->array[i].stop));
+		PyList_SetItem(pylist, i, range);
 	}
-
-	Py_DECREF(pylist);
+	return pylist;
 
 }
 
-
 PyObject* get_memory_read(vm_mngr_t* vm_mngr)
 {
-	return vm_mngr->memory_r;
+	return get_memory_pylist(vm_mngr, &vm_mngr->memory_r);
 }
 
 PyObject* get_memory_write(vm_mngr_t* vm_mngr)
 {
-	return vm_mngr->memory_w;
+	return get_memory_pylist(vm_mngr, &vm_mngr->memory_w);
 }
 
 PyObject* addr2BlocObj(vm_mngr_t* vm_mngr, uint64_t addr)
@@ -1542,9 +1518,8 @@ void init_code_bloc_pool(vm_mngr_t* vm_mngr)
 	vm_mngr->code_bloc_pool_ad_min = 0xffffffff;
 	vm_mngr->code_bloc_pool_ad_max = 0;
 
-	vm_mngr->memory_r = PyList_New(0);
-	vm_mngr->memory_w = PyList_New(0);
-	vm_mngr->write_num = 0;
+	memory_access_list_init(&(vm_mngr->memory_r));
+	memory_access_list_init(&(vm_mngr->memory_w));
 
 
 }
@@ -1586,13 +1561,8 @@ void reset_code_bloc_pool(vm_mngr_t* vm_mngr)
 
 void reset_memory_access(vm_mngr_t* vm_mngr)
 {
-	reset_pylist(vm_mngr->memory_r);
-	vm_mngr->memory_r = PyList_New(0);
-
-	reset_pylist(vm_mngr->memory_w);
-	vm_mngr->memory_w = PyList_New(0);
-	vm_mngr->write_num = 0;
-
+	memory_access_list_reset(&(vm_mngr->memory_r));
+	memory_access_list_reset(&(vm_mngr->memory_w));
 }
 
 void reset_memory_breakpoint(vm_mngr_t* vm_mngr)
diff --git a/miasm2/jitter/vm_mngr.h b/miasm2/jitter/vm_mngr.h
index fa264b27..13ec065a 100644
--- a/miasm2/jitter/vm_mngr.h
+++ b/miasm2/jitter/vm_mngr.h
@@ -73,7 +73,16 @@ struct memory_page_node {
 	char* name;
 };
 
+struct memory_access {
+	uint64_t start;
+	uint64_t stop;
+};
 
+struct memory_access_list {
+	struct memory_access *array;
+	uint64_t allocated;
+	uint64_t num;
+};
 
 typedef struct {
 	int sex;
@@ -91,8 +100,10 @@ typedef struct {
 	uint64_t exception_flags_new;
 	PyObject *addr2obj;
 
-	PyObject* memory_r;
-	PyObject* memory_w;
+
+	struct memory_access_list memory_r;
+	struct memory_access_list memory_w;
+
 
 	int write_num;
 
@@ -266,6 +277,11 @@ unsigned int rcr_rez_op(unsigned int size, unsigned int a, unsigned int b, unsig
 	    }
 
 
+void memory_access_list_init(struct memory_access_list * access);
+void memory_access_list_reset(struct memory_access_list * access);
+void memory_access_list_add(struct memory_access_list * access, uint64_t start, uint64_t stop);
+
+
 void hexdump(char* m, unsigned int l);
 
 struct code_bloc_node * create_code_bloc_node(uint64_t ad_start, uint64_t ad_stop);