implement dynamic laoder with overrides

CMakeLists.txt

@@ -19,6 +19,7 @@ project(lsfg-vk
 add_subdirectory(lsfg-vk-gen)
 
 file(GLOB SOURCES
+    "src/loader/*.cpp"
     "src/*.cpp"
 )
 

include/loader/dl.hpp

@@ -0,0 +1,133 @@
+#ifndef DL_HPP
+#define DL_HPP
+
+#include <string>
+#include <unordered_map>
+
+//
+// This dynamic loader replaces the standard dlopen, dlsym, and dlclose functions.
+// On initialize, the original functions are obtained via dlvsym (glibc exclusive)
+// and made available under functions with the "o" prefix.
+//
+// Any call to regular dlopen, dlsym or dlclose is intercepted and may be
+// overriden by registering a File override via `Loader::DL::registerFile`.
+//
+
+namespace Loader::DL {
+
+    /// Dynamic loader override structure.
+    class File {
+    public:
+        ///
+        /// Create a dynamic loader override for a specific file.
+        ///
+        /// @param filename The name of the file to override.
+        ///
+        File(std::string filename)
+            : filename(std::move(filename)) {}
+
+        ///
+        /// Append a symbol to the dynamic loader override.
+        ///
+        /// @param symbol The name of the symbol to add.
+        /// @param address The address of the symbol.
+        ///
+        void defineSymbol(const std::string& symbol, void* address) {
+            symbols[symbol] = address;
+        }
+
+        /// Get the filename
+        [[nodiscard]] const std::string& getFilename() const { return filename; }
+        /// Get all overriden symbols
+        [[nodiscard]] const std::unordered_map<std::string, void*>& getSymbols() const { return symbols; }
+
+        // Find a specific symbol
+        [[nodiscard]] void* findSymbol(const std::string& symbol) const {
+            auto it = symbols.find(symbol);
+            return (it != symbols.end()) ? it->second : nullptr;
+        }
+
+        /// Get the fake handle
+        [[nodiscard]] void* getHandle() const { return handle; }
+        /// Get the real handle
+        [[nodiscard]] void* getOriginalHandle() const { return handle_orig; }
+
+        /// Set the fake handle
+        void setHandle(void* new_handle) { handle = new_handle; }
+        /// Set the real handle
+        void setOriginalHandle(void* new_handle) { handle_orig = new_handle; }
+
+        /// Copyable, moveable, default destructor
+        File(const File&) = default;
+        File(File&&) = default;
+        File& operator=(const File&) = default;
+        File& operator=(File&&) = default;
+        ~File() = default;
+    private:
+        std::string filename;
+        std::unordered_map<std::string, void*> symbols;
+
+        void* handle = nullptr;
+        void* handle_orig = nullptr;
+    };
+
+    ///
+    /// Initialize the dynamic loader
+    ///
+    void initialize();
+
+    ///
+    /// Register a file with the dynamic loader.
+    ///
+    /// @param file The file to register.
+    ///
+    void registerFile(const File& file);
+
+    ///
+    /// Disable hooks temporarily. This may be useful
+    /// when loading third-party libraries you wish not
+    /// to hook.
+    ///
+    void disableHooks();
+
+    ///
+    /// Re-enable hooks after they were disabled.
+    ///
+    void enableHooks();
+
+    ///
+    /// Call the original dlopen function.
+    ///
+    /// @param filename The name of the file to open.
+    /// @param flag The flags to use when opening the file.
+    /// @return A handle to the opened file, or NULL on failure.
+    ///
+    void* odlopen(const char* filename, int flag);
+
+    ///
+    /// Call the original dlsym function.
+    ///
+    /// @param handle The handle to the opened file.
+    /// @param symbol The name of the symbol to look up.
+    /// @return A pointer to the symbol, or NULL on failure.
+    ///
+    void* odlsym(void* handle, const char* symbol);
+
+    ///
+    /// Call the original dlclose function.
+    ///
+    /// @param handle The handle to the opened file.
+    /// @return 0 on success, or -1 on failure.
+    ///
+    int odlclose(void* handle);
+
+}
+
+/// Modified version of the dlopen function.
+extern "C" void* dlopen(const char* filename, int flag);
+/// Modified version of the dlsym function.
+extern "C" void* dlsym(void* handle, const char* symbol);
+/// Modified version of the dlclose function.
+extern "C" int dlclose(void* handle);
+
+#endif // DL_HPP

src/init.cpp

@@ -1,3 +1,4 @@
+#include "loader/dl.hpp"
 #include "log.hpp"
 
 extern "C" void __attribute__((constructor)) init();
@@ -5,6 +6,9 @@ extern "C" [[noreturn]] void __attribute__((destructor)) deinit();
 
 void init() {
     Log::info("lsfg-vk: init() called");
+
+    // hook loaders
+    Loader::DL::initialize();
 }
 
 void deinit() {

src/loader/dl.cpp

@@ -0,0 +1,182 @@
+#include "loader/dl.hpp"
+#include "log.hpp"
+
+#include <vector>
+
+using namespace Loader;
+
+using dlopen_t  = void* (*)(const char*, int);
+using dlsym_t   = void* (*)(void*, const char*);
+using dlclose_t = int   (*)(void*);
+
+// glibc exclusive function to get versioned symbols
+extern "C" void* dlvsym(long, const char*, const char*);
+
+namespace {
+    // original function pointers
+    dlopen_t  dlopen_ptr;
+    dlsym_t   dlsym_ptr;
+    dlclose_t dlclose_ptr;
+
+    // map of all registered overrides
+    auto& overrides() {
+        // this has to be a function rather than a static variable
+        // because of weird initialization order issues.
+        static std::unordered_map<std::string, DL::File> overrides;
+        return overrides;
+    }
+
+    // vector of loaded handles
+    auto& handles() {
+        static std::vector<void*> handles;
+        return handles;
+    }
+
+    bool enable_hooks{true};
+}
+
+void DL::initialize() {
+    if (dlopen_ptr || dlsym_ptr || dlclose_ptr) {
+        Log::warn("lsfg-vk(dl): Dynamic loader already initialized, did you call it twice?");
+        return;
+    }
+
+    dlopen_ptr  = reinterpret_cast<dlopen_t> (dlvsym(-1, "dlopen",  "GLIBC_2.2.5"));
+    dlsym_ptr   = reinterpret_cast<dlsym_t>  (dlvsym(-1, "dlsym",   "GLIBC_2.2.5"));
+    dlclose_ptr = reinterpret_cast<dlclose_t>(dlvsym(-1, "dlclose", "GLIBC_2.2.5"));
+    if (!dlopen_ptr || !dlsym_ptr || !dlclose_ptr) {
+        Log::error("lsfg-vk(dl): Failed to initialize dynamic loader, missing symbols");
+        exit(EXIT_FAILURE);
+    }
+
+    Log::debug("lsfg-vk(dl): Initialized dynamic loader with original functions");
+}
+
+void DL::registerFile(const File& file) {
+    auto& files = overrides();
+
+    auto it = files.find(file.getFilename());
+    if (it == files.end()) {
+        // simply register if the file hasn't been registered yet
+        files.emplace(file.getFilename(), file);
+        return;
+    }
+
+    // merge the new file's symbols into the previously registered one
+    auto& existing_file = it->second;
+    for (const auto& [symbol, func] : file.getSymbols())
+        if (existing_file.findSymbol(symbol) == nullptr)
+            existing_file.defineSymbol(symbol, func);
+        else
+            Log::warn("lsfg-vk(dl): Tried registering symbol {}::{}, but it is already defined",
+                existing_file.getFilename(), symbol);
+}
+
+void DL::disableHooks() { enable_hooks = false; }
+void DL::enableHooks()  { enable_hooks = true; }
+
+extern "C" void* dlopen(const char* filename, int flag) {
+    auto& files = overrides();
+    auto& loaded = handles();
+
+    // ALWAYS load the library and ensure it's tracked
+    auto* handle = dlopen_ptr(filename, flag);
+    if (handle && std::ranges::find(loaded, handle) == loaded.end())
+        loaded.push_back(handle);
+
+    // no need to check for overrides if hooks are disabled
+    if (!enable_hooks || !filename)
+        return handle;
+
+    // try to find an override for this filename
+    const std::string filename_str(filename);
+    auto it = files.find(filename_str);
+    if (it == files.end())
+        return handle;
+
+    auto& file = it->second;
+    file.setOriginalHandle(handle);
+    file.setHandle(reinterpret_cast<void*>(&file));
+
+    Log::debug("lsfg-vk(dl): Intercepted module load for {}", file.getFilename());
+    return file.getHandle();
+}
+
+extern "C" void* dlsym(void* handle, const char* symbol) {
+    const auto& files = overrides();
+
+    if (!enable_hooks || !handle || !symbol)
+        return dlsym_ptr(handle, symbol);
+
+    // see if handle is a fake one
+    const auto it = std::ranges::find_if(files, [handle](const auto& pair) {
+        return pair.second.getHandle() == handle;
+    });
+    if (it == files.end())
+        return dlsym_ptr(handle, symbol);
+    const auto& file = it->second;
+
+    // find a symbol override
+    const std::string symbol_str(symbol);
+    auto* func = file.findSymbol(symbol_str);
+    if (func == nullptr)
+        return dlsym_ptr(file.getOriginalHandle(), symbol);
+
+    Log::debug("lsfg-vk(dl): Intercepted symbol {}::{}", file.getFilename(), symbol_str);
+    return func;
+}
+
+extern "C" int dlclose(void* handle) {
+    auto& files = overrides();
+    auto& loaded = handles();
+
+    // no handle, let the original dlclose handle it
+    if (!handle)
+        return dlclose_ptr(handle);
+
+    // see if the handle is a fake one
+    auto it = std::ranges::find_if(files, [handle](const auto& pair) {
+        return pair.second.getHandle() == handle;
+    });
+    if (it == files.end()) {
+        // if the handle is not fake, check if it's still loaded.
+        // this is necessary to avoid double closing when
+        // one handle was acquired while hooks were disabled
+        auto l_it = std::ranges::find(loaded, handle);
+        if (l_it == loaded.end())
+            return 0;
+        loaded.erase(l_it);
+        return dlclose_ptr(handle);
+    }
+
+    auto& file = it->second;
+    handle = file.getOriginalHandle();
+    file.setHandle(nullptr);
+    file.setOriginalHandle(nullptr);
+
+    // similarly, if it is fake, check if it's still loaded
+    // before unloading it again.
+    auto l_it = std::ranges::find(loaded, handle);
+    if (l_it == loaded.end()) {
+        Log::debug("lsfg-vk(dl): Skipping unload for {} (already unloaded)", file.getFilename());
+        return 0;
+    }
+    loaded.erase(l_it);
+
+    Log::debug("lsfg-vk(dl): Unloaded {}", file.getFilename());
+    return dlclose_ptr(handle);
+}
+
+// original function calls
+
+void* DL::odlopen(const char* filename, int flag) {
+    return dlopen_ptr(filename, flag);
+}
+
+void* DL::odlsym(void* handle, const char* symbol) {
+    return dlsym_ptr(handle, symbol);
+}
+
+int DL::odlclose(void* handle) {
+    return dlclose_ptr(handle);
+}