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();

main.cpp

@@ -7,11 +7,12 @@ int main(int argc, char* argv[])
 {
     if (argc < 2)
     {
-        cerr << "usage: cfgrip [--subs-only] <binary>" << endl;
+        cerr << "usage: cfgrip [--subs-only] [--clean] <binary>" << endl;
         return 1;
     }
 
     bool subsOnly = false;
+    bool clean = false;
     string path;
 
     for (int i = 1; i < argc; i++)
@@ -19,13 +20,15 @@ int main(int argc, char* argv[])
         string arg = argv[i];
         if (arg == "--subs-only")
             subsOnly = true;
+        else if (arg == "--clean")
+            clean = true;
         else
             path = arg;
     }
 
     if (path.empty())
     {
-        cerr << "usage: cfgrip [--subs-only] <binary>" << endl;
+        cerr << "usage: cfgrip [--subs-only] [--clean] <binary>" << endl;
         return 1;
     }
 
@@ -75,6 +78,12 @@ int main(int argc, char* argv[])
     if (subsOnly) cout << "mode: subs-only" << endl;
     CFG cfg = builder.build();
 
+    if (clean)
+    {
+        CFGBuilder::clean(cfg);
+        cout << "clean: xrefs=" << cfg.xrefs.size() << " functions=" << cfg.functions.size() << endl;
+    }
+
     cout << "functions: " << cfg.functions.size() << endl;
     cout << "indirect targets: " << cfg.indirect_targets.size() << endl;
 

types.hpp

@@ -28,6 +28,12 @@ struct IndirectTarget
     bool resolved;
 };
 
+struct XrefEntry
+{
+    addr_t target;
+    vector<pair<addr_t, string>> callers;
+};
+
 struct Instruction
 {
     addr_t address;
@@ -35,6 +41,7 @@ struct Instruction
     string mnemonic;
     string operands;
     uint8_t size;
+    int64_t stack_offset = 0;
 };
 
 struct BasicBlock
@@ -62,6 +69,7 @@ struct CFG
     vector<Function> functions;
     vector<ImportEntry> imports;
     vector<IndirectTarget> indirect_targets;
+    vector<XrefEntry> xrefs;
 
     static string escapeJSON(const string& s)
     {
@@ -140,7 +148,8 @@ struct CFG
                     os << "              \"address\": \"0x" << hex << inst.address << "\",\n";
                     os << "              \"size\": " << dec << (int)inst.size << ",\n";
                     os << "              \"mnemonic\": \"" << escapeJSON(inst.mnemonic) << "\",\n";
-                    os << "              \"operands\": \"" << escapeJSON(inst.operands) << "\"\n";
+                    os << "              \"operands\": \"" << escapeJSON(inst.operands) << "\",\n";
+                    os << "              \"stack_offset\": " << dec << inst.stack_offset << "\n";
                     os << "            }";
                     if (k < b.instructions.size() - 1) os << ",";
                     os << "\n";
@@ -162,7 +171,34 @@ struct CFG
             if (i < functions.size() - 1) os << ",";
             os << "\n";
         }
-        os << "  ]\n";
+        if (!xrefs.empty())
+        {
+            os << "  ],\n";
+            os << "  \"xrefs\": [\n";
+            for (size_t i = 0; i < xrefs.size(); ++i) {
+                const auto& x = xrefs[i];
+                os << "    {\n";
+                os << "      \"target\": \"0x" << hex << x.target << "\",\n";
+                os << "      \"callers\": [\n";
+                for (size_t j = 0; j < x.callers.size(); ++j) {
+                    os << "        {\n";
+                    os << "          \"address\": \"0x" << hex << x.callers[j].first << "\",\n";
+                    os << "          \"type\": \"" << escapeJSON(x.callers[j].second) << "\"\n";
+                    os << "        }";
+                    if (j < x.callers.size() - 1) os << ",";
+                    os << "\n";
+                }
+                os << "      ]\n";
+                os << "    }";
+                if (i < xrefs.size() - 1) os << ",";
+                os << "\n";
+            }
+            os << "  ]\n";
+        }
+        else
+        {
+            os << "  ]\n";
+        }
         os << "}\n";
         return os.str();
     }