Clean up disassemble() function

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

diff --git a/src/disassembler/llvm/LLVMDisassembler.cxx b/src/disassembler/llvm/LLVMDisassembler.cxx
index 6e76c91bc27c8121f7088460401ae3c20d4faf47..7980a9e1377a315dd97fbe1fb945b64edf8ead8e 100644 (file)
--- a/src/disassembler/llvm/LLVMDisassembler.cxx
+++ b/src/disassembler/llvm/LLVMDisassembler.cxx
@@ -121,14 +121,12 @@ LLVMDisassembler::~LLVMDisassembler() {
                   });
 }
 
-/*
- * TODO: If we jump into some Basic Block we need to split it there into two
- */
 void LLVMDisassembler::disassemble() {
     std::stack<LLVMFunction*> remaining_functions;
     std::stack<LLVMBasicBlock*> remaining_blocks;
     SectionRef text_section = sections[".text"];
 
+       // Assume all function symbols actually start a real function
     for (auto x = symbols.begin(); x != symbols.end(); ++x) {
         uint64_t result;
         bool contains;
@@ -160,9 +158,6 @@ void LLVMDisassembler::disassemble() {
 
         LOG4CXX_DEBUG(logger, "Handling function " << current_function->getName());
 
-        // if ("_start" != current_function->getName())
-        //  continue;
-
         LLVMBasicBlock * block = new LLVMBasicBlock(current_function->getStartAddress(), this);
         remaining_blocks.push(block);
         blocks.insert(std::make_pair(block->getStartAddress(), block));
@@ -185,19 +180,6 @@ void LLVMDisassembler::disassemble() {
                 if(llvm::MCDisassembler::Success ==
                    DisAsm->getInstruction(inst, inst_size, ref, current_address, nulls(), nulls())) {
 
-                    uint8_t bytes[inst_size+2];
-                    ref.readBytes(current_address, inst_size, bytes);
-                    s << '\t';
-                    for(uint8_t* cur = bytes; cur < bytes + inst_size; ++cur) {
-                        s.write_hex(*cur);
-                        s << ' ';
-                    }
-                    s << '\t';
-
-                    IP->printInst(&inst, s, "");
-
-                    LOG4CXX_DEBUG(logger, std::hex << current_address + base_address << s.str());
-
                     uint64_t jmptarget;
                     if (MIA->evaluateBranch(inst, current_address, inst_size, jmptarget)) {
                         jmptarget += base_address;
@@ -245,6 +227,39 @@ void LLVMDisassembler::disassemble() {
         }
         LOG4CXX_DEBUG(logger, "Finished function " << current_function->getName());
     }
+
+       // Split blocks where jumps are going inside the block
+       for (auto it = blocks.begin(); it != blocks.end(); ++it) {
+               LLVMBasicBlock * current_block = it->second;
+               uint64_t inst_size;
+               uint64_t base_address;
+               text_section.getAddress(base_address);
+               uint64_t current_address = current_block->getStartAddress() - base_address;
+               while(current_block->getEndAddress() - base_address != current_address) {
+                       MCInst inst;
+                       std::string buf;
+                       llvm::raw_string_ostream s(buf);
+
+                       if(llvm::MCDisassembler::Success ==
+                          DisAsm->getInstruction(inst, inst_size, ref, current_address, nulls(), nulls())) {
+                               auto other = blocks.find(current_address + inst_size + base_address);
+
+                               if (other != blocks.end()) {
+                                       LOG4CXX_DEBUG(logger, "Shortening block starting at "
+                                                                 << std::hex
+                                                                 << current_block->getStartAddress()
+                                                                 << " now ending at "
+                                                                 << other->first);
+                                       current_block->setEndAddress(current_address + inst_size + base_address);
+                                       current_block->setNextBlock(0, other->first);
+                                       current_block->setNextBlock(1, 0);
+                               }
+                       } else {
+                               inst_size = 1;
+                       }
+                       current_address += inst_size;
+               }
+       }
 }
 
 void LLVMDisassembler::readSymbols() {
@@ -283,7 +298,8 @@ void LLVMDisassembler::forEachFunction(std::function<void (uint64_t, Function*)>
                   });
 }
 
-void LLVMDisassembler::printEachInstruction(uint64_t start, uint64_t end, std::function<void (const std::string&)> fun) {
+void LLVMDisassembler::printEachInstruction(uint64_t start, uint64_t end,
+                                                                                       std::function<void (uint8_t*, size_t, const std::string&)> fun) {
     SectionRef text_section = sections[".text"];
     uint64_t base_address;
     text_section.getAddress(base_address);
@@ -304,16 +320,11 @@ void LLVMDisassembler::printEachInstruction(uint64_t start, uint64_t end, std::f
 
             uint8_t bytes[inst_size+2];
             ref.readBytes(current_address, inst_size, bytes);
-            for(uint8_t* cur = bytes; cur < bytes + inst_size; ++cur) {
-                s.write_hex(*cur);
-                s << ' ';
-            }
-            s << '\t';
 
             IP->printInst(&inst, s, "");
-                       fun(s.str());
+                       fun(bytes, inst_size, s.str());
         } else {
-                       fun("Invalid Byte");
+                       fun(NULL, 0, "Invalid Byte");
                        inst_size = 1;
                }