Fix breakpoint handler exceptions

Author: dshikashio

pybag/dbgeng/breakpoints.py

@@ -24,45 +24,56 @@ def __getitem__(self, id):
 
     def __call__(self, *args):
         rawbp = DebugBreakpoint(args[0])
+        bpid = rawbp.GetId()
+
         try:
-            bpfn = self.__getitem__(rawbp.GetId())
+            (count,bpfn) = self.__getitem__(bpid)
         except KeyError:
+            count = 0
             bpfn = None
         if bpfn:
             ret = bpfn(rawbp)
         else:
             ret = DbgEng.DEBUG_STATUS_NO_CHANGE
 
-        if rawbp.GetFlags() & DbgEng.DEBUG_BREAKPOINT_ONE_SHOT:
-            self._bp.pop(rawbp.GetId(), None)
-            self._control.RemoveBreakpoint(rawbp)
-
+        count -= 1
+        if count == 0:
+            # We can disable here and but not delete
+            rawbp.RemoveFlags(DbgEng.DEBUG_BREAKPOINT_ENABLED)
         return ret
 
     def set(self, expr, handler=None, type=None, windbgcmd=None, oneshot=False,
-                  passcount=None, threadid=None, size=None, access=None):
-        # XXX - would be nice to check if bp with expr already exists
-        # if so, then enable existing, else add
-        if type is None:
-            type = DbgEng.DEBUG_BREAKPOINT_CODE
-        if isinstance(expr, int):
-            expr = b"0x%x" % expr
-        elif isinstance(expr, str):
-            expr = expr.encode()
-        bp = self._control.AddBreakpoint(type)
-        id = bp.GetId()
-        bp.SetOffsetExpression(expr)
-        if windbgcmd:
-            pass
+                  passcount=None, threadid=None, size=None, access=None,
+                  count=0xffffffff):
+
+        existingid = self.find(expr)
+        if existingid != -1:
+            # Just enable
+            self.enable(existingid)
+            return
+        else:
+            if type is None:
+                type = DbgEng.DEBUG_BREAKPOINT_CODE
+            if isinstance(expr, int):
+                expr = b"0x%x" % expr
+            elif isinstance(expr, str):
+                expr = expr.encode()
+            bp = self._control.AddBreakpoint(type)
+            id = bp.GetId()
+            bp.SetOffsetExpression(expr)
+            if windbgcmd:
+                pass
+            if type == DbgEng.DEBUG_BREAKPOINT_DATA:
+                bp.SetDataParameters(size, access)
         if passcount is not None:
             bp.SetPassCount(passcount)
         if threadid:
             bp.SetMatchThreadId(threadid)
-        if type == DbgEng.DEBUG_BREAKPOINT_DATA:
-            bp.SetDataParameters(size, access)
         if oneshot:
-            bp.AddFlags(DbgEng.DEBUG_BREAKPOINT_ONE_SHOT)
-        self._bp[id] = handler
+            count = 1
+            # comtypes can't properly handle true one shot
+            #bp.AddFlags(DbgEng.DEBUG_BREAKPOINT_ONE_SHOT)
+        self._bp[id] = (count,handler)
         if not bp.GetFlags() & DbgEng.DEBUG_BREAKPOINT_DEFERRED:
             self.enable(id)
         return id
@@ -98,3 +109,45 @@ def remove_all(self):
         for bpid in self:
             self.remove(bpid)
 
+    def find_id(self, fid):
+        for bpid in self:
+            if bpid == fid:
+                return fid
+        return -1
+
+    def find_expr(self, expr):
+        ids = [bpid for bpid in self]
+        for bpid in ids:
+            try:
+                bp = self._control.GetBreakpointById(bpid)
+            except exception.E_NOINTERFACE_Error:
+                self.breakpoints._remove_stale(bpid)
+                continue
+
+            bpexpr = bp.GetOffsetExpression()
+            print(bpexpr)
+            #expr = self.get_name_by_offset(bp.GetOffset())
+            if bpexpr == expr:
+                return bpid
+        return -1
+
+    def find_offset(self, offset):
+        ids = [bpid for bpid in self]
+        for bpid in ids:
+            try:
+                bp = self._control.GetBreakpointById(bpid)
+            except exception.E_NOINTERFACE_Error:
+                self.breakpoints._remove_stale(bpid)
+                continue
+
+            bpoff = bp.GetOffset()
+            if bpoff == offset:
+                return bpid
+        return -1
+
+
+    def find(self, expr):
+        bpid = self.find_expr(expr)
+        if bpid == -1:
+            bpid = self.find_offset(expr)
+        return bpid

pybag/dbgeng/idebugcontrol.py

@@ -6,6 +6,9 @@
 from .idebugbreakpoint import DebugBreakpoint
 from .util import logger
 
+from comtypes.gen.DbgEng import IDebugBreakpoint2
+from comtypes import _compointer_base
+
 
 class DebugControl(object):
     def __init__(self, controlobj):
@@ -392,7 +395,15 @@ def AddBreakpoint(self, type=DbgEng.DEBUG_BREAKPOINT_CODE):
         return DebugBreakpoint(bp)
 
     def RemoveBreakpoint(self, bp):
-        hr = self._ctrl.RemoveBreakpoint(bp._bp)
+        # We have to let comtypes get out of the way
+        # So we get the raw pointer value, release the object, and call with that
+        tmp = super(_compointer_base, bp._bp).value
+        del bp._bp
+        import gc
+        gc.collect()
+        #print(f"{tmp=:#016x}")
+        hr = self._ctrl.RemoveBreakpoint(cast(tmp, POINTER(IDebugBreakpoint2)))
+        #hr = self._ctrl.RemoveBreakpoint(bp._bp)
         exception.check_err(hr)
 
     def AddExtension(self):

pybag/pydbg.py

@@ -256,7 +256,7 @@ def dispatch_events(self):
     def _worker_wait(self, msg, timeout=DbgEng.WAIT_INFINITE, args=None):
         if timeout == -1:
             timeout = 0xffffffff
-        elif timeout != 0xffffffff:
+        if timeout != 0xffffffff:
             timeout *= 1000
 
         item = WorkItem(msg, timeout, args)