add --clean flag: jump-threading, dead-block pruning, stack deltas, xrefs

cfg/builder.cpp

@@ -534,3 +534,175 @@ CFG CFGBuilder::build()
     cfg.indirect_targets = m_targets;
     return cfg;
 }
+
+static int64_t stackDelta(const Instruction& inst)
+{
+    if (inst.mnemonic == "push") return -8;
+    if (inst.mnemonic == "pop") return 8;
+    if (inst.mnemonic == "call") return 0;
+    if (inst.mnemonic == "ret" || inst.mnemonic == "retf") return 8;
+    if (inst.mnemonic == "sub")
+    {
+        size_t p = inst.operands.find("rsp");
+        if (p == string::npos) p = inst.operands.find("esp");
+        if (p != string::npos && (p == 0 || inst.operands[p-1] == ' ' || inst.operands[p-1] == ','))
+        {
+            size_t x = inst.operands.rfind("0x");
+            if (x != string::npos && x > p)
+            {
+                try { return -(int64_t)stoull(inst.operands.substr(x), nullptr, 16); }
+                catch (...) {}
+            }
+        }
+    }
+    if (inst.mnemonic == "add")
+    {
+        size_t p = inst.operands.find("rsp");
+        if (p == string::npos) p = inst.operands.find("esp");
+        if (p != string::npos && (p == 0 || inst.operands[p-1] == ' ' || inst.operands[p-1] == ','))
+        {
+            size_t x = inst.operands.rfind("0x");
+            if (x != string::npos && x > p)
+            {
+                try { return (int64_t)stoull(inst.operands.substr(x), nullptr, 16); }
+                catch (...) {}
+            }
+        }
+    }
+    return 0;
+}
+
+static void computeStackDeltas(CFG& cfg)
+{
+    for (auto& func : cfg.functions)
+    {
+        for (auto& block : func.blocks)
+        {
+            int64_t delta = 0;
+            for (auto& inst : block.instructions)
+            {
+                inst.stack_offset = delta;
+                delta += stackDelta(inst);
+            }
+        }
+    }
+}
+
+static void jumpThread(CFG& cfg)
+{
+    map<addr_t, BasicBlock*> addr_to_block;
+    map<addr_t, Function*> addr_to_func;
+    for (auto& func : cfg.functions)
+    {
+        for (auto& block : func.blocks)
+        {
+            addr_to_block[block.address] = &block;
+            addr_to_func[block.address] = &func;
+        }
+    }
+
+    for (auto& func : cfg.functions)
+    {
+        bool changed = true;
+        while (changed)
+        {
+            changed = false;
+            for (auto& block : func.blocks)
+            {
+                if (block.successors.size() != 1) continue;
+                addr_t t = block.successors[0];
+                auto it = addr_to_block.find(t);
+                if (it == addr_to_block.end()) continue;
+                if (addr_to_func[t] != &func) continue;
+
+                BasicBlock* target = it->second;
+                if (target->instructions.size() == 1 &&
+                    target->instructions[0].mnemonic == "jmp" &&
+                    target->successors.size() == 1)
+                {
+                    addr_t redirect = target->successors[0];
+                    if (redirect != block.address && redirect != t)
+                    {
+                        block.successors[0] = redirect;
+                        changed = true;
+                    }
+                }
+            }
+        }
+    }
+}
+
+static void pruneDeadBlocks(CFG& cfg)
+{
+    for (auto& func : cfg.functions)
+    {
+        if (func.blocks.empty()) continue;
+        set<addr_t> reachable;
+        queue<addr_t> q;
+        reachable.insert(func.blocks[0].address);
+        q.push(func.blocks[0].address);
+
+        map<addr_t, BasicBlock*> addr_map;
+        for (auto& b : func.blocks) addr_map[b.address] = &b;
+
+        while (!q.empty())
+        {
+            addr_t cur = q.front(); q.pop();
+            auto it = addr_map.find(cur);
+            if (it == addr_map.end()) continue;
+            for (auto s : it->second->successors)
+                if (reachable.insert(s).second) q.push(s);
+        }
+
+        vector<BasicBlock> alive;
+        for (auto& b : func.blocks)
+            if (reachable.count(b.address))
+                alive.push_back(move(b));
+        func.blocks = move(alive);
+    }
+}
+
+static void computeXrefs(CFG& cfg)
+{
+    map<addr_t, vector<pair<addr_t, string>>> xmap;
+    for (const auto& func : cfg.functions)
+    {
+        for (const auto& block : func.blocks)
+        {
+            for (const auto& inst : block.instructions)
+            {
+                if (inst.mnemonic == "call" || inst.mnemonic == "jmp" ||
+                    (inst.mnemonic.size() == 2 && inst.mnemonic[0] == 'j' && inst.mnemonic != "jn"))
+                {
+                    size_t pos = inst.operands.find("0x");
+                    if (pos != string::npos)
+                    {
+                        try
+                        {
+                            addr_t target = stoull(inst.operands.substr(pos), nullptr, 16);
+                            string type = inst.mnemonic;
+                            xmap[target].push_back({inst.address, type});
+                        }
+                        catch (...) {}
+                    }
+                }
+            }
+        }
+    }
+
+    for (auto& kv : xmap)
+    {
+        XrefEntry xe;
+        xe.target = kv.first;
+        xe.callers = move(kv.second);
+        cfg.xrefs.push_back(move(xe));
+    }
+}
+
+void CFGBuilder::clean(CFG& cfg)
+{
+    jumpThread(cfg);
+    pruneDeadBlocks(cfg);
+    computeStackDeltas(cfg);
+    computeXrefs(cfg);
+}

cfg/builder.hpp

@@ -11,6 +11,7 @@ class CFGBuilder
     public:
         CFGBuilder(Binary* binary, Disassembler* disasm, bool subsOnly = false);
         CFG build();
+        static void clean(CFG& cfg);
 
     private:
         void buildGotMap();