All Resource Redirect - Experimental (https://discord.com/channels/923944342991818753/923954164583759873/1233458845070790687)

all-resource-redirect-experimental.wh.cpp

// ==WindhawkMod==
// @id              all-resource-redirect-experimental
// @name            All Resource Redirect - Experimental
// @description     Define alternative files for loading resources (e.g. instead of imageres.dll) for simple theming without having to modify system files
// @version         1.0
// @author          m417z
// @github          https://github.com/m417z
// @twitter         https://twitter.com/m417z
// @homepage        https://m417z.com/
// @include         *
// ==/WindhawkMod==

// ==WindhawkModReadme==
/*
# All Resource Redirect

Define alternative  files for loading resources (e.g. instead of imageres.dll)
for simple theming without having to modify system files.
*/
// ==/WindhawkModReadme==

// ==WindhawkModSettings==
/*
- redirectionResourcePaths:
  - - original: 'C:\Windows\System32\imageres.dll'
      $name: The original resource file
      $description: The original file from which icons are loaded
    - redirect: 'C:\my-themes\theme-1\imageres.dll'
      $name: The custom resource file
      $description: The custom resource file that will be used instead
  $name: Redirection resource paths
*/
// ==/WindhawkModSettings==

#include <windhawk_utils.h>

#include <psapi.h>
#include <shlobj.h>

#include <functional>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <vector>

std::shared_mutex g_redirectionResourcePathsMutex;
std::unordered_map<std::wstring, std::vector<std::wstring>>
    g_redirectionResourcePaths;

std::shared_mutex g_redirectionResourceModulesMutex;
std::unordered_map<std::wstring, HMODULE> g_redirectionResourceModules;

bool DevicePathToDosPath(const WCHAR* device_path,
                         WCHAR* dos_path,
                         size_t dos_path_size) {
    WCHAR drive_strings[MAX_PATH];
    if (!GetLogicalDriveStrings(ARRAYSIZE(drive_strings), drive_strings)) {
        return false;
    }

    // Drive strings are stored as a set of null terminated strings, with an
    // extra null after the last string. Each drive string is of the form "C:\".
    // We convert it to the form "C:", which is the format expected by
    // QueryDosDevice().
    WCHAR drive_colon[3] = L" :";
    for (const WCHAR* next_drive_letter = drive_strings; *next_drive_letter;
         next_drive_letter += wcslen(next_drive_letter) + 1) {
        // Dos device of the form "C:".
        *drive_colon = *next_drive_letter;
        WCHAR device_name[MAX_PATH];
        if (!QueryDosDevice(drive_colon, device_name, ARRAYSIZE(device_name))) {
            continue;
        }

        size_t name_length = wcslen(device_name);
        if (_wcsnicmp(device_path, device_name, name_length) == 0) {
            size_t dos_path_size_required =
                2 + wcslen(device_path + name_length) + 1;
            if (dos_path_size < dos_path_size_required) {
                return false;
            }

            // Construct DOS path.
            wcscpy(dos_path, drive_colon);
            wcscat(dos_path, device_path + name_length);
            return true;
        }
    }

    return false;
}

HMODULE GetRedirectedModule(std::wstring_view fileName) {
    std::wstring fileNameUpper{fileName};
    LCMapStringEx(LOCALE_NAME_USER_DEFAULT, LCMAP_UPPERCASE, &fileNameUpper[0],
                  static_cast<int>(fileNameUpper.length()), &fileNameUpper[0],
                  static_cast<int>(fileNameUpper.length()), nullptr, nullptr,
                  0);

    {
        std::shared_lock lock{g_redirectionResourceModulesMutex};
        const auto it = g_redirectionResourceModules.find(fileNameUpper);
        if (it != g_redirectionResourceModules.end()) {
            return it->second;
        }
    }

    HINSTANCE module = LoadLibraryEx(
        fileNameUpper.c_str(), nullptr,
        LOAD_LIBRARY_AS_DATAFILE | LOAD_LIBRARY_AS_IMAGE_RESOURCE);
    if (!module) {
        DWORD dwError = GetLastError();
        Wh_Log(L"LoadLibraryEx failed with error %u", dwError);
        return nullptr;
    }

    {
        std::unique_lock lock{g_redirectionResourceModulesMutex};
        g_redirectionResourceModules.try_emplace(fileNameUpper, module);
    }

    return module;
}

void FreeAndClearRedirectedModules() {
    std::unordered_map<std::wstring, HMODULE> modules;

    {
        std::unique_lock lock{g_redirectionResourceModulesMutex};
        modules.swap(g_redirectionResourceModules);
    }

    for (const auto& it : modules) {
        FreeLibrary(it.second);
    }
}

bool RedirectModule(HINSTANCE hInstance,
                    std::function<void()> beforeFirstRedirectionFunction,
                    std::function<bool(HINSTANCE)> redirectFunction) {
    WCHAR szFileName[MAX_PATH];
    DWORD fileNameLen;
    if ((ULONG_PTR)hInstance & 3) {
        WCHAR szNtFileName[MAX_PATH * 2];

        if (!GetMappedFileName(GetCurrentProcess(), (void*)hInstance,
                               szNtFileName, ARRAYSIZE(szNtFileName))) {
            DWORD dwError = GetLastError();
            Wh_Log(L"> GetMappedFileName(%p) failed with error %u", hInstance,
                   dwError);
            return false;
        }

        if (!DevicePathToDosPath(szNtFileName, szFileName,
                                 ARRAYSIZE(szFileName))) {
            Wh_Log(L"> DevicePathToDosPath failed");
            return false;
        }

        fileNameLen = wcslen(szFileName);
    } else {
        fileNameLen =
            GetModuleFileName(hInstance, szFileName, ARRAYSIZE(szFileName));
        switch (fileNameLen) {
            case 0: {
                DWORD dwError = GetLastError();
                Wh_Log(L"> GetModuleFileName(%p) failed with error %u",
                       hInstance, dwError);
                return false;
            }

            case ARRAYSIZE(szFileName):
                Wh_Log(L"> GetModuleFileName(%p) failed, name too long",
                       hInstance);
                return false;
        }
    }

    Wh_Log(L"> Module: %s", szFileName);

    LCMapStringEx(LOCALE_NAME_USER_DEFAULT, LCMAP_UPPERCASE, &szFileName[0],
                  fileNameLen, &szFileName[0], fileNameLen, nullptr, nullptr,
                  0);

    bool triedRedirection = false;

    {
        std::shared_lock lock{g_redirectionResourcePathsMutex};

        if (const auto it = g_redirectionResourcePaths.find(szFileName);
            it != g_redirectionResourcePaths.end()) {
            const auto& redirects = it->second;
            for (const auto& redirect : redirects) {
                if (!triedRedirection) {
                    beforeFirstRedirectionFunction();
                    triedRedirection = true;
                }

                Wh_Log(L"Trying %s", redirect.c_str());

                HINSTANCE hInstanceRedirect = GetRedirectedModule(redirect);
                if (!hInstanceRedirect) {
                    Wh_Log(L"GetRedirectedModule failed");
                    continue;
                }

                if (redirectFunction(hInstanceRedirect)) {
                    return true;
                }
            }
        }
    }

    if (triedRedirection) {
        Wh_Log(L"No redirection succeeded, falling back to original");
    }

    return false;
}

using LdrpSearchResourceSection_U_t = NTSTATUS(WINAPI*)(HINSTANCE hInstance,
                                                        PVOID param2,
                                                        PVOID param3,
                                                        PVOID param4,
                                                        PVOID param5);
LdrpSearchResourceSection_U_t LdrpSearchResourceSection_U_Original;
NTSTATUS WINAPI LdrpSearchResourceSection_U_Hook(HINSTANCE hInstance,
                                                 PVOID param2,
                                                 PVOID param3,
                                                 PVOID param4,
                                                 PVOID param5) {
    Wh_Log(L">");

    NTSTATUS result;

    bool redirected = RedirectModule(
        hInstance, [&]() {},
        [&](HINSTANCE hInstanceRedirect) {
            result = LdrpSearchResourceSection_U_Original(
                hInstanceRedirect, param2, param3, param4, param5);
            if (result >= 0) {
                Wh_Log(L"Redirected successfully");
                return true;
            }

            Wh_Log(L"LdrpSearchResourceSection_U failed with error %08X",
                   result);
            return false;
        });
    if (redirected) {
        return result;
    }

    return LdrpSearchResourceSection_U_Original(hInstance, param2, param3,
                                                param4, param5);
}

using LdrpAccessResourceDataNoMultipleLanguage_t = NTSTATUS(
    WINAPI*)(HINSTANCE hInstance, PVOID param2, PVOID param3, PVOID param4);
LdrpAccessResourceDataNoMultipleLanguage_t
    LdrpAccessResourceDataNoMultipleLanguage_Original;
NTSTATUS WINAPI
LdrpAccessResourceDataNoMultipleLanguage_Hook(HINSTANCE hInstance,
                                              PVOID param2,
                                              PVOID param3,
                                              PVOID param4) {
    Wh_Log(L">");

    NTSTATUS result;

    bool redirected = RedirectModule(
        hInstance, [&]() {},
        [&](HINSTANCE hInstanceRedirect) {
            result = LdrpAccessResourceDataNoMultipleLanguage_Original(
                hInstanceRedirect, param2, param3, param4);
            if (result >= 0) {
                Wh_Log(L"Redirected successfully");
                return true;
            }

            Wh_Log(
                L"LdrpAccessResourceDataNoMultipleLanguage failed with error "
                L"%08X",
                result);
            return false;
        });
    if (redirected) {
        return result;
    }

    return LdrpAccessResourceDataNoMultipleLanguage_Original(hInstance, param2,
                                                             param3, param4);
}

void LoadSettings() {
    std::unordered_map<std::wstring, std::vector<std::wstring>> paths;

    for (int i = 0;; i++) {
        PCWSTR original =
            Wh_GetStringSetting(L"redirectionResourcePaths[%d].original", i);
        PCWSTR redirect =
            Wh_GetStringSetting(L"redirectionResourcePaths[%d].redirect", i);

        bool hasName = *original || *redirect;

        if (hasName) {
            std::wstring originalUpper{original};
            LCMapStringEx(
                LOCALE_NAME_USER_DEFAULT, LCMAP_UPPERCASE, &originalUpper[0],
                static_cast<int>(originalUpper.length()), &originalUpper[0],
                static_cast<int>(originalUpper.length()), nullptr, nullptr, 0);

            paths[originalUpper].push_back(redirect);
        }

        Wh_FreeStringSetting(original);
        Wh_FreeStringSetting(redirect);

        if (!hasName) {
            break;
        }
    }

    std::unique_lock lock{g_redirectionResourcePathsMutex};
    g_redirectionResourcePaths = std::move(paths);
}

BOOL Wh_ModInit() {
    Wh_Log(L">");

    LoadSettings();

    HMODULE ntdllModule = GetModuleHandle(L"ntdll.dll");
    if (!ntdllModule) {
        Wh_Log(L"Couldn't load ntdll.dll");
        return FALSE;
    }

    WindhawkUtils::SYMBOL_HOOK symbolHooks[] = {
        {
            {L"LdrpSearchResourceSection_U"},
            (void**)&LdrpSearchResourceSection_U_Original,
            (void*)LdrpSearchResourceSection_U_Hook,
            true,
        },
        {
            {L"LdrpAccessResourceDataNoMultipleLanguage"},
            (void**)&LdrpAccessResourceDataNoMultipleLanguage_Original,
            (void*)LdrpAccessResourceDataNoMultipleLanguage_Hook,
            true,
        },
    };

    if (!HookSymbols(ntdllModule, symbolHooks, ARRAYSIZE(symbolHooks))) {
        return FALSE;
    }

    return TRUE;
}

void Wh_ModUninit() {
    Wh_Log(L">");

    FreeAndClearRedirectedModules();
}

void Wh_ModSettingsChanged() {
    Wh_Log(L">");

    LoadSettings();

    FreeAndClearRedirectedModules();

    // Invalidate icon cache.
    SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, nullptr, nullptr);
}