fixed some race conditions and some other small bugs

medioxor · May 14, 2024 · 83aef2b · 83aef2b
1 parent 0c2bcb7
commit 83aef2b
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 13 deletions.
diff --git a/src/Gadgify.cpp b/src/Gadgify.cpp
@@ -54,7 +54,7 @@ void Gadgify::GetGadgets(const std::function<void(uint64_t offset, const std::st
 }
 
 bool Gadgify::Disassemble(const std::function<void(std::vector<Instruction>, uint64_t)> &callback, cs_arch arch, cs_mode mode) {
-    ThreadPool pool(4);
+    ThreadPool pool;
     switch (arch)
     {
         case CS_ARCH_X86:
@@ -105,8 +105,9 @@ bool Gadgify::Disassemble(const std::function<void(std::vector<Instruction>, uin
     {
         return false;
     }
-    cs_option(handle, CS_OPT_DETAIL, CS_OPT_ON);
+    //cs_option(handle, CS_OPT_DETAIL, CS_OPT_ON);
     cs_insn *insn = cs_malloc(handle);
+    uint32_t returnCounter = 0;
     for (Section& section : sections_)
     {
         if (section.isExecutable_)
@@ -125,15 +126,26 @@ bool Gadgify::Disassemble(const std::function<void(std::vector<Instruction>, uin
                                 insn->op_str,
                                 insn->address
                     );
+                    totalDisassembled += insn->size;
                     if (strcmp(insn->mnemonic, returnMnemonic_.data()) == 0)
                     {
-                        pool.Enqueue([instructionChunk, section, callback]() {
-                            callback(instructionChunk, section.virtualAddress_);
+                        returnCounter++;
+                    }
+                    if (returnCounter == 40)
+                    {
+                        pool.Enqueue([chunk = std::move(instructionChunk), virtualAddress = section.virtualAddress_, callback]() {
+                            callback(chunk, virtualAddress);
                         });
-                        instructionChunk.clear();
+                        returnCounter = 0;
                     }
                     totalDisassembled += insn->size;
                 }
+                if (!instructionChunk.empty())
+                {
+                    pool.Enqueue([chunk = std::move(instructionChunk), virtualAddress = section.virtualAddress_, callback]() {
+                        callback(chunk, virtualAddress);
+                    });
+                }
                 totalDisassembled += 1;
             }
         }

diff --git a/src/ThreadPool.cpp b/src/ThreadPool.cpp
@@ -3,15 +3,19 @@
 ThreadPool::ThreadPool(uint32_t threads) {
     for (int i = 0; i < threads; i++)
     {
-        threads_.emplace_back(std::bind(&ThreadPool::Thread, this));
+        threads_.emplace_back([this] { Thread(); });
     }
 }
 
 void ThreadPool::Thread() {
     while (!stopThreads_)
     {
         std::unique_lock<std::mutex> latch(taskQueueMutex_);
-        // wait here for newTaskCondition_ to be notified via ThreadPool::Enqueue
+        // wait here for newTaskCondition_ to be notified via ThreadPool::Enqueue OR
+        // unblock if the taskQueue is not empty OR stopThreads_ is true
+        // note: Right after wait returns, latch.owns_lock() is true, and
+        // latch.mutex() is locked by the calling thread.
+        // (https://en.cppreference.com/w/cpp/thread/condition_variable/wait)
         newTaskCondition_.wait(latch, [&](){
             return !taskQueue_.empty() || stopThreads_;
         });
@@ -22,7 +26,6 @@ void ThreadPool::Thread() {
             taskQueue_.pop_front();
             latch.unlock();
             task();
-            latch.lock();
             runningTasks_--;
             finishedTaskCondition_.notify_one();
         }
@@ -47,11 +50,12 @@ ThreadPool::~ThreadPool() {
         thread.join();
 }
 
-void ThreadPool::Enqueue(const std::function<void(void)>& task) {
+void ThreadPool::Enqueue(const std::function<void()>& task) {
     // lock the task queue
     std::unique_lock<std::mutex> lock(taskQueueMutex_);
     // place a job on it
     taskQueue_.emplace_back(task);
-    // notify any on of the threads waiting on newTaskCondition_ to start a task
+    // notify any one of the threads waiting on newTaskCondition_ to start a task
     newTaskCondition_.notify_one();
+    // lock is unlocked during deconstructor on return
 }
diff --git a/src/ThreadPool.h b/src/ThreadPool.h
@@ -12,7 +12,7 @@
 class ThreadPool {
 public:
     explicit ThreadPool(uint32_t threads = std::thread::hardware_concurrency());
-    void Enqueue(const std::function<void(void)>& task);
+    void Enqueue(const std::function<void()>& task);
     void Wait();
     virtual ~ThreadPool();
 private:

diff --git a/src/main.cpp b/src/main.cpp
@@ -38,7 +38,7 @@ int main(int argc, char *argv[]) {
 
     bool isRaw = program.get<bool>("--raw");
     std::stringstream results;
-
+    std::mutex stringStreamMutex;
     if (isRaw)
     {
         std::vector<char> fileContents(File::Read(program.get<std::string>("binaryPath")));
@@ -54,8 +54,9 @@ int main(int argc, char *argv[]) {
     }
     else
     {
-        Gadgify::GetGadgets([&results](uint64_t offset, const std::string &gadget)
+        Gadgify::GetGadgets([&results, &stringStreamMutex](uint64_t offset, const std::string &gadget)
             {
+                std::lock_guard<std::mutex> streamLock(stringStreamMutex);
                 results << "0x" << std::hex << std::setfill('0') << std::setw(8) << offset << ": " << gadget << std::endl;
             },
             program.get<std::string>("binaryPath"),