Add interface to iterate over functions

[frida/frida.git] / src / disassembler / llvm / LLVMDisassembler.cxx

diff --git a/src/disassembler/llvm/LLVMDisassembler.cxx b/src/disassembler/llvm/LLVMDisassembler.cxx
index e1570758da27831e02d1ca63fe85fa326297bda1..3494b79bd8a370814d35eba28fec53db0e1543d6 100644 (file)
--- a/src/disassembler/llvm/LLVMDisassembler.cxx
+++ b/src/disassembler/llvm/LLVMDisassembler.cxx
@@ -13,7 +13,6 @@ using namespace llvm::object;
  * ist sondern z.B. einfach nur Instruktionen oder ein Bootsektor oder
  * foo
  */
-
 LLVMDisassembler::LLVMDisassembler(const std::string& filename)
     : Disassembler(filename)
     , logger(log4cxx::Logger::getLogger("LLVMDisassembler"))
@@ -117,24 +116,29 @@ void LLVMDisassembler::disassemble() {
     std::stack<LLVMBasicBlock*> remaining_blocks;
     SectionRef text_section = sections[".text"];
 
-    std::for_each(symbols.begin(), symbols.end(),
-                  [&](std::pair<const std::string, SymbolRef> x) {
-                      uint64_t result;
-                      bool contains;
-                      SymbolRef::Type symbol_type;
+       for (auto x = symbols.begin(); x != symbols.end(); ++x) {
+               uint64_t result;
+               bool contains;
+               SymbolRef::Type symbol_type;
+
+/*
+ * TODO: If we jump into some Basic Block we need to split it there into two
+ */
 
-                      if (text_section.containsSymbol(x.second, contains) || !contains)
-                          return;
+               if (text_section.containsSymbol(x->second, contains) || !contains)
+                       continue;
 
-                      if (x.second.getType(symbol_type)
-                          || SymbolRef::ST_Function != symbol_type)
-                          return;
+               if (x->second.getType(symbol_type)
+                       || SymbolRef::ST_Function != symbol_type)
+                       continue;
 
-                      if (!x.second.getAddress(result)) {
-                          remaining_functions.push(new LLVMFunction(x.first, result));
-                          LOG4CXX_DEBUG(logger, "Disasembling " << x.first);
-                      }
-                  });
+               if (!x->second.getAddress(result)) {
+                       LLVMFunction * fun = new LLVMFunction(x->first, result);
+                       remaining_functions.push(fun);
+                       functions.insert(std::make_pair(result, fun));
+                       LOG4CXX_DEBUG(logger, "Disasembling " << x->first);
+               }
+       }
 
     StringRef bytes;
     text_section.getContents(bytes);
@@ -149,7 +153,9 @@ void LLVMDisassembler::disassemble() {
         // if ("_start" != current_function->getName())
         //  continue;
 
-        remaining_blocks.push(new LLVMBasicBlock(current_function->getStartAddress()));
+               LLVMBasicBlock * block = new LLVMBasicBlock(current_function->getStartAddress());
+        remaining_blocks.push(block);
+               blocks.insert(std::make_pair(block->getStartAddress(), block));
 
         while (remaining_blocks.size()) {
             LLVMBasicBlock * current_block = remaining_blocks.top();
@@ -187,14 +193,24 @@ void LLVMDisassembler::disassemble() {
                                                jmptarget += base_address;
                         if (!MIA->isIndirectBranch(inst)) {
                             if (MIA->isCall(inst)) {
-                                                               if (blocks.find(jmptarget) == blocks.end())
-                                                                       remaining_functions.push(new LLVMFunction("<Unnamed>", jmptarget));
+                                                               if (functions.find(jmptarget) == functions.end()) {
+                                                                       std::stringstream s;
+                                                                       s << "<Unnamed 0x" << std::hex << jmptarget << ">";
+                                                                       LLVMFunction * fun = new LLVMFunction(s.str(), jmptarget);
+                                                                       functions.insert(std::make_pair(jmptarget, fun));
+                                                                       remaining_functions.push(fun);
+                                                               }
                             } else {
-                                                               if (blocks.find(jmptarget) == blocks.end())
-                                remaining_blocks.push(new LLVMBasicBlock(jmptarget));
+                                                               if (blocks.find(jmptarget) == blocks.end()) {
+                                                                       LLVMBasicBlock * block = new LLVMBasicBlock(jmptarget);
+                                                                       blocks.insert(std::make_pair(block->getStartAddress(), block));
+                                                                       remaining_blocks.push(block);
+                                                               }
                                 if (MIA->isConditionalBranch(inst)) {
                                                                        jmptarget = base_address + current_address + inst_size;
-                                                                       if (blocks.find(jmptarget) == blocks.end())
+                                                                       if (blocks.find(jmptarget) == blocks.end()) {
+                                                                               LLVMBasicBlock * block = new LLVMBasicBlock(jmptarget);
+                                                                               blocks.insert(std::make_pair(block->getStartAddress(), block));
                                                                                remaining_blocks.push(new LLVMBasicBlock(jmptarget));
                                 }
                             }
@@ -207,8 +223,9 @@ void LLVMDisassembler::disassemble() {
 
                 if (inst_size == 0 || MIA->isTerminator(inst) || MIA->isBranch(inst)) {
                     current_block->setEndAddress(current_address + base_address);
-                                       blocks.insert(std::make_pair(current_block->getStartAddress(), current_block));
-                    LOG4CXX_DEBUG(logger, "Finished Block at " << current_block->getEndAddress());
+                    LOG4CXX_DEBUG(logger, "Finished Block at " << std::hex << 
+                                                                 current_block->getEndAddress());
+                                                                       }
                     break;
                 }
                 current_address += inst_size;
@@ -247,6 +264,15 @@ void LLVMDisassembler::readSections() {
 
 }
 
-BasicBlock * LLVMDisassembler::generateControlFlowGraph(uint64_t address) {
+void LLVMDisassembler::forEachFunction(std::function<void (uint64_t, Function*)> callback) {
+       std::for_each(functions.begin(), functions.end(),
+                                 [&](std::pair<uint64_t, LLVMFunction*> x) {
+                                         callback(x.first, x.second);
+                                 });
+}
+
+
+
+void LLVMDisassembler::generateControlFlowGraph(uint64_t address) {
 
 }