implement Vulkan loader with symbol overrides

include/loader/vk.hpp

@@ -0,0 +1,58 @@
+#ifndef VK_HPP
+#define VK_HPP
+
+#include <vulkan/vulkan.h>
+
+#include <string>
+
+//
+// Similar to the dynamic loader, the Vulkan loader replaces the standard
+// vkGetInstanceProcAddr and vkGetDeviceProcAddr functions.
+//
+// One thing that should be noted, is that not every application uses the
+// Vulkan loader for every method. On Linux it's not unusual to see dlsym
+// being used for vulkan functions, so make sure to register the same
+// symbol on both loaders.
+//
+
+namespace Loader::VK {
+
+    ///
+    /// Initialize the Vulkan loader.
+    ///
+    void initialize();
+
+    ///
+    /// Register a symbol to the Vulkan loader.
+    ///
+    /// @param symbol The name of the Vulkan function to override.
+    /// @param address The address of the Vulkan function.
+    ///
+    void registerSymbol(const std::string& symbol, void* address);
+
+    ///
+    /// Call the original vkGetInstanceProcAddr function.
+    ///
+    /// @param instance The (optional) Vulkan instance.
+    /// @param pName The name of the function to retrieve.
+    /// @return The address of the function, or nullptr if not found.
+    ///
+    PFN_vkVoidFunction ovkGetInstanceProcAddr(VkInstance instance, const char* pName);
+
+    ///
+    /// Call the original vkGetDeviceProcAddr function.
+    ///
+    /// @param device The Vulkan device.
+    /// @param pName The name of the function to retrieve.
+    /// @return The address of the function, or nullptr if not found.
+    ///
+    PFN_vkVoidFunction ovkGetDeviceProcAddr(VkDevice device, const char* pName);
+
+}
+
+/// Modified version of the vkGetInstanceProcAddr function.
+extern "C" PFN_vkVoidFunction myvkGetInstanceProcAddr(VkInstance instance, const char* pName);
+/// Modified version of the vkGetDeviceProcAddr function.
+extern "C" PFN_vkVoidFunction myvkGetDeviceProcAddr(VkDevice device, const char* pName);
+
+#endif // VK_HPP

src/init.cpp

@@ -1,4 +1,5 @@
 #include "loader/dl.hpp"
+#include "loader/vk.hpp"
 #include "log.hpp"
 
 extern "C" void __attribute__((constructor)) init();
@@ -9,6 +10,7 @@ void init() {
 
     // hook loaders
     Loader::DL::initialize();
+    Loader::VK::initialize();
 }
 
 void deinit() {

src/loader/vk.cpp

@@ -0,0 +1,109 @@
+#include "loader/vk.hpp"
+#include "loader/dl.hpp"
+#include "log.hpp"
+
+using namespace Loader;
+
+namespace {
+    // original function pointers
+    PFN_vkGetInstanceProcAddr vkGetInstanceProcAddr_ptr{};
+    PFN_vkGetDeviceProcAddr   vkGetDeviceProcAddr_ptr{};
+
+    // map of all overridden symbols
+    auto& symbols() {
+        static std::unordered_map<std::string, void*> symbols;
+        return symbols;
+    }
+}
+
+void VK::initialize() {
+    if (vkGetInstanceProcAddr_ptr || vkGetDeviceProcAddr_ptr) {
+        Log::warn("lsfg-vk(vk): Vulkan loader already initialized, did you call it twice?");
+        return;
+    }
+
+    // get original function pointers
+    auto* handle = DL::odlopen("libvulkan.so.1", 0x2);
+    vkGetInstanceProcAddr_ptr =
+        reinterpret_cast<PFN_vkGetInstanceProcAddr>(DL::odlsym(handle, "vkGetInstanceProcAddr"));
+    vkGetDeviceProcAddr_ptr =
+        reinterpret_cast<PFN_vkGetDeviceProcAddr>  (DL::odlsym(handle, "vkGetDeviceProcAddr"));
+    if (!vkGetInstanceProcAddr_ptr || !vkGetDeviceProcAddr_ptr) {
+        Log::error("lsfg-vk(vk): Failed to initialize Vulkan loader, missing symbols");
+        exit(EXIT_FAILURE);
+    }
+
+    // register dynamic loader overrides
+    DL::File vulkanLib{"libvulkan.so.1"};
+    vulkanLib.defineSymbol("vkGetInstanceProcAddr",
+        reinterpret_cast<void*>(myvkGetInstanceProcAddr));
+    vulkanLib.defineSymbol("vkGetDeviceProcAddr",
+        reinterpret_cast<void*>(myvkGetDeviceProcAddr));
+    DL::registerFile(vulkanLib);
+
+    DL::File vulkanLib2{"libvulkan.so"};
+    vulkanLib2.defineSymbol("vkGetInstanceProcAddr",
+        reinterpret_cast<void*>(myvkGetInstanceProcAddr));
+    vulkanLib2.defineSymbol("vkGetDeviceProcAddr",
+        reinterpret_cast<void*>(myvkGetDeviceProcAddr));
+    DL::registerFile(vulkanLib2);
+
+    // register vulkan loader overrides
+    VK::registerSymbol("vkGetInstanceProcAddr", reinterpret_cast<void*>(myvkGetInstanceProcAddr));
+    VK::registerSymbol("vkGetDeviceProcAddr", reinterpret_cast<void*>(myvkGetDeviceProcAddr));
+
+    Log::debug("lsfg-vk(vk): Initialized Vulkan loader with original functions");
+}
+
+void VK::registerSymbol(const std::string& symbol, void* address) {
+    auto& syms = symbols();
+
+    const auto it = syms.find(symbol);
+    if (it != syms.end()) {
+        Log::warn("lsfg-vk(vk): Tried registering symbol {}, but it is already defined", symbol);
+        return;
+    }
+
+    syms.emplace(symbol, address);
+}
+
+PFN_vkVoidFunction myvkGetInstanceProcAddr(VkInstance instance, const char* pName) {
+    const auto& syms = symbols();
+
+    if (!pName)
+        return vkGetInstanceProcAddr_ptr(instance, pName);
+
+    // try to find an override
+    const std::string pName_str(pName);
+    const auto it = syms.find(pName_str);
+    if (it == syms.end())
+        return vkGetInstanceProcAddr_ptr(instance, pName);
+
+    Log::debug("lsfg-vk(vk): Intercepted Vulkan symbol {}", pName_str);
+    return reinterpret_cast<PFN_vkVoidFunction>(it->second);
+}
+
+PFN_vkVoidFunction myvkGetDeviceProcAddr(VkDevice device, const char* pName) {
+    const auto& syms = symbols();
+
+    if (!pName)
+        return vkGetDeviceProcAddr_ptr(device, pName);
+
+    const std::string pName_str(pName);
+    auto it = syms.find(pName_str);
+    if (it == syms.end())
+        return vkGetDeviceProcAddr_ptr(device, pName);
+
+    Log::debug("lsfg-vk(vk): Intercepted Vulkan symbol {}", pName_str);
+    return reinterpret_cast<PFN_vkVoidFunction>(it->second);
+}
+
+// original function calls
+
+PFN_vkVoidFunction VK::ovkGetInstanceProcAddr(VkInstance instance, const char* pName) {
+    return vkGetInstanceProcAddr_ptr(instance, pName);
+}
+
+PFN_vkVoidFunction VK::ovkGetDeviceProcAddr(VkDevice device, const char* pName) {
+    return vkGetDeviceProcAddr_ptr(device, pName);
+}