WIP: vmi_mmap_guest_2

plugins/inmemoryscanner/src/lib/IYaraInterface.h

@@ -2,26 +2,25 @@
 #define INMEMORYSCANNER_IYARAINTERFACE_H
 
 #include "Common.h"
-#include <exception>
-#include <string>
+#include <span>
+#include <stdexcept>
 #include <vector>
 #include <vmicore/vmi/IMemoryMapping.h>
 
 namespace InMemoryScanner
 {
     class YaraException : public std::runtime_error
     {
-      public:
-        explicit YaraException(const std::string& Message) : std::runtime_error(Message.c_str()){};
+        using std::runtime_error::runtime_error;
     };
 
     class IYaraInterface
     {
       public:
         virtual ~IYaraInterface() = default;
 
-        virtual std::unique_ptr<std::vector<Rule>>
-        scanMemory(const std::vector<VmiCore::MappedRegion>& mappedRegions) = 0;
+        virtual std::vector<Rule> scanMemory(VmiCore::addr_t regionBase,
+                                             std::span<const VmiCore::MappedRegion> mappedRegions) = 0;
 
       protected:
         IYaraInterface() = default;

plugins/inmemoryscanner/src/lib/Scanner.cpp

@@ -1,10 +1,9 @@
 #include "Scanner.h"
 #include "Common.h"
 #include "Filenames.h"
+#include <algorithm>
 #include <fmt/core.h>
 #include <future>
-#include <iterator>
-#include <span>
 #include <vmicore/callback.h>
 #include <vmicore/os/PagingDefinitions.h>
 
@@ -49,20 +48,19 @@ namespace InMemoryScanner
 
     bool Scanner::shouldRegionBeScanned(const MemoryRegion& memoryRegionDescriptor)
     {
-        bool verdict = true;
         if (configuration->isScanAllRegionsActivated())
         {
             return true;
         }
         if (memoryRegionDescriptor.isSharedMemory && !memoryRegionDescriptor.isProcessBaseImage)
         {
-            verdict = false;
             logger->info("Skipping: Is shared memory and not the process base image.");
+            return false;
         }
-        return verdict;
+        return true;
     }
 
-    std::vector<uint8_t> Scanner::constructPaddedMemoryRegion(const std::vector<MappedRegion>& regions)
+    std::vector<uint8_t> Scanner::constructPaddedMemoryRegion(std::span<const MappedRegion> regions)
     {
         std::vector<uint8_t> result;
 
@@ -74,22 +72,24 @@ namespace InMemoryScanner
         std::size_t regionSize = 0;
         for (const auto& region : regions)
         {
-            regionSize += region.mapping.size();
+            regionSize += region.asSpan().size();
             regionSize += pageSizeInBytes;
         }
         // last region should not have succeeding padding page
         regionSize -= pageSizeInBytes;
 
         result.reserve(regionSize);
         // copy first region
-        std::copy(regions.front().mapping.begin(), regions.front().mapping.end(), std::back_inserter(result));
+        auto frontRegionSpan = regions.front().asSpan();
+        std::ranges::copy(frontRegionSpan.begin(), frontRegionSpan.end(), std::back_inserter(result));
 
         for (std::size_t i = 1; i < regions.size(); i++)
         {
             const auto& region = regions[i];
             // padding page
             result.insert(result.end(), pageSizeInBytes, 0);
-            std::copy(region.mapping.begin(), region.mapping.end(), std::back_inserter(result));
+            auto regionSpan = region.asSpan();
+            std::ranges::copy(regionSpan.begin(), regionSpan.end(), std::back_inserter(result));
         }
 
         return result;
@@ -105,47 +105,48 @@ namespace InMemoryScanner
                       {"Size", memoryRegionDescriptor.size},
                       {"Module", memoryRegionDescriptor.moduleName}});
 
-        if (shouldRegionBeScanned(memoryRegionDescriptor))
+        if (!shouldRegionBeScanned(memoryRegionDescriptor))
         {
-            auto memoryMapping = pluginInterface->mapProcessMemoryRegion(
-                memoryRegionDescriptor.base, dtb, bytesToNumberOfPages(memoryRegionDescriptor.size));
-            auto mappedRegions = memoryMapping->getMappedRegions().lock();
+            return;
+        }
 
-            if (mappedRegions->empty())
-            {
-                logger->debug("Extracted memory region has size 0, skipping");
-            }
-            else
-            {
-                if (configuration->isDumpingMemoryActivated())
-                {
-                    logger->debug("Start dumpVadRegionToFile", {{"Size", memoryMapping->getSizeInGuest()}});
+        auto memoryMapping = pluginInterface->mapProcessMemoryRegion(
+            memoryRegionDescriptor.base, dtb, bytesToNumberOfPages(memoryRegionDescriptor.size));
+        auto mappedRegions = memoryMapping->getMappedRegions();
 
-                    auto paddedRegion = constructPaddedMemoryRegion(*mappedRegions);
+        if (mappedRegions.empty())
+        {
+            logger->debug("Extracted memory region has size 0, skipping");
+            return;
+        }
 
-                    dumping->dumpMemoryRegion(processName, pid, memoryRegionDescriptor, paddedRegion);
-                }
+        if (configuration->isDumpingMemoryActivated())
+        {
+            logger->debug("Start dumpVadRegionToFile", {{"Size", memoryRegionDescriptor.size}});
 
-                logger->debug("Start scanMemory", {{"Size", memoryMapping->getSizeInGuest()}});
+            auto paddedRegion = constructPaddedMemoryRegion(mappedRegions);
 
-                // The semaphore protects the yara rules from being accessed more than YR_MAX_THREADS (32 atm.) times in
-                // parallel.
-                semaphore.wait();
-                auto results = yaraInterface->scanMemory(*mappedRegions);
-                semaphore.notify();
+            dumping->dumpMemoryRegion(processName, pid, memoryRegionDescriptor, paddedRegion);
+        }
 
-                logger->debug("End scanMemory");
+        logger->debug("Start scanMemory", {{"Size", memoryRegionDescriptor.size}});
 
-                if (!results->empty())
-                {
-                    for (const auto& result : *results)
-                    {
-                        pluginInterface->sendInMemDetectionEvent(result.ruleName);
-                    }
-                    outputXml.addResult(processName, pid, memoryRegionDescriptor.base, *results);
-                    logInMemoryResultToTextFile(processName, pid, memoryRegionDescriptor.base, *results);
-                }
+        // The semaphore protects the yara rules from being accessed more than YR_MAX_THREADS (32 atm.) times in
+        // parallel.
+        semaphore.wait();
+        auto results = yaraInterface->scanMemory(memoryRegionDescriptor.base, mappedRegions);
+        semaphore.notify();
+
+        logger->debug("End scanMemory");
+
+        if (!results.empty())
+        {
+            for (const auto& result : results)
+            {
+                pluginInterface->sendInMemDetectionEvent(result.ruleName);
             }
+            outputXml.addResult(processName, pid, memoryRegionDescriptor.base, results);
+            logInMemoryResultToTextFile(processName, pid, memoryRegionDescriptor.base, results);
         }
     }
 

plugins/inmemoryscanner/src/lib/Scanner.h

@@ -7,6 +7,7 @@
 #include "Semaphore.h"
 #include <condition_variable>
 #include <memory>
+#include <span>
 #include <vmicore/plugins/PluginInterface.h>
 #include <yara/limits.h> // NOLINT(modernize-deprecated-headers)
 
@@ -41,7 +42,7 @@ namespace InMemoryScanner
 
         [[nodiscard]] bool shouldRegionBeScanned(const VmiCore::MemoryRegion& memoryRegionDescriptor);
 
-        static std::vector<uint8_t> constructPaddedMemoryRegion(const std::vector<VmiCore::MappedRegion>& regions);
+        static std::vector<uint8_t> constructPaddedMemoryRegion(std::span<const VmiCore::MappedRegion> regions);
 
         void scanMemoryRegion(pid_t pid,
                               uint64_t dtb,

plugins/inmemoryscanner/src/lib/YaraInterface.cpp

@@ -1,8 +1,42 @@
 #include "YaraInterface.h"
 #include <fmt/core.h>
 
+using VmiCore::addr_t;
 using VmiCore::MappedRegion;
-using BlockIteratorPair = std::pair<std::vector<YR_MEMORY_BLOCK>::iterator, std::vector<YR_MEMORY_BLOCK>::iterator>;
+using VmiCore::PagingDefinitions::pageSizeInBytes;
+
+namespace
+{
+    struct YaraIteratorContext
+    {
+        std::vector<YR_MEMORY_BLOCK> blocks;
+        std::size_t index;
+    };
+
+    YR_MEMORY_BLOCK* get_next_block(YR_MEMORY_BLOCK_ITERATOR* iterator)
+    {
+        if (auto* iteratorContext = static_cast<YaraIteratorContext*>(iterator->context);
+            ++iteratorContext->index < iteratorContext->blocks.size())
+        {
+            return &iteratorContext->blocks[iteratorContext->index];
+        }
+
+        return nullptr;
+    }
+
+    YR_MEMORY_BLOCK* get_first_block(YR_MEMORY_BLOCK_ITERATOR* iterator)
+    {
+        auto* iteratorContext = static_cast<YaraIteratorContext*>(iterator->context);
+        iteratorContext->index = 0;
+
+        return &iteratorContext->blocks[iteratorContext->index];
+    }
+
+    const uint8_t* fetch_block_data(YR_MEMORY_BLOCK* block)
+    {
+        return static_cast<const uint8_t*>(block->context);
+    }
+}
 
 namespace InMemoryScanner
 {
@@ -11,74 +45,52 @@ namespace InMemoryScanner
         auto err = yr_initialize();
         if (err != ERROR_SUCCESS)
         {
-            throw YaraException("Cannot initialize Yara. Error code: " + std::to_string(err));
+            throw YaraException(fmt::format("Cannot initialize Yara. Error code: {}", err));
         }
 
         err = yr_rules_load(rulesFile.c_str(), &rules);
         if (err != ERROR_SUCCESS)
         {
-            throw YaraException("Cannot load rules. Error code: " + std::to_string(err));
+            throw YaraException(fmt::format("Cannot load rules. Error code: {}", err));
         }
     }
 
     YaraInterface::~YaraInterface()
     {
-        yr_rules_destroy(rules);
-        yr_finalize();
-    }
-
-    YR_MEMORY_BLOCK* get_next_block(YR_MEMORY_BLOCK_ITERATOR* iterator)
-    {
-        auto* blockVectorIterators = reinterpret_cast<BlockIteratorPair*>(iterator->context);
-        blockVectorIterators->first++;
-
-        if (blockVectorIterators->first == blockVectorIterators->second)
+        if (rules)
         {
-            return nullptr;
+            yr_rules_destroy(rules);
+            yr_finalize();
         }
-
-        return &*blockVectorIterators->first;
-    }
-
-    YR_MEMORY_BLOCK* get_first_block(YR_MEMORY_BLOCK_ITERATOR* iterator)
-    {
-        return &*reinterpret_cast<BlockIteratorPair*>(iterator->context)->first;
     }
 
-    const uint8_t* fetch_block_data(YR_MEMORY_BLOCK* block)
+    std::vector<Rule> YaraInterface::scanMemory(addr_t regionBase, std::span<const MappedRegion> mappedRegions)
     {
-        return reinterpret_cast<const uint8_t*>(block->context);
-    }
+        std::vector<Rule> results;
 
-    std::unique_ptr<std::vector<Rule>> YaraInterface::scanMemory(const std::vector<MappedRegion>& mappedRegions)
-    {
-        auto results = std::make_unique<std::vector<Rule>>();
-
-        std::vector<YR_MEMORY_BLOCK> blocks;
-        blocks.reserve(mappedRegions.size());
+        YaraIteratorContext iteratorContext{};
+        iteratorContext.blocks.reserve(mappedRegions.size());
         for (const auto& mappedRegion : mappedRegions)
         {
-            blocks.emplace_back(mappedRegion.mapping.size(),
-                                mappedRegion.guestBaseVA - mappedRegions.front().guestBaseVA,
-                                reinterpret_cast<void*>(mappedRegion.mapping.data()),
-                                &fetch_block_data);
+            iteratorContext.blocks.emplace_back(mappedRegion.num_pages * pageSizeInBytes,
+                                                mappedRegion.guestBaseVA - regionBase,
+                                                mappedRegion.mappingBase,
+                                                &fetch_block_data);
         }
-        auto blockIterators = std::make_pair(blocks.begin(), blocks.end());
 #ifdef LIBYARA_4_1
-        YR_MEMORY_BLOCK_ITERATOR iterator{.context = &blockIterators,
+        YR_MEMORY_BLOCK_ITERATOR iterator{.context = &iteratorContext,
                                           .first = &get_first_block,
                                           .next = &get_next_block,
                                           .file_size = nullptr,
                                           .last_error = ERROR_SUCCESS};
 #else
         YR_MEMORY_BLOCK_ITERATOR iterator{
-            .context = &blockIterators, .first = &get_first_block, .next = &get_next_block};
+            .context = &iteratorContext, .first = &get_first_block, .next = &get_next_block};
 #endif
 
-        auto err = yr_rules_scan_mem_blocks(rules, &iterator, 0, yaraCallback, results.get(), 0);
-        if (err != ERROR_SUCCESS)
+        if (auto err = yr_rules_scan_mem_blocks(rules, &iterator, 0, yaraCallback, &results, 0); err != ERROR_SUCCESS)
         {
-            throw YaraException("Error scanning memory. Error code: " + std::to_string(err));
+            throw YaraException(fmt::format("Error scanning memory. Error code: {}", err));
         }
 
         return results;

plugins/inmemoryscanner/src/lib/YaraInterface.h

@@ -2,9 +2,7 @@
 
 #include "Common.h"
 #include "IYaraInterface.h"
-#include <memory>
 #include <string>
-#include <vector>
 #include <yara.h>
 
 namespace InMemoryScanner
@@ -14,9 +12,32 @@ namespace InMemoryScanner
       public:
         explicit YaraInterface(const std::string& rulesFile);
 
+        YaraInterface(const YaraInterface& other) = delete;
+
+        YaraInterface(YaraInterface&& other) noexcept : rules(other.rules)
+        {
+            other.rules = nullptr;
+        }
+
+        YaraInterface& operator=(const YaraInterface& other) = delete;
+
+        YaraInterface& operator=(YaraInterface&& other) noexcept
+        {
+            if (this == &other)
+            {
+                return *this;
+            }
+
+            rules = other.rules;
+            other.rules = nullptr;
+
+            return *this;
+        }
+
         ~YaraInterface() override;
 
-        std::unique_ptr<std::vector<Rule>> scanMemory(const std::vector<VmiCore::MappedRegion>& mappedRegions) override;
+        std::vector<Rule> scanMemory(VmiCore::addr_t regionBase,
+                                     std::span<const VmiCore::MappedRegion> mappedRegions) override;
 
       private:
         YR_RULES* rules = nullptr;

plugins/inmemoryscanner/test/CMakeLists.txt

@@ -1,7 +1,7 @@
 add_executable(inmemoryscanner-test
         FakeYaraInterface.cpp
         Scanner_unittest.cpp
-        YaraInterface_unitttest.cpp)
+        YaraInterface_unittest.cpp)
 target_link_libraries(inmemoryscanner-test inmemoryscanner-obj pthread)
 
 # Setup bundled google test framework