Implements initial register analysis

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third-party/capstone"]
+	path = third-party/capstone
+	url = https://github.com/aquynh/capstone.git

librw/container.py

@@ -1,7 +1,7 @@
 from collections import defaultdict
 import struct
 
-from capstone import CS_OP_IMM, CS_OP_MEM, CS_GRP_JUMP
+from capstone import CS_OP_IMM, CS_OP_MEM, CS_GRP_JUMP, CS_OP_REG
 
 from . import disasm
 
@@ -130,12 +130,15 @@ def __init__(self, name, start, sz, bytes, bind="STB_LOCAL"):
         self.bbstarts = set()
         self.bind = bind
 
+        # Populated during symbolization.
+        # Invalidated by any instrumentation.
+        self.nexts = defaultdict(list)
+
         self.bbstarts.add(start)
 
     def disasm(self):
         assert not self.cache
         for decoded in disasm.disasm_bytes(self.bytes, self.start):
-            #if not decoded.mnemonic.startswith("nop"):
             self.cache.append(InstructionWrapper(decoded))
 
     def is_valid_instruction(self, address):
@@ -184,6 +187,16 @@ def __str__(self):
 
         return "\n".join(results)
 
+    def next_of(self, instruction_idx):
+        print(self.cache[instruction_idx], self.nexts[instruction_idx])
+        nexts = list()
+        for x in self.nexts[instruction_idx]:
+            if isinstance(x, str):
+                nexts.append(x)
+            else:
+                nexts.append(self.cache[x].address)
+        return nexts
+
 
 class InstructionWrapper():
     def __init__(self, instruction):
@@ -202,6 +215,14 @@ def get_mem_access_op(self):
                 return (op.mem, idx)
         return (None, None)
 
+    def reg_reads(self):
+        regs = self.cs.regs_access()[0]
+        return [self.cs.reg_name(x) for x in regs]
+
+    def reg_writes(self):
+        regs = self.cs.regs_access()[1]
+        return [self.cs.reg_name(x) for x in regs]
+
 
 class InstrumentedInstruction():
     def __init__(self, code, label=None, forinst=None):

librw/rw.py

@@ -146,12 +146,33 @@ def symbolize_text_section(self, container, context):
         self.symbolize_switch_tables(container, context)
 
     def symbolize_cf_transfer(self, container, context=None):
+        addr_to_idx = dict()
         for _, function in container.functions.items():
-            for instruction in function.cache:
+            for inst_idx, instruction in enumerate(function.cache):
+                addr_to_idx[instruction.address] = inst_idx
+
+        for _, function in container.functions.items():
+            for inst_idx, instruction in enumerate(function.cache):
+
                 if (CS_GRP_JUMP not in instruction.cs.groups
                         and CS_GRP_CALL not in instruction.cs.groups):
+                    # Simple, next is idx + 1
+                    if instruction.mnemonic.startswith('ret'):
+                        function.nexts[inst_idx].append("ret")
+                    else:
+                        function.nexts[inst_idx].append(inst_idx + 1)
                     continue
 
+                if (CS_GRP_JUMP in instruction.cs.groups and
+                        not instruction.mnemonic.startswith("jmp")):
+
+                    if inst_idx + 1 < len(function.cache):
+                        function.nexts[inst_idx].append(inst_idx + 1)
+                    else:
+                        function.nexts[inst_idx].append("undef")
+                elif CS_GRP_CALL in instruction.cs.groups:
+                    function.nexts[inst_idx].append("call")
+
                 if instruction.cs.operands[0].type == CS_OP_IMM:
                     target = instruction.cs.operands[0].imm
                     # Check if the target is in .text section.
@@ -176,6 +197,17 @@ def symbolize_cf_transfer(self, container, context=None):
                         else:
                             print("[x] Missed call target: %x" % (target))
 
+                    if target in addr_to_idx:
+                        idx = addr_to_idx[target]
+                        function.nexts[inst_idx].append(idx)
+                    else:
+                        function.nexts[inst_idx].append("undef")
+
+                elif CS_GRP_CALL in instruction.cs.groups:
+                    function.nexts[inst_idx].append("call")
+                else:
+                    function.nexts[inst_idx].append("undef")
+
     def symbolize_switch_tables(self, container, context):
         rodata = container.sections.get(".rodata", None)
         if not rodata:
@@ -343,6 +375,7 @@ def symbolize_data_sections(self, container, context=None):
 
 if __name__ == "__main__":
     from .loader import Loader
+    from .analysis import register
 
     argp = argparse.ArgumentParser()
 
@@ -369,4 +402,13 @@ def symbolize_data_sections(self, container, context=None):
 
     rw = Rewriter(loader.container, args.outfile)
     rw.symbolize()
-    rw.dump()
+
+    for f, func in loader.container.functions.items():
+        if func.name != "P7AllocTrace":
+            continue
+        ra = register.RegisterAnalysis()
+        ra.analyze_function(func)
+        print("===== FREE REGS:", func.name)
+        for instruction in func.cache:
+            print(instruction, ra.free_regs[instruction.address])
+    #rw.dump()

requirements.txt

@@ -1,3 +1,9 @@
-capstone==3.0.5
+archinfo==8.18.10.5
+bitstring==3.1.5
+capstone==4.0.0rc1
+cffi==1.11.5
+future==0.16.0
+pycparser==2.19
 pyelftools==0.25
+pyvex==8.18.10.5
 yapf==0.24.0

rwtools/asan/instrument.py

@@ -67,7 +67,7 @@ def _access2(self):
         rest = copy.copy(sp.MEM_REG_RESTORE)
         save = copy.copy(sp.MEM_REG_SAVE)
 
-        ac1[2] = "incl %(clob1_32)s"
+        ac1[2] = "incl {clob1_32}"
 
         return "\n".join(save + common + ac1 + rest)
 
@@ -84,7 +84,7 @@ def _access8(self):
         rest = copy.copy(sp.MEM_REG_RESTORE)
         save = copy.copy(sp.MEM_REG_SAVE)
 
-        common[3] = "cmpb $0, 2147450880(%(tgt)s)"
+        common[3] = "cmpb $0, 2147450880({tgt})"
         common[4] = common[5]
         common[5] = sp.MEM_LOAD_SZ[-1]
 
@@ -170,7 +170,7 @@ def instrument_mem_accesses(self):
 
                 enter_lbl = "%s_%x" % (sp.ASAN_MEM_ENTER, instruction.address)
 
-                iinstr = InstrumentedInstruction(instrument % (args),
+                iinstr = InstrumentedInstruction(instrument.format(**args),
                                                  enter_lbl, str(instruction))
 
                 # TODO: Replace original instruction for efficiency

rwtools/asan/snippets.py

@@ -37,34 +37,46 @@
 ]
 
 MEM_LOAD_COMMON = [
-    "\tleaq %(lexp)s, %(clob1)s",
-    "\tmovq %(clob1)s, %(tgt)s",
-    "\tshrq $3, %(tgt)s",
-    "\tmovb 2147450880(%(tgt)s), %(tgt_8)s",
-    "\ttestb %(tgt_8)s, %(tgt_8)s",
-    "\tje {0}_%(addr)x".format(ASAN_MEM_EXIT),
+    "\tleaq {lexp}, {clob1}",
+    "\tmovq {clob1}, {tgt}",
+    "\tshrq $3, {tgt}",
+    "\tmovb 2147450880({tgt}), {tgt_8}",
+    "\ttestb {tgt_8}, {tgt_8}",
+    "\tje {0}_{{addr}}".format(ASAN_MEM_EXIT),
 ]
 
 MEM_LOAD_SZ = [
-    "\tandl $7, %(clob1_32)s",
-    "\taddl $%(acsz_1)d, %(clob1_32)s",
-    "\tmovsbl %(tgt_8)s, %(tgt_32)s",
-    "\tcmpl %(tgt_32)s, %(clob1_32)s",
-    "\tjl {0}_%(addr)x".format(ASAN_MEM_EXIT),
-    "\tcallq __asan_report_load%(acsz)d@PLT",
+    "\tandl $7, {clob1_32}",
+    "\taddl ${acsz_1}, {clob1_32}",
+    "\tmovsbl {tgt_8}, {tgt_32}",
+    "\tcmpl {tgt_32}, {clob1_32}",
+    "\tjl {0}_{{addr}}".format(ASAN_MEM_EXIT),
+    "\tcallq __asan_report_load{acsz}@PLT",
 ]
 
 MEM_REG_SAVE = [
-    "\tpushf",
-    "\tpushq %(clob1)s",
-    "\tpushq %(tgt)s",
+    # Save Regs
+    "\tpushq {clob1}",
+    "\tpushq {tgt}",
+    # Save Flags
+    "\tpushq %rax",
+    "\tlahf",
+    "\tseto %al",
+    "\txchg %ah, %al",
+    "\txchg %rax, 0(%rsp)",
 ]
 
 MEM_REG_RESTORE = [
-    "{0}_%(addr)x:".format(ASAN_MEM_EXIT),
-    "\tpopq %(tgt)s",
-    "\tpopq %(clob1)s",
-    "\tpopf",
+    "{0}_{{addr}}:".format(ASAN_MEM_EXIT),
+    # Restore Flags
+    "\txchg %rax, 0(%rsp)",
+    "\txchg %ah, %al",
+    "\tadd $0x7f, %al",
+    "\tsahf",
+    "\tpopq %rax",
+    # Restore others
+    "\tpopq {tgt}",
+    "\tpopq {clob1}",
 ]
 
 STACK_POISON_BASE = [

setup.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+PATH=`pwd`
+
+python3 -m venv retro
+source retro/bin/activate
+pip install --upgrade pip
+pip install -r reqruirements.txt
+git submodule update --init --checkout third-party/capstone
+cd third-party/capstone
+make
+cd bindings/python/ && make && make install
+cd $PATH