Jitter: use mem read/write for memory breakpoint/automod

6 files changed, 162 insertions, 87 deletions

diff --git a/miasm2/jitter/arch/JitCore_x86.c b/miasm2/jitter/arch/JitCore_x86.c
index 6a503d83..21d0e04c 100644
--- a/miasm2/jitter/arch/JitCore_x86.c
+++ b/miasm2/jitter/arch/JitCore_x86.c
@@ -322,13 +322,6 @@ IMOD(16)
 IMOD(32)
 IMOD(64)
 
-void check_automod(JitCpu* jitcpu, uint64_t addr, uint64_t size)
-{
-	if (!(((VmMngr*)jitcpu->pyvm)->vm_mngr.exception_flags & EXCEPT_CODE_AUTOMOD))
-		return;
-	code_bloc_add_write(&((VmMngr*)jitcpu->pyvm)->vm_mngr, addr, size/8);
-}
-
 void MEM_WRITE_08(JitCpu* jitcpu, uint64_t addr, uint8_t src)
 {
 	vm_MEM_WRITE_08(&((VmMngr*)jitcpu->pyvm)->vm_mngr, addr, src);
@@ -376,7 +369,6 @@ PyObject* vm_set_mem(JitCpu *self, PyObject* args)
        ret = vm_write_mem(&(((VmMngr*)self->pyvm)->vm_mngr), addr, buffer, size);
        if (ret < 0)
 	       RAISE(PyExc_TypeError,"arg must be str");
-       check_automod(self, addr, size*8);
 
        Py_INCREF(Py_None);
        return Py_None;
diff --git a/miasm2/jitter/codegen.py b/miasm2/jitter/codegen.py
index 068839e7..4f688896 100644
--- a/miasm2/jitter/codegen.py
+++ b/miasm2/jitter/codegen.py
@@ -89,6 +89,8 @@ class CGen(object):
     """
 
     CODE_VM_EXCEPTION_POST_INSTR = r"""
+    check_memory_breakpoint(&((VmMngr*)jitcpu->pyvm)->vm_mngr);
+    check_invalid_code_blocs(&((VmMngr*)jitcpu->pyvm)->vm_mngr);
     if (VM_exception_flag) {
         %s = %s;
         BlockDst->address = DST_value;
@@ -332,7 +334,7 @@ class CGen(object):
             out += (self.CODE_CPU_EXCEPTION_POST_INSTR % (self.C_PC, dst)).split('\n')
 
         if attrib.mem_read | attrib.mem_write:
-            out.append("reset_code_bloc_write(&((VmMngr*)jitcpu->pyvm)->vm_mngr);")
+            out.append("reset_memory_access(&((VmMngr*)jitcpu->pyvm)->vm_mngr);")
 
         return out
 
diff --git a/miasm2/jitter/jitcore.py b/miasm2/jitter/jitcore.py
index e81acc4e..2e3411c0 100644
--- a/miasm2/jitter/jitcore.py
+++ b/miasm2/jitter/jitcore.py
@@ -265,8 +265,7 @@ class JitCore(object):
         """Remove code jitted in range self.addr_mod
         @vm: VmMngr instance
         """
-        for addr_start, addr_stop in vm.get_code_bloc_write():
+        for addr_start, addr_stop in vm.get_memory_write():
             self.del_bloc_in_range(addr_start, addr_stop)
         self.__updt_jitcode_mem_range(vm)
         self.addr_mod = interval()
-        vm.reset_code_bloc_write()
diff --git a/miasm2/jitter/vm_mngr.c b/miasm2/jitter/vm_mngr.c
index 2b0ae1fd..208343fb 100644
--- a/miasm2/jitter/vm_mngr.c
+++ b/miasm2/jitter/vm_mngr.c
@@ -393,67 +393,88 @@ void dump_code_bloc(vm_mngr_t* vm_mngr)
 
 }
 
-void code_bloc_add_write(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t size)
+void add_range_to_pylist(PyObject* pylist, uint64_t addr1, uint64_t addr2)
 {
 	PyObject* range;
 	PyObject* element;
 	int list_size;
 	uint64_t addr_start, addr_stop;
 
-	list_size = PyList_Size(vm_mngr->code_bloc_memory_w);
+	list_size = PyList_Size(pylist);
 
 	if (list_size > 0) {
-		/* check match on upper bound */
-		 element = PyList_GetItem(vm_mngr->code_bloc_memory_w, list_size - 1);
+		/* 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 == addr) {
+		 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)addr+size));
-			 PyList_SetItem(vm_mngr->code_bloc_memory_w, list_size - 1, range);
+			 PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr2));
+			 PyList_SetItem(pylist, list_size - 1, range);
 			 return;
 
 		 }
 
-		/* check match on lower bound */
-		 element = PyList_GetItem(vm_mngr->code_bloc_memory_w, 0);
+		/* Check match on lower bound */
+		 element = PyList_GetItem(pylist, 0);
 		 addr_start = (uint64_t)PyLong_AsUnsignedLongLong(PyTuple_GetItem(element, 0));
 
-		 if (addr_start == addr + size) {
+		 if (addr_start == addr2) {
 			 range = PyTuple_New(2);
-			 PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr));
+			 PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr1));
 			 PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr_start));
-			 PyList_SetItem(vm_mngr->code_bloc_memory_w, 0, range);
+			 PyList_SetItem(pylist, 0, range);
 			 return;
 		 }
 
 	}
+
+	/* No merge, add to the list */
 	range = PyTuple_New(2);
-	PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr));
-	PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr+size));
+	PyTuple_SetItem(range, 0, PyLong_FromUnsignedLongLong((uint64_t)addr1));
+	PyTuple_SetItem(range, 1, PyLong_FromUnsignedLongLong((uint64_t)addr2));
+
+	PyList_Append(pylist, range);
+}
+
+
+void vm_mngr_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);
+}
 
-	PyList_Append(vm_mngr->code_bloc_memory_w, range);
+void vm_mngr_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);
 }
 
-void check_write_code_bloc(vm_mngr_t* vm_mngr, uint64_t my_size, uint64_t addr)
+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->exception_flags & EXCEPT_CODE_AUTOMOD)
-		return;
+	list_size = PyList_Size(vm_mngr->memory_w);
+
+	LIST_FOREACH(cbp, &vm_mngr->code_bloc_pool, next){
+		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 (!(addr + my_size/8 <= vm_mngr->code_bloc_pool_ad_min ||
-	      addr >=vm_mngr->code_bloc_pool_ad_max)){
-		LIST_FOREACH(cbp, &vm_mngr->code_bloc_pool, next){
-			if ((cbp->ad_start < addr + my_size/8) &&
-			    (addr < cbp->ad_stop)){
+			if ((cbp->ad_start < addr_stop) &&
+			    (addr_start < cbp->ad_stop)){
 #ifdef DEBUG_MIASM_AUTOMOD_CODE
 				fprintf(stderr, "**********************************\n");
 				fprintf(stderr, "self modifying code %"PRIX64" %"PRIX64"\n",
-				       addr, my_size);
+					addr, my_size);
 				fprintf(stderr, "**********************************\n");
 #endif
 				vm_mngr->exception_flags |= EXCEPT_CODE_AUTOMOD;
@@ -463,27 +484,78 @@ void check_write_code_bloc(vm_mngr_t* vm_mngr, uint64_t my_size, uint64_t addr)
 	}
 }
 
-void reset_code_bloc_write(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;
 
-	list_size = PyList_Size(vm_mngr->code_bloc_memory_w);
+	/* Check Write memory breakpoint */
+	list_size = PyList_Size(vm_mngr->memory_w);
+	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)) {
+					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)) {
+					vm_mngr->exception_flags |= EXCEPT_BREAKPOINT_INTERN;
+					break;
+				}
+			}
+		}
+	}
+}
+
+
+
+void reset_pylist(PyObject* pylist)
+{
+	int i;
+	int list_size;
+	PyObject* element;
+
+	list_size = PyList_Size(pylist);
 
 	for (i=0;i<list_size; i++) {
-		element = PyList_GetItem(vm_mngr->code_bloc_memory_w, i);
+		element = PyList_GetItem(pylist, i);
 		Py_DECREF(element);
 	}
 
-	Py_DECREF(vm_mngr->code_bloc_memory_w);
-	vm_mngr->code_bloc_memory_w = PyList_New(0);
+	Py_DECREF(pylist);
+
+}
+
 
+PyObject* get_memory_read(vm_mngr_t* vm_mngr)
+{
+	return vm_mngr->memory_r;
 }
 
-PyObject* get_code_bloc_write(vm_mngr_t* vm_mngr)
+PyObject* get_memory_write(vm_mngr_t* vm_mngr)
 {
-	return vm_mngr->code_bloc_memory_w;
+	return vm_mngr->memory_w;
 }
 
 PyObject* addr2BlocObj(vm_mngr_t* vm_mngr, uint64_t addr)
@@ -505,53 +577,53 @@ PyObject* addr2BlocObj(vm_mngr_t* vm_mngr, uint64_t addr)
 
 void vm_MEM_WRITE_08(vm_mngr_t* vm_mngr, uint64_t addr, unsigned char src)
 {
-	check_write_code_bloc(vm_mngr, 8, addr);
-	code_bloc_add_write(vm_mngr, addr, 1);
+	vm_mngr_add_mem_write(vm_mngr, addr, 1);
 	memory_page_write(vm_mngr, 8, addr, src);
 }
 
 void vm_MEM_WRITE_16(vm_mngr_t* vm_mngr, uint64_t addr, unsigned short src)
 {
-	check_write_code_bloc(vm_mngr, 16, addr);
-	code_bloc_add_write(vm_mngr, addr, 2);
+	vm_mngr_add_mem_write(vm_mngr, addr, 2);
 	memory_page_write(vm_mngr, 16, addr, src);
 }
 void vm_MEM_WRITE_32(vm_mngr_t* vm_mngr, uint64_t addr, unsigned int src)
 {
-	check_write_code_bloc(vm_mngr, 32, addr);
-	code_bloc_add_write(vm_mngr, addr, 4);
+	vm_mngr_add_mem_write(vm_mngr, addr, 4);
 	memory_page_write(vm_mngr, 32, addr, src);
 }
 void vm_MEM_WRITE_64(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t src)
 {
-	check_write_code_bloc(vm_mngr, 64, addr);
-	code_bloc_add_write(vm_mngr, addr, 8);
+	vm_mngr_add_mem_write(vm_mngr, addr, 8);
 	memory_page_write(vm_mngr, 64, addr, src);
 }
 
 unsigned char vm_MEM_LOOKUP_08(vm_mngr_t* vm_mngr, uint64_t addr)
 {
-    unsigned char ret;
-    ret = memory_page_read(vm_mngr, 8, addr);
-    return ret;
+	unsigned char ret;
+	vm_mngr_add_mem_read(vm_mngr, addr, 1);
+	ret = memory_page_read(vm_mngr, 8, addr);
+	return ret;
 }
 unsigned short vm_MEM_LOOKUP_16(vm_mngr_t* vm_mngr, uint64_t addr)
 {
-    unsigned short ret;
-    ret = memory_page_read(vm_mngr, 16, addr);
-    return ret;
+	unsigned short ret;
+	vm_mngr_add_mem_read(vm_mngr, addr, 2);
+	ret = memory_page_read(vm_mngr, 16, addr);
+	return ret;
 }
 unsigned int vm_MEM_LOOKUP_32(vm_mngr_t* vm_mngr, uint64_t addr)
 {
-    unsigned int ret;
-    ret = memory_page_read(vm_mngr, 32, addr);
-    return ret;
+	unsigned int ret;
+	vm_mngr_add_mem_read(vm_mngr, addr, 4);
+	ret = memory_page_read(vm_mngr, 32, addr);
+	return ret;
 }
 uint64_t vm_MEM_LOOKUP_64(vm_mngr_t* vm_mngr, uint64_t addr)
 {
-    uint64_t ret;
-    ret = memory_page_read(vm_mngr, 64, addr);
-    return ret;
+	uint64_t ret;
+	vm_mngr_add_mem_read(vm_mngr, addr, 8);
+	ret = memory_page_read(vm_mngr, 64, addr);
+	return ret;
 }
 
 
@@ -592,8 +664,6 @@ int vm_write_mem(vm_mngr_t* vm_mngr, uint64_t addr, char *buffer, uint64_t size)
        uint64_t len;
        struct memory_page_node * mpn;
 
-       check_write_code_bloc(vm_mngr, size * 8, addr);
-
        /* write is multiple page wide */
        while (size){
 	      mpn = get_memory_page_from_address(vm_mngr, addr, 1);
@@ -1465,7 +1535,9 @@ 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->code_bloc_memory_w = PyList_New(0);
+	vm_mngr->memory_r = PyList_New(0);
+	vm_mngr->memory_w = PyList_New(0);
+
 
 }
 
@@ -1504,6 +1576,14 @@ void reset_code_bloc_pool(vm_mngr_t* vm_mngr)
 	vm_mngr->code_bloc_pool_ad_max = 0;
 }
 
+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);
+}
 
 void reset_memory_breakpoint(vm_mngr_t* vm_mngr)
 {
@@ -1517,6 +1597,8 @@ void reset_memory_breakpoint(vm_mngr_t* vm_mngr)
 
 }
 
+
+
 /* We don't use dichotomy here for the insertion */
 int is_mpn_in_tab(vm_mngr_t* vm_mngr, struct memory_page_node* mpn_a)
 {
diff --git a/miasm2/jitter/vm_mngr.h b/miasm2/jitter/vm_mngr.h
index 67f0bac2..9267a462 100644
--- a/miasm2/jitter/vm_mngr.h
+++ b/miasm2/jitter/vm_mngr.h
@@ -91,10 +91,8 @@ typedef struct {
 	uint64_t exception_flags_new;
 	PyObject *addr2obj;
 
-
-
-	PyObject* code_bloc_memory_w;
-
+	PyObject* memory_r;
+	PyObject* memory_w;
 
 }vm_mngr_t;
 
@@ -287,10 +285,13 @@ void remove_memory_breakpoint(vm_mngr_t* vm_mngr, uint64_t ad, unsigned int acce
 
 void add_memory_page(vm_mngr_t* vm_mngr, struct memory_page_node* mpn);
 
-void check_write_code_bloc(vm_mngr_t* vm_mngr, uint64_t my_size, uint64_t addr);
-void code_bloc_add_write(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t size);
-void reset_code_bloc_write(vm_mngr_t* vm_mngr);
-PyObject* get_code_bloc_write(vm_mngr_t* vm_mngr);
+void vm_mngr_add_mem_read(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t size);
+void vm_mngr_add_mem_write(vm_mngr_t* vm_mngr, uint64_t addr, uint64_t size);
+void check_invalid_code_blocs(vm_mngr_t* vm_mngr);
+void check_memory_breakpoint(vm_mngr_t* vm_mngr);
+void reset_memory_access(vm_mngr_t* vm_mngr);
+PyObject* get_memory_read(vm_mngr_t* vm_mngr);
+PyObject* get_memory_write(vm_mngr_t* vm_mngr);
 
 
 char* dump(vm_mngr_t* vm_mngr);
diff --git a/miasm2/jitter/vm_mngr_py.c b/miasm2/jitter/vm_mngr_py.c
index 9e370465..3e0c820c 100644
--- a/miasm2/jitter/vm_mngr_py.c
+++ b/miasm2/jitter/vm_mngr_py.c
@@ -186,7 +186,8 @@ PyObject* vm_set_mem(VmMngr* self, PyObject* args)
        if (ret < 0)
 	      RAISE(PyExc_TypeError, "Error in set_mem");
 
-       check_write_code_bloc(&self->vm_mngr, size*8, addr);
+       vm_mngr_add_mem_write(&self->vm_mngr, addr, size);
+       check_invalid_code_blocs(&self->vm_mngr);
 
        Py_INCREF(Py_None);
        return Py_None;
@@ -460,26 +461,24 @@ PyObject* vm_is_mapped(VmMngr* self, PyObject* args)
 	return PyLong_FromUnsignedLongLong((uint64_t)ret);
 }
 
-PyObject* vm_reset_code_bloc_write(VmMngr* self, PyObject* args)
+PyObject* vm_get_memory_read(VmMngr* self, PyObject* args)
 {
-	reset_code_bloc_write(&self->vm_mngr);
-	Py_INCREF(Py_None);
-	return Py_None;
+	PyObject* result;
+	result = get_memory_read(&self->vm_mngr);
+	Py_INCREF(result);
+	return result;
 }
 
-PyObject* vm_get_code_bloc_write(VmMngr* self, PyObject* args)
+PyObject* vm_get_memory_write(VmMngr* self, PyObject* args)
 {
 	PyObject* result;
-
-	result = get_code_bloc_write(&self->vm_mngr);
+	result = get_memory_write(&self->vm_mngr);
 	Py_INCREF(result);
-
 	return result;
 }
 
 
 
-
 static PyObject *
 vm_set_big_endian(VmMngr *self, PyObject *value, void *closure)
 {
@@ -548,8 +547,6 @@ static PyMethodDef VmMngr_methods[] = {
 	 "X"},
 	{"is_mapped", (PyCFunction)vm_is_mapped, METH_VARARGS,
 	 "X"},
-	{"reset_code_bloc_write", (PyCFunction)vm_reset_code_bloc_write, METH_VARARGS,
-	 "X"},
 	{"add_code_bloc",(PyCFunction)vm_add_code_bloc, METH_VARARGS,
 	 "X"},
 	{"get_mem", (PyCFunction)vm_get_mem, METH_VARARGS,
@@ -572,8 +569,6 @@ static PyMethodDef VmMngr_methods[] = {
 	 "X"},
 	{"reset_code_bloc_pool", (PyCFunction)vm_reset_code_bloc_pool, METH_VARARGS,
 	 "X"},
-	{"get_code_bloc_write", (PyCFunction)vm_get_code_bloc_write, METH_VARARGS,
-	 "X"},
 	{"set_alarm", (PyCFunction)set_alarm, METH_VARARGS,
 	 "X"},
 	{"get_exception",(PyCFunction)vm_get_exception, METH_VARARGS,
@@ -584,6 +579,10 @@ static PyMethodDef VmMngr_methods[] = {
 	 "X"},
 	{"set_little_endian",(PyCFunction)vm_set_little_endian, METH_VARARGS,
 	 "X"},
+	{"get_memory_read",(PyCFunction)vm_get_memory_read, METH_VARARGS,
+	 "X"},
+	{"get_memory_write",(PyCFunction)vm_get_memory_write, METH_VARARGS,
+	 "X"},
 
 	{NULL}  /* Sentinel */
 };