From 40c8a7bbd831fcb8192f3a47d2be6b601e2cea98 Mon Sep 17 00:00:00 2001
From: Sebastian Poeplau <sebastian.poeplau@eurecom.fr>
Date: Tue, 6 Nov 2018 16:06:46 +0100
Subject: [PATCH 2/2] Finish the tool; ugly modification of the translator
 library

A better solution would be to preserve the old interface for other tools and
create a new method with our semantics...
---
 include/Translator/TranslatorInternal.h | 85 +++++++++++++------------
 tools/liftandcount/LiftAndCount.cpp     | 21 +++---
 2 files changed, 55 insertions(+), 51 deletions(-)

diff --git a/include/Translator/TranslatorInternal.h b/include/Translator/TranslatorInternal.h
index 1533c27..a3750d2 100644
--- a/include/Translator/TranslatorInternal.h
+++ b/include/Translator/TranslatorInternal.h
@@ -533,52 +533,57 @@ TranslatedBlock *X86Translator::translate(uint64_t address, uint64_t lastAddress
     /// Sometimes, we have things like mov eax, constant; jmp eax,
     /// which is reported as indirect jump, but really is a direct
     /// one. Need to handle this case.
-    switch (tb.se_tb_type) {
-        case TB_DEFAULT:
-        case TB_JMP:
-        case TB_CALL:
-            if (info.staticBranchTargets.size() != 1) {
-                LOGERROR("TB static targets mismatch TB type\n");
-                return NULL;
-            }
-            break;
-
-        case TB_COND_JMP:
-        case TB_REP:
-            if (info.staticBranchTargets.size() != 2) {
-                LOGERROR("TB static targets mismatch TB type\n");
-                return NULL;
-            }
-            break;
-
-        case TB_JMP_IND:
-        case TB_COND_JMP_IND:
-        case TB_CALL_IND:
-
-        case TB_RET:
-        case TB_IRET:
-        case TB_EXCP:
-        case TB_SYSENTER:
-            if (info.staticBranchTargets.size() != 0) {
-                LOGERROR("TB static targets mismatch TB type\n");
-                return NULL;
-            }
-            break;
-
-        default:
-            abort();
-    }
+    // switch (tb.se_tb_type) {
+    //     case TB_DEFAULT:
+    //     case TB_JMP:
+    //     case TB_CALL:
+    //         if (info.staticBranchTargets.size() != 1) {
+    //             LOGERROR("TB static targets mismatch TB type\n");
+    //             return NULL;
+    //         }
+    //         break;
+
+    //     case TB_COND_JMP:
+    //     case TB_REP:
+    //         if (info.staticBranchTargets.size() != 2) {
+    //             LOGERROR("TB static targets mismatch TB type\n");
+    //             return NULL;
+    //         }
+    //         break;
+
+    //     case TB_JMP_IND:
+    //     case TB_COND_JMP_IND:
+    //     case TB_CALL_IND:
+
+    //     case TB_RET:
+    //     case TB_IRET:
+    //     case TB_EXCP:
+    //     case TB_SYSENTER:
+    //         if (info.staticBranchTargets.size() != 0) {
+    //             LOGERROR("TB static targets mismatch TB type\n");
+    //             return NULL;
+    //         }
+    //         break;
+
+    //     default:
+    //         abort();
+    // }
 
     /* Make sure that we don't ask to translate a block that is actually shorter */
     if (tb.pcOfLastInstr != lastAddress) {
         /* XXX: sometimes TBs may be to big, QEMU stops translating when
            it runs out of buffer space. Could split the BB in multiple chunks. */
-        LOGERROR("Translation didn't finish at bb end (" << hexval(tb.pcOfLastInstr) << " instead of "
-                                                         << hexval(lastAddress) << ")\n");
-
-        return NULL;
+        // LOGERROR("Translation didn't finish at bb end (" << hexval(tb.pcOfLastInstr) << " instead of "
+        //                                                  << hexval(lastAddress) << ")\n");
+
+        /* Jumps at the end of a basic block don't get translated */
+        if (lastAddress - tb.pcOfLastInstr > 8) {
+            LOGERROR("Large discrepancy between bb end and actual end ("
+                     << hexval(tb.pcOfLastInstr) << " instead of " << hexval(lastAddress) << ")\n");
+            return NULL;
+        }
     }
-    assert(tb.pcOfLastInstr == lastAddress);
+    // assert(tb.pcOfLastInstr == lastAddress);
 
     /* Handle call $+5 instructions, used to push the address of the next instruction */
     if (tb.se_tb_type == TB_CALL) {
diff --git a/tools/liftandcount/LiftAndCount.cpp b/tools/liftandcount/LiftAndCount.cpp
index 8251f6f..ab934be 100644
--- a/tools/liftandcount/LiftAndCount.cpp
+++ b/tools/liftandcount/LiftAndCount.cpp
@@ -31,6 +31,8 @@ int main(int argc, char *argv[])
         return -1;
     }
 
+    long instCount = 0;
+    auto translator = s2etools::X86Translator(bitcodeLibraryName, binary);
     std::ifstream addressFile(addressFileName);
     while (!addressFile.eof()) {
         uint64_t start, end;
@@ -38,21 +40,18 @@ int main(int argc, char *argv[])
         if (!addressFile.good())
             break;
 
-        auto translator = s2etools::X86Translator(bitcodeLibraryName, binary);
-        auto translatedBlock = translator.translate(start, end);
-        auto llvmFunction = translatedBlock->getFunction();
-        if (llvmFunction->size() != 1) {
-            std::cerr << "Error: translation yielded more than one basic block" << std::endl;
+        std::unique_ptr<s2etools::TranslatedBlock> translatedBlock(
+            translator.translate(start, end));
+        if (!translatedBlock) {
+            std::cerr << "Failed to translate basic block at 0x"
+                      << std::hex << start << std::endl;
             return -1;
         }
 
-        // for (auto& inst : llvmFunction->front()) {
-        //     std::cout << inst.getName().str() << std::endl;
-        // }
-        llvmFunction->dump();
-        std::cout << llvmFunction->front().size() << std::endl;
+        for (auto& bb : *translatedBlock->getFunction())
+            instCount += bb.size();
     }
 
-    std::cout << "Done" << std::endl;
+    std::cout << instCount << std::endl;
     return 0;
 }
-- 
2.23.0

