From 109042ff52b0fdbef9e63265776e386cbccc2903 Mon Sep 17 00:00:00 2001
From: Deukhoofd <Deukhoofd@gmail.com>
Date: Thu, 13 May 2021 12:47:02 +0200
Subject: [PATCH] Adds support for including angelscript files using #include
 pragma, as long as they're within a set source directory.

---
 .../AngelScript/AngelScriptResolver.cpp       | 27 ++++++++++++++++++-
 .../AngelScript/AngelScriptResolver.hpp       | 14 +++++++---
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp b/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp
index 692552b..a589545 100644
--- a/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp
+++ b/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp
@@ -1,5 +1,6 @@
 #include "AngelScriptResolver.hpp"
 #include <CreatureLib/Battling/Models/Creature.hpp>
+#include <filesystem>
 #include <regex>
 #include "../../../extern/angelscript_addons/scriptdictionary/scriptdictionary.h"
 #include "../../../extern/angelscript_addons/scripthandle/scripthandle.h"
@@ -85,6 +86,7 @@ void AngelScriptResolver::Initialize(CreatureLib::Battling::BattleLibrary* arg,
         r = _engine->RegisterGlobalFunction("void print(const constString &in)", asFUNCTION(PrintConst), asCALL_CDECL);
         if (r < 0)
             throw ArbUt::Exception("Registering print function failed.");
+        _builder.SetIncludeCallback(IncludeCallback, this);
     }
 
     _builder.StartNewModule(_engine, "pkmn");
@@ -401,4 +403,27 @@ void AngelScriptResolver::InitializeByteCode(
         }
     }
     _typeDatabase = typeDatabase;
-}
\ No newline at end of file
+}
+
+i32 AngelScriptResolver::IncludeCallback(const char* include, const char*, CScriptBuilder* builder, void* userParam) {
+    auto* r = reinterpret_cast<AngelScriptResolver*>(userParam);
+    // If source directory is not set, bail out.
+    if (r->_sourceDirectory.empty()) {
+        return -100;
+    }
+    auto root = std::filesystem::path(r->_sourceDirectory);
+    // Resolve any special operators, to get the actual path.
+    auto path = (root / std::filesystem::path(include)).lexically_normal();
+
+    // Validate the path is inside the root directory. If not, bail out.
+    auto [rootEnd, nothing] = std::mismatch(root.begin(), root.end(), path.begin());
+    if (rootEnd != root.end()) {
+        return -101;
+    }
+
+    // If the file doesn't exist, bail out.
+    if (!std::filesystem::exists(path)) {
+        return -102;
+    }
+    return builder->AddSectionFromFile(path.c_str());
+}
diff --git a/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp b/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp
index d01b179..447a771 100644
--- a/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp
+++ b/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp
@@ -17,16 +17,21 @@ private:
     asIScriptEngine* _engine = nullptr;
     asIScriptModule* _mainModule = nullptr;
     ContextPool* _contextPool = nullptr;
-    CScriptBuilder _builder;
+    CScriptBuilder _builder = {};
+    std::string _sourceDirectory = {};
 
-    static void MessageCallback(const asSMessageInfo* msg, void* param);
-    static void Print(const std::string& str) { std::cout << str << std::endl; }
-    static void PrintConst(const ArbUt::StringView& str) { Print(std::string(str.std_str())); }
     ArbUt::Dictionary<ScriptCategory, ArbUt::Dictionary<ArbUt::StringView, AngelScriptTypeInfo*>> _typeDatabase;
     ArbUt::Dictionary<ArbUt::StringView, asITypeInfo*> _baseTypes;
     ArbUt::Dictionary<ArbUt::StringView, asITypeInfo*> _itemUseTypes;
     ArbUt::Dictionary<const CreatureLib::Library::Item*, AngelScriptItemUseScript*> _itemUseScripts;
 
+
+    static void MessageCallback(const asSMessageInfo* msg, void* param);
+    static void Print(const std::string& str) { std::cout << str << std::endl; }
+    static void PrintConst(const ArbUt::StringView& str) { Print(std::string(str.std_str())); }
+    static i32 IncludeCallback(const char *include, const char *from, CScriptBuilder *builder, void *userParam);
+
+
     void RegisterTypes();
     void
     InitializeByteCode(const ArbUt::Dictionary<ScriptCategory, ArbUt::Dictionary<ArbUt::StringView, uint32_t>>& types);
@@ -50,6 +55,7 @@ public:
     void Initialize(CreatureLib::Battling::BattleLibrary* library) override { Initialize(library, true); }
     void Initialize(CreatureLib::Battling::BattleLibrary* library, bool includeStandard);
     void CreateScript(const char* name, const char* script);
+    inline void SetSourceDirectory(std::string const& path) noexcept { _sourceDirectory = path; }
     const asIScriptModule* GetMainModule() const noexcept { return _mainModule; }
 
     void FinalizeModule();