Merge pull request #757 from commial/fix/dse

Fix DSE and ensure API type for SymbolPool loc_key_to_*

3 files changed, 5 insertions, 8 deletions

diff --git a/example/symbol_exec/dse_strategies.py b/example/symbol_exec/dse_strategies.py
index a981853a..5a4be321 100644
--- a/example/symbol_exec/dse_strategies.py
+++ b/example/symbol_exec/dse_strategies.py
@@ -67,7 +67,7 @@ jitter.init_run(run_addr)
 # Init a DSE instance with a given strategy
 dse = DSEPathConstraint(machine, produce_solution=strategy)
 dse.attach(jitter)
-# Concretize everything exept the argument
+# Concretize everything except the argument
 dse.update_state_from_concrete()
 regs = jitter.ir_arch.arch.regs
 arg = ExprId("ARG", 32)
diff --git a/miasm2/analysis/dse.py b/miasm2/analysis/dse.py
index 0502ea42..6eaf1e91 100644
--- a/miasm2/analysis/dse.py
+++ b/miasm2/analysis/dse.py
@@ -63,7 +63,6 @@ from miasm2.expression.expression_helper import possible_values
 from miasm2.ir.translators import Translator
 from miasm2.analysis.expression_range import expr_range
 from miasm2.analysis.modularintervals import ModularIntervals
-from miasm2.core.asmblock import AsmBlockBad
 
 DriftInfo = namedtuple("DriftInfo", ["symbol", "computed", "expected"])
 
@@ -326,8 +325,7 @@ class DSEEngine(object):
 
             ## Update current state
             asm_block = self.mdis.dis_block(cur_addr)
-            if not isinstance(asm_block, AsmBlockBad):
-                self.ir_arch.add_block(asm_block)
+            self.ir_arch.add_block(asm_block)
             self.addr_to_cacheblocks[cur_addr] = dict(self.ir_arch.blocks)
 
         # Emulate the current instruction
@@ -351,14 +349,11 @@ class DSEEngine(object):
 
                 if not (isinstance(next_addr_concrete, ExprLoc) and
                         self.ir_arch.symbol_pool.loc_key_to_offset(
-                            next_addr_concrete
+                            next_addr_concrete.loc_key
                         ) is None):
                     # Not a lbl_gen, exit
                     break
 
-                if self.symb.ir_arch.get_block(cur_addr) is None:
-                    break
-
                 # Call handle with lbl_gen state
                 self.handle(next_addr_concrete)
                 cur_addr = next_addr_concrete
diff --git a/miasm2/core/asmblock.py b/miasm2/core/asmblock.py
index c2b7aa0a..9a0f8081 100644
--- a/miasm2/core/asmblock.py
+++ b/miasm2/core/asmblock.py
@@ -374,6 +374,7 @@ class AsmSymbolPool(object):
         Return offset of @loc_key, None otherwise.
         @loc_key: LocKey instance
         """
+        assert isinstance(loc_key, LocKey)
         return self._loc_key_to_offset.get(loc_key)
 
     def loc_key_to_name(self, loc_key):
@@ -381,6 +382,7 @@ class AsmSymbolPool(object):
         Return name of @loc_key.
         @loc_key: LocKey instance
         """
+        assert isinstance(loc_key, LocKey)
         return self._loc_key_to_name[loc_key]
 
     def add_location(self, name, offset=None):