patterns: Added recursion and array size limit pragma

WerWolv · Sep 21, 2021 · 7556428 · 4g3v · Sep 26, 2021 · WerWolv
1 parent d1c0517
commit 7556428
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 5 deletions.
diff --git a/plugins/libimhex/include/hex/pattern_language/ast_node.hpp b/plugins/libimhex/include/hex/pattern_language/ast_node.hpp
@@ -821,6 +821,10 @@ namespace hex::pl {
                             [](auto &&size) -> u128 { return size; }
                     }, literal->getValue());
 
+                    auto limit = evaluator->getArrayLimit();
+                    if (entryCount > limit)
+                        LogConsole::abortEvaluation(hex::format("array grew past set limit of {}", limit), this);
+
                     for (u64 i = 0; i < entryCount; i++) {
                         auto pattern = this->m_type->createPatterns(evaluator).front();
 
@@ -833,6 +837,10 @@ namespace hex::pl {
                     }
                 } else if (auto whileStatement = dynamic_cast<ASTNodeWhileStatement*>(sizeNode)) {
                     while (whileStatement->evaluateCondition(evaluator)) {
+                        auto limit = evaluator->getArrayLimit();
+                        if (entryCount > limit)
+                            LogConsole::abortEvaluation(hex::format("array grew past set limit of {}", limit), this);
+
                         auto pattern = this->m_type->createPatterns(evaluator).front();
 
                         pattern->setVariableName(hex::format("[{}]", entryCount));
@@ -846,6 +854,10 @@ namespace hex::pl {
                 }
             } else {
                 while (true) {
+                    auto limit = evaluator->getArrayLimit();
+                    if (entryCount > limit)
+                        LogConsole::abortEvaluation(hex::format("array grew past set limit of {}", limit), this);
+
                     auto pattern = this->m_type->createPatterns(evaluator).front();
                     std::vector<u8> buffer(pattern->getSize());
 

diff --git a/plugins/libimhex/include/hex/pattern_language/evaluator.hpp b/plugins/libimhex/include/hex/pattern_language/evaluator.hpp
@@ -8,6 +8,8 @@
 #include <hex/pattern_language/log_console.hpp>
 #include <hex/api/content_registry.hpp>
 
+#include <hex/helpers/fmt.hpp>
+
 namespace hex::prv { class Provider; }
 
 namespace hex::pl {
@@ -27,8 +29,17 @@ namespace hex::pl {
         }
 
         struct Scope { PatternData *parent; std::vector<PatternData*>* scope; };
-        void pushScope(PatternData *parent, std::vector<PatternData*> &scope) { this->m_scopes.push_back({ parent, &scope }); }
-        void popScope() { this->m_scopes.pop_back(); }
+        void pushScope(PatternData *parent, std::vector<PatternData*> &scope) {
+            if (this->m_scopes.size() > this->m_evalDepth)
+                LogConsole::abortEvaluation(hex::format("recursion limit of {} reached", this->m_evalDepth));
+
+            this->m_scopes.push_back({ parent, &scope });
+        }
+
+        void popScope() {
+            this->m_scopes.pop_back();
+        }
+
         const Scope& getScope(s32 index) {
             static Scope empty;
 
@@ -58,6 +69,24 @@ namespace hex::pl {
             return this->m_defaultEndian;
         }
 
+        void setEvaluationDepth(u32 evalDepth) {
+            this->m_evalDepth = evalDepth;
+        }
+
+        [[nodiscard]]
+        u32 getEvaluationDepth() const {
+            return this->m_evalDepth;
+        }
+
+        void setArrayLimit(u32 arrayLimit) {
+            this->m_arrayLimit = arrayLimit;
+        }
+
+        [[nodiscard]]
+        u32 getArrayLimit() const {
+            return this->m_arrayLimit;
+        }
+
         u64& dataOffset() { return this->m_currOffset; }
 
         bool addCustomFunction(const std::string &name, u32 numParams, const ContentRegistry::PatternLanguageFunctions::Callback &function) {
@@ -86,6 +115,8 @@ namespace hex::pl {
         LogConsole m_console;
 
         std::endian m_defaultEndian = std::endian::native;
+        u32 m_evalDepth;
+        u32 m_arrayLimit;
 
         std::vector<Scope> m_scopes;
         std::map<std::string, ContentRegistry::PatternLanguageFunctions::Function> m_customFunctions;

diff --git a/plugins/libimhex/include/hex/pattern_language/pattern_language.hpp b/plugins/libimhex/include/hex/pattern_language/pattern_language.hpp
@@ -45,7 +45,8 @@ namespace hex::pl {
 
         prv::Provider *m_provider = nullptr;
         std::endian m_defaultEndian = std::endian::native;
-        u32 m_recursionLimit = 32;
+        u32 m_evalDepth;
+        u32 m_arrayLimit;
 
         std::optional<std::pair<u32, std::string>> m_currError;
     };

diff --git a/plugins/libimhex/source/pattern_language/pattern_language.cpp b/plugins/libimhex/source/pattern_language/pattern_language.cpp
@@ -43,7 +43,17 @@ namespace hex::pl {
             if (limit <= 0)
                 return false;
 
-            this->m_recursionLimit = limit;
+            this->m_evalDepth = limit;
+            return true;
+        });
+
+        this->m_preprocessor->addPragmaHandler("array_limit", [this](std::string value) {
+            auto limit = strtol(value.c_str(), nullptr, 0);
+
+            if (limit <= 0)
+                return false;
+
+            this->m_arrayLimit = limit;
             return true;
         });
 
@@ -69,6 +79,8 @@ namespace hex::pl {
         this->m_currError.reset();
         this->m_evaluator->getConsole().clear();
         this->m_evaluator->setProvider(provider);
+        this->m_evalDepth = 32;
+        this->m_arrayLimit = 0x1000;
 
         for (auto &node : this->m_currAST)
             delete node;
@@ -81,7 +93,8 @@ namespace hex::pl {
         }
 
         this->m_evaluator->setDefaultEndian(this->m_defaultEndian);
-        // this->m_evaluator->setRecursionLimit(this->m_recursionLimit);
+        this->m_evaluator->setEvaluationDepth(this->m_evalDepth);
+        this->m_evaluator->setArrayLimit(this->m_arrayLimit);
 
         auto tokens = this->m_lexer->lex(preprocessedCode.value());
         if (!tokens.has_value()) {