From 99f8efac9ae3d27fc38c5c9a0983e9a563a46a2d Mon Sep 17 00:00:00 2001
From: WerWolv <werwolv98@gmail.com>
Date: Sun, 6 Feb 2022 15:28:20 +0100
Subject: [PATCH] pattern: Fixed pragma once memory corruption

---
 .../hex/pattern_language/preprocessor.hpp     |  7 ++++
 .../source/pattern_language/preprocessor.cpp  | 39 ++++++++++++-------
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/lib/libimhex/include/hex/pattern_language/preprocessor.hpp b/lib/libimhex/include/hex/pattern_language/preprocessor.hpp
index 1c732ef84..b92655a10 100644
--- a/lib/libimhex/include/hex/pattern_language/preprocessor.hpp
+++ b/lib/libimhex/include/hex/pattern_language/preprocessor.hpp
@@ -22,10 +22,15 @@ namespace hex::pl {
         std::optional<std::string> preprocess(const std::string &code, bool initialRun = true);
 
         void addPragmaHandler(const std::string &pragmaType, const std::function<bool(const std::string &)> &function);
+        void removePragmaHandler(const std::string &pragmaType);
         void addDefaultPragmaHandlers();
 
         const std::optional<PatternLanguageError> &getError() { return this->m_error; }
 
+        bool shouldOnlyIncludeOnce() {
+            return this->m_onlyIncludeOnce;
+        }
+
     private:
         [[noreturn]] static void throwPreprocessorError(const std::string &error, u32 lineNumber) {
             throw PatternLanguageError(lineNumber, "Preprocessor: " + error);
@@ -39,6 +44,8 @@ namespace hex::pl {
         std::set<fs::path> m_onceIncludedFiles;
 
         std::optional<PatternLanguageError> m_error;
+
+        bool m_onlyIncludeOnce = false;
     };
 
 }
\ No newline at end of file
diff --git a/lib/libimhex/source/pattern_language/preprocessor.cpp b/lib/libimhex/source/pattern_language/preprocessor.cpp
index 63c6fc418..f5638f146 100644
--- a/lib/libimhex/source/pattern_language/preprocessor.cpp
+++ b/lib/libimhex/source/pattern_language/preprocessor.cpp
@@ -79,25 +79,30 @@ namespace hex::pl {
                                 throwPreprocessorError(hex::format("{0}: No such file or directory", includeFile.c_str()), lineNumber);
                         }
 
-                        bool shouldInclude = true;
-                        this->addPragmaHandler("once", [&, includePath, this](const std::string &value) {
-                            auto [iter, added] = this->m_onceIncludedFiles.insert(includePath);
-                            if (!added) shouldInclude = false;
-                            return value.empty();
-                        });
+                        Preprocessor preprocessor;
+                        preprocessor.addDefaultPragmaHandlers();
+                        preprocessor.m_defines           = this->m_defines;
+                        preprocessor.m_onceIncludedFiles = this->m_onceIncludedFiles;
+
+                        auto preprocessedInclude = preprocessor.preprocess(file.readString(), true);
 
-                        auto preprocessedInclude = this->preprocess(file.readString(), false);
                         if (!preprocessedInclude.has_value())
-                            throw this->m_error;
+                            throw *preprocessor.m_error;
 
-                        if (shouldInclude) {
-                            auto content = preprocessedInclude.value();
+                        if (preprocessor.shouldOnlyIncludeOnce()) {
+                            auto [iter, added] = this->m_onceIncludedFiles.insert(includePath);
+                            if (added) {
+                                auto content = preprocessedInclude.value();
 
-                            std::replace(content.begin(), content.end(), '\n', ' ');
-                            std::replace(content.begin(), content.end(), '\r', ' ');
+                                std::replace(content.begin(), content.end(), '\n', ' ');
+                                std::replace(content.begin(), content.end(), '\r', ' ');
 
-                            output += content;
+                                output += content;
+                            }
                         }
+
+                        this->m_defines           = preprocessor.m_defines;
+                        this->m_onceIncludedFiles = preprocessor.m_onceIncludedFiles;
                     } else if (code.substr(offset, 6) == "define") {
                         offset += 6;
 
@@ -234,6 +239,10 @@ namespace hex::pl {
         this->m_pragmaHandlers[pragmaType] = function;
     }
 
+    void Preprocessor::removePragmaHandler(const std::string &pragmaType) {
+        this->m_pragmaHandlers.erase(pragmaType);
+    }
+
     void Preprocessor::addDefaultPragmaHandlers() {
         this->addPragmaHandler("MIME", [](const std::string &value) {
             return !std::all_of(value.begin(), value.end(), isspace) && !value.ends_with('\n') && !value.ends_with('\r');
@@ -241,7 +250,9 @@ namespace hex::pl {
         this->addPragmaHandler("endian", [](const std::string &value) {
             return value == "big" || value == "little" || value == "native";
         });
-        this->addPragmaHandler("once", [](const std::string &value) {
+        this->addPragmaHandler("once", [this](const std::string &value) {
+            this->m_onlyIncludeOnce = true;
+
             return value.empty();
         });
     }