IRGenerator::explore()

yegord · May 13, 2015 · 7f1e836 · 7f1e836
1 parent f87528b
commit 7f1e836
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 0 deletions.
diff --git a/src/nc/core/irgen/IRGenerator.cpp b/src/nc/core/irgen/IRGenerator.cpp
@@ -24,6 +24,10 @@
 
 #include "IRGenerator.h"
 
+#include <queue>
+
+#include <boost/unordered_set.hpp>
+
 #include <nc/common/Foreach.h>
 #include <nc/common/Range.h>
 #include <nc/common/Warnings.h>
@@ -243,6 +247,102 @@ void IRGenerator::addJumpToDirectSuccessor(ir::BasicBlock *basicBlock) {
     }
 }
 
+std::unique_ptr<arch::Instructions> IRGenerator::explore(const image::Image *image, ByteAddr startAddress, bool followCalls, const CancellationToken &canceled) {
+    assert(image);
+
+    auto instructions = std::make_unique<arch::Instructions>();
+    ir::Program program;
+
+    auto disassembler = image->architecture()->createDisassembler();
+    auto instructionAnalyzer = image->architecture()->createInstructionAnalyzer();
+
+    auto exploreAddress = [&](ByteAddr address){
+        if (instructions->get(address)) {
+            return;
+        }
+
+        if (auto instruction = disassembler->disassembleSingleInstruction(address, image)) {
+            try {
+                instructionAnalyzer->createStatements(instruction.get(), &program);
+            } catch (const InvalidInstructionException &e) {
+                ncWarning(e.unicodeWhat());
+            }
+            instructions->add(std::move(instruction));
+        }
+    };
+
+    exploreAddress(startAddress);
+
+    if (program.basicBlocks().empty()) {
+        return instructions;
+    }
+
+    IRGenerator generator(image, instructions.get(), &program);
+
+    auto exploreBasicBlock = [&](ir::BasicBlock *basicBlock){
+        while (basicBlock->successorAddress() && !basicBlock->getTerminator()) {
+            auto lastSuccessorAddress = basicBlock->successorAddress();
+            exploreAddress(*basicBlock->successorAddress());
+            if (basicBlock->successorAddress() == lastSuccessorAddress) {
+                break;
+            }
+            canceled.poll();
+        }
+
+        generator.computeJumpTargets(basicBlock);
+    };
+
+    std::queue<ir::BasicBlock *> queue;
+    boost::unordered_set<ir::BasicBlock *> enqueued;
+
+    auto enqueue = [&](ir::BasicBlock *basicBlock) {
+        if (!nc::contains(enqueued, basicBlock)) {
+            queue.push(basicBlock);
+            enqueued.insert(basicBlock);
+        }
+    };
+
+    auto enqueueJumpTarget = [&](const ir::JumpTarget &target) {
+        if (target.basicBlock()) {
+            enqueue(target.basicBlock());
+        } else if (target.table()) {
+            foreach(const auto &entry, *target.table()) {
+                if (entry.basicBlock()) {
+                    enqueue(entry.basicBlock());
+                }
+            }
+        }
+    };
+
+    auto enqueueJumpTargets = [&](const ir::Jump *jump){
+        enqueueJumpTarget(jump->thenTarget());
+        enqueueJumpTarget(jump->elseTarget());
+    };
+
+    enqueue(program.basicBlocks().front());
+
+    while (!queue.empty()) {
+        auto basicBlock = queue.front();
+        queue.pop();
+
+        exploreBasicBlock(basicBlock);
+
+        if (auto jump = basicBlock->getJump()) {
+            enqueueJumpTargets(jump);
+        }
+
+        if (followCalls && queue.empty()) {
+            foreach (auto address, program.calledAddresses()) {
+                enqueue(program.createBasicBlock(address));
+            }
+        }
+
+        canceled.poll();
+    }
+
+    return instructions;
+}
+
 } // namespace irgen
 } // namespace core
 } // namespace nc

diff --git a/src/nc/core/irgen/IRGenerator.h b/src/nc/core/irgen/IRGenerator.h
@@ -83,9 +83,23 @@ class IRGenerator {
     /**
      * Builds a program control flow graph from the instructions
      * given to the constructor.
+     *
+     * \param[in] canceled Cancellation token.
      */
     void generate(const CancellationToken &canceled);
 
+    /**
+     * Tries to find all instructions reachable from the given one.
+     *
+     * \param[in] image Valid pointer to the executable image.
+     * \param[in] startAddress Virtual address of the instruction to start from.
+     * \param[in] followCalls Pass true to follow calls, pass false to stay within one function.
+     * \param[in] canceled Cancellation token.
+     *
+     * \return Valid pointer to the set of instructions discovered.
+     */
+    static std::unique_ptr<arch::Instructions> explore(const image::Image *image, ByteAddr startAddress, bool followCalls, const CancellationToken &canceled);
+
 private:
     /**
      * Computes jump targets in the basic block.