Merge pull request #14 from DFHack/resteam

Add support for redirected DLLs using .ini files

Author: ab9rf

ChainLoader.cpp

@@ -2,35 +2,64 @@
 
 #include <filesystem>
 #include <queue>
+#include <fstream>
+#include <string_view>
 
 using std::string;
 using std::vector;
 
-#if _WIN32
-static const string search_prefix = "dfhooks_";
-static const string search_ext = ".dll";
-#else
-static const string search_prefix = "libdfhooks_";
-#  if __APPLE__
-static const string search_ext = ".dylib";
-#  else
-static const string search_ext = ".so";
-#  endif
-#endif
+using std::string_view_literals::operator""sv;
+
+// look for libraries with these prefixes and extensions.
+// .ini files are expected to contain a single line with the path to the library to load, which allows for loading libraries with non-standard names or from non-standard locations
+// note that if both a .ini and a .dll/.so are present _both_ will be loaded, which may have unintended consequences if they point to the same library, so use with caution
+
+constexpr auto search_prefix = {"dfhooks_"sv, "libdfhooks_"sv};
+constexpr auto search_ext = {".dll"sv, ".so"sv};
+constexpr auto search_ext_ini = ".ini"sv;
 
 ChainLoader::ChainLoader() {
     auto cmp = [](const LibWrapper * lhs, const LibWrapper * rhs) { return lhs->priority < rhs->priority; };
     std::priority_queue<LibWrapper*, vector<LibWrapper*>, decltype(cmp)> priq(cmp);
 
-    for (auto const& dir_entry : std::filesystem::directory_iterator{"."}) {
+    for (auto const& dir_entry : std::filesystem::directory_iterator{"."})
+    {
         auto fname = dir_entry.path().filename();
-        if (fname.extension() != search_ext || fname.stem().string().find(search_prefix) != 0)
-            continue;
-        auto wrapper = new LibWrapper(fname.string());
-        if (wrapper->handle)
-            priq.emplace(wrapper);
-        else
-            delete wrapper;
+        bool fnd = false;
+        for (auto& prefix : search_prefix) {
+            if (fname.stem().string().find(prefix) == 0) {
+                fnd = true;
+                break;
+            }
+        }
+        if (!fnd) continue;
+
+        if (fname.extension() == search_ext_ini)
+        {
+            std::string indirection;
+            std::ifstream ini_file{fname};
+            std::getline(ini_file, indirection);
+            if (!indirection.empty())
+            {
+                std::filesystem::path lib{indirection};
+                auto wrapper = new LibWrapper(lib);
+                if (wrapper->handle)
+                    priq.emplace(wrapper);
+                else
+                    delete wrapper;
+            }
+        }
+        else for (auto& ext : search_ext)
+        {
+            if (fname.extension() == ext)
+            {
+                auto wrapper = new LibWrapper(fname);
+                if (wrapper->handle)
+                    priq.emplace(wrapper);
+                else
+                    delete wrapper;
+            }
+        }
     }
     while (!priq.empty()) {
         auto lib = priq.top();

LibWrapper.cpp

@@ -1,16 +1,25 @@
 #include "LibWrapper.h"
 
+#include <filesystem>
+#include <iostream>
+
 using std::string;
+using std::filesystem::path;
 
 #if _WIN32
 #  include <windows.h>
+#  include <libloaderapi.h>
 #else
 #  include <dlfcn.h>
 #endif
 
-static void* open_library(const string& fname) {
+static void* open_library(std::filesystem::path fname) {
 #if _WIN32
-    return LoadLibrary(fname.c_str());
+    std::filesystem::path path = std::filesystem::canonical(fname);
+    std::filesystem::path ppath = path.parent_path();
+    [[maybe_unused]] DLL_DIRECTORY_COOKIE cookie = AddDllDirectory(ppath.wstring().c_str());
+    auto handle = LoadLibraryExW(path.filename().wstring().c_str(), NULL, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
+    return handle;
 #else
     return dlopen(fname.c_str(), RTLD_LAZY);
 #endif
@@ -35,21 +44,34 @@ static void* load_sym(void* handle, const char* sym) {
 #endif
 }
 
-LibWrapper::LibWrapper(const string& fname) {
+extern "C" void dfhooks_init();
+
+LibWrapper::LibWrapper(const std::filesystem::path& fname) {
     handle = open_library(fname);
     if (!handle)
         return;
 
     int32_t * priority_sym = (int *)load_sym(handle, "dfhooks_priority");
     if (priority_sym) priority = *priority_sym;
 
-    init = (dfhooks_init_fn)load_sym(handle, "dfhooks_init");
-    shutdown = (dfhooks_shutdown_fn)load_sym(handle, "dfhooks_shutdown");
-    update = (dfhooks_update_fn)load_sym(handle, "dfhooks_update");
-    prerender = (dfhooks_prerender_fn)load_sym(handle, "dfhooks_prerender");
-    sdl_event = (dfhooks_sdl_event_fn)load_sym(handle, "dfhooks_sdl_event");
-    sdl_loop = (dfhooks_sdl_loop_fn)load_sym(handle, "dfhooks_sdl_loop");
-    ncurses_key = (dfhooks_ncurses_key_fn)load_sym(handle, "dfhooks_ncurses_key");
+    preinit = (dfhooks_preinit_fn*)load_sym(handle, "dfhooks_preinit");
+    init = (dfhooks_init_fn*)load_sym(handle, "dfhooks_init");
+    shutdown = (dfhooks_shutdown_fn*)load_sym(handle, "dfhooks_shutdown");
+    update = (dfhooks_update_fn*)load_sym(handle, "dfhooks_update");
+    prerender = (dfhooks_prerender_fn*)load_sym(handle, "dfhooks_prerender");
+    sdl_event = (dfhooks_sdl_event_fn*)load_sym(handle, "dfhooks_sdl_event");
+    sdl_loop = (dfhooks_sdl_loop_fn*)load_sym(handle, "dfhooks_sdl_loop");
+    ncurses_key = (dfhooks_ncurses_key_fn*)load_sym(handle, "dfhooks_ncurses_key");
+
+    if (preinit)
+        preinit(fname);
+
+    if (init == &dfhooks_init)
+    {
+        // if the library exports the same init function as the main dfhooks, initing the library will can an infinite recursion. disallow this.
+        close_library(handle);
+        handle = nullptr;
+    }
 }
 
 LibWrapper::~LibWrapper() {

LibWrapper.h

@@ -1,28 +1,31 @@
 #pragma once
 
 #include <string>
+#include <filesystem>
 
 union SDL_Event;
 
-typedef void (*dfhooks_init_fn)();
-typedef void (*dfhooks_shutdown_fn)();
-typedef void (*dfhooks_update_fn)();
-typedef void (*dfhooks_prerender_fn)();
-typedef bool (*dfhooks_sdl_event_fn)(SDL_Event* event);
-typedef void (*dfhooks_sdl_loop_fn)();
-typedef bool (*dfhooks_ncurses_key_fn)(int key);
+using dfhooks_preinit_fn = auto (std::filesystem::path) -> void;
+using dfhooks_init_fn = auto () -> void;
+using dfhooks_shutdown_fn = auto () -> void;
+using dfhooks_update_fn = auto () -> void;
+using dfhooks_prerender_fn = auto () -> void;
+using dfhooks_sdl_event_fn = auto (SDL_Event* event) -> bool;
+using dfhooks_sdl_loop_fn = auto () -> void;
+using dfhooks_ncurses_key_fn = auto (int key) -> bool;
 
 struct LibWrapper {
     void* handle = nullptr;
     int32_t priority = 0;
-    dfhooks_init_fn init = nullptr;
-    dfhooks_shutdown_fn shutdown = nullptr;
-    dfhooks_update_fn update = nullptr;
-    dfhooks_prerender_fn prerender = nullptr;
-    dfhooks_sdl_event_fn sdl_event = nullptr;
-    dfhooks_sdl_loop_fn sdl_loop = nullptr;
-    dfhooks_ncurses_key_fn ncurses_key = nullptr;
+    dfhooks_preinit_fn* preinit = nullptr;
+    dfhooks_init_fn* init = nullptr;
+    dfhooks_shutdown_fn* shutdown = nullptr;
+    dfhooks_update_fn* update = nullptr;
+    dfhooks_prerender_fn* prerender = nullptr;
+    dfhooks_sdl_event_fn* sdl_event = nullptr;
+    dfhooks_sdl_loop_fn* sdl_loop = nullptr;
+    dfhooks_ncurses_key_fn* ncurses_key = nullptr;
 
-    LibWrapper(const std::string& fname);
+    LibWrapper(const std::filesystem::path& fname);
     virtual ~LibWrapper();
 };