[REFACT] Simplified DisasmCache

hasherezade · Aug 25, 2024 · 3b0bd21 · 3b0bd21
1 parent bf31b43
commit 3b0bd21
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 53 deletions.
diff --git a/DisasmCache.h b/DisasmCache.h
@@ -4,43 +4,18 @@
 #include <string>
 #include <map>
 
-#include "crc.h"
+#include "Crc.h"
 
 //----
 
-struct InstrInfo
+struct DisasmCache
 {
     static uint64_t calcChecks(const std::string& _disasm)
     {
         uint64_t init = 0;
         return crc64(init, (unsigned char*)_disasm.c_str(), _disasm.length());
     }
 
-    //---
-
-    InstrInfo() : disasm(nullptr), checks(0) { }
-
-    InstrInfo(const InstrInfo& other) {
-        this->disasm = new std::string(other.disasm->c_str());
-        this->checks = other.checks;
-    }
-
-    InstrInfo(const std::string& _disasm)
-        : disasm(nullptr), checks(0)
-    {
-        this->disasm = new std::string(_disasm.c_str());
-        this->checks = calcChecks(_disasm);
-    }
-
-	~InstrInfo() { delete disasm; }
-
-	//---
-	std::string* disasm;
-	uint64_t checks;
-};
-
-struct DisasmCache
-{
 	DisasmCache() {}
 
 	~DisasmCache() {
@@ -50,7 +25,7 @@ struct DisasmCache
 		}
 	}
 
-    InstrInfo* get(uint32_t crc32)
+    std::string* get(uint32_t crc32)
     {
         auto itr = disasmLines.find(crc32);
         if (itr == disasmLines.end()) {
@@ -59,18 +34,17 @@ struct DisasmCache
         return disasmLines[crc32];
     }
 
-    InstrInfo* put(const std::string& _disasm)
+    std::string* put(const std::string& _disasm)
     {
-        auto checks = InstrInfo::calcChecks(_disasm);
-        auto itr = disasmLines.find(checks);
-        if (itr != disasmLines.end()) {
-            return itr->second;
+        auto checks = calcChecks(_disasm);
+        auto infoPtr = get(checks);
+        if (infoPtr) {
+            return infoPtr;
         }
-        InstrInfo* info = new InstrInfo(_disasm);
-        disasmLines[checks] = info;
-        return info;
+        disasmLines[checks] = new std::string(_disasm);
+        return disasmLines[checks];
     }
 
 	//---
-	std::map<uint64_t, InstrInfo*> disasmLines;
+	std::map<uint64_t, std::string*> disasmLines;
 };
diff --git a/TinyTracer.cpp b/TinyTracer.cpp
@@ -841,9 +841,9 @@ VOID MonitorFunctionArgs(IMG Image, const WFuncInfo& funcInfo)
 
 DisasmCache m_disasmCache;
 
-VOID LogInstruction(const CONTEXT* ctxt, THREADID tid, InstrInfo* disasmInfo)//, std::string* str)
+VOID LogInstruction(const CONTEXT* ctxt, THREADID tid, std::string* disasm)
 {
-    if (!disasmInfo) return;
+    if (!disasm) return;
 
     PinLocker locker;
 
@@ -872,18 +872,10 @@ VOID LogInstruction(const CONTEXT* ctxt, THREADID tid, InstrInfo* disasmInfo)//,
         base = query_region_base(Address);
         rva = Address - base;
     }
-/*
-    if (str->compare(*disasmInfo->disasm) != 0) {
-        std::cout << "Different strings!!! "
-            << disasmInfo->disasm->c_str() << " (" <<std::hex << InstrInfo::calcChecks(*disasmInfo->disasm) << ") != "
-            << *str << " (" << std::hex << InstrInfo::calcChecks(*str) << ")" << std::endl;
-    }
-*/
     if (base != UNKNOWN_ADDR && rva != UNKNOWN_ADDR) {
         std::stringstream ss;
         ss << "[" << std::dec << tid << "] ";
-        ss << std::hex << disasmInfo->checks << " : ";
-        ss << disasmInfo->disasm->c_str();
+        ss << disasm->c_str();
         if (!base && rva == (ADDRINT)m_Settings.disasmStart) {
             ss << " # disasm start";
         }
@@ -909,21 +901,22 @@ VOID InstrumentInstruction(INS ins, VOID *v)
     IMG pImg = IMG_FindByAddress(INS_Address(ins));
     const BOOL isMyImg = pInfo.isMyImg(pImg);
     BOOL shouldTrace = isMyImg;
-    if (m_Settings.followShellcode != t_shellc_options::SHELLC_DO_NOT_FOLLOW && !IMG_Valid(pImg)) {
+    if (m_Settings.followShellcode != t_shellc_options::SHELLC_DO_NOT_FOLLOW
+        && !IMG_Valid(pImg))
+    {
         shouldTrace = TRUE;
     }
     if (!shouldTrace) return;
 
     if (m_Settings.disasmStart) {
-        InstrInfo* info = m_disasmCache.put(INS_Disassemble(ins));
-        if (info) {
+        std::string* disasm = m_disasmCache.put(INS_Disassemble(ins));
+        if (disasm) {
             INS_InsertCall(
                 ins,
                 IPOINT_BEFORE, (AFUNPTR)LogInstruction,
                 IARG_CONTEXT,
                 IARG_THREAD_ID,
-                IARG_PTR, info,
-                //IARG_PTR, new std::string(INS_Disassemble(ins)),
+                IARG_PTR, disasm,
                 IARG_END
             );
         }