fix(libkernel): reprotect bundled guest module images on Windows

core/libs/prx/libc/include/GuestAllocations.hpp

@@ -25,6 +25,10 @@ using Lease = std::vector<std::shared_ptr<const Range>>;
 extern "C" {
 void* GuestAllocationsBegin_nid_postfix();
 void GuestAllocationsRegisterMainImage_nid_postfix(void* mutation);
+#ifdef _WIN32
+void GuestAllocationsRegisterImage_nid_postfix(void* mutation, const void* image);
+void GuestAllocationsUnregisterImage_nid_postfix(void* mutation, const void* image);
+#endif
 void GuestAllocationsEnd_nid_postfix(void* mutation) noexcept;
 void GuestAllocationsAdd_nid_postfix(void* mutation, void* pointer, std::size_t bytes, bool readable, bool writable);
 void GuestAllocationsRequireUnpinned_nid_postfix(void* mutation, const void* pointer, std::size_t bytes);
@@ -58,6 +62,10 @@ class Mutation {
     Mutation(const Mutation&) = delete;
     Mutation& operator=(const Mutation&) = delete;
     void RegisterMainImage() { GuestAllocationsRegisterMainImage_nid_postfix(handle); }
+#ifdef _WIN32
+    void RegisterImage(const void* image) { GuestAllocationsRegisterImage_nid_postfix(handle, image); }
+    void UnregisterImage(const void* image) { GuestAllocationsUnregisterImage_nid_postfix(handle, image); }
+#endif
     void Add(void* pointer, std::size_t bytes, bool readable, bool writable) { GuestAllocationsAdd_nid_postfix(handle, pointer, bytes, readable, writable); }
     void RequireUnpinned(const void* pointer, std::size_t bytes) const { GuestAllocationsRequireUnpinned_nid_postfix(handle, pointer, bytes); }
     void RequireAvailable(const void* pointer, std::size_t bytes) const { GuestAllocationsRequireAvailable_nid_postfix(handle, pointer, bytes); }

core/libs/prx/libc/src/GuestAllocations.cpp

@@ -34,6 +34,7 @@ struct Registry {
     std::mutex mutex;
     std::map<std::uint64_t, std::shared_ptr<const Range>> ranges;
     bool mainImageRegistered = false;
+    std::map<std::uintptr_t, std::uintptr_t> registeredImages;
 };
 
 Registry& registry() {
@@ -134,17 +135,16 @@ void GuestAllocationsSetPinWaiter_nid_postfix(bool (*callback)(std::uintptr_t, s
 }
 
 #ifdef _WIN32
-void GuestAllocationsRegisterMainImage_nid_postfix(void*) {
+namespace {
+void registerImage(const void* image) {
     auto& state = registry();
-    if (state.mainImageRegistered) return;
-    const auto image = GetModuleHandleW(nullptr);
-    require(image != nullptr, "cannot locate the main guest image");
+    if (state.registeredImages.contains(reinterpret_cast<std::uintptr_t>(image))) return;
     auto replacement = state.ranges;
     auto cursor = reinterpret_cast<std::uintptr_t>(image);
     bool registered = false;
     for (;;) {
         MEMORY_BASIC_INFORMATION memory{};
-        require(VirtualQuery(reinterpret_cast<const void*>(cursor), &memory, sizeof(memory)) == sizeof(memory), "cannot query the main guest image");
+        require(VirtualQuery(reinterpret_cast<const void*>(cursor), &memory, sizeof(memory)) == sizeof(memory), "cannot query a guest image");
         if (memory.AllocationBase != image) break;
         require(memory.Type == MEM_IMAGE && (memory.State == MEM_COMMIT || memory.State == MEM_RESERVE), "unsupported guest image mapping");
         require(reinterpret_cast<std::uintptr_t>(memory.BaseAddress) == cursor && memory.RegionSize != 0 && memory.RegionSize <= std::numeric_limits<std::uintptr_t>::max() - cursor, "invalid guest image range");
@@ -165,10 +165,36 @@ void GuestAllocationsRegisterMainImage_nid_postfix(void*) {
         }
         cursor += memory.RegionSize;
     }
-    require(registered, "main guest image has no committed pages");
+    require(registered, "guest image has no committed pages");
     state.ranges.swap(replacement);
+    state.registeredImages.emplace(reinterpret_cast<std::uintptr_t>(image), cursor);
+}
+}
+
+void GuestAllocationsRegisterMainImage_nid_postfix(void*) {
+    auto& state = registry();
+    if (state.mainImageRegistered) return;
+    const auto image = GetModuleHandleW(nullptr);
+    require(image != nullptr, "cannot locate the main guest image");
+    registerImage(image);
     state.mainImageRegistered = true;
 }
+
+void GuestAllocationsRegisterImage_nid_postfix(void*, const void* image) {
+    require(image != nullptr, "cannot register a guest image without a base");
+    registerImage(image);
+}
+
+void GuestAllocationsUnregisterImage_nid_postfix(void* mutation, const void* image) {
+    auto& state = registry();
+    const auto found = state.registeredImages.find(reinterpret_cast<std::uintptr_t>(image));
+    if (found == state.registeredImages.end()) return;
+    const auto [begin, end] = *found;
+    GuestAllocationsRequireUnpinned_nid_postfix(mutation, image, end - begin);
+    recordChange(mutation, image, end - begin);
+    std::erase_if(state.ranges, [&](const auto& entry) { return !entry.second->releasable && entry.second->allocationAddress >= begin && entry.second->allocationAddress < end; });
+    state.registeredImages.erase(found);
+}
 #else
 void GuestAllocationsRegisterMainImage_nid_postfix(void*) {
     auto& state = registry();

core/libs/prx/libkernel/CMakeLists.txt

@@ -691,6 +691,10 @@ if(BUILD_TESTING)
     configure_windows_unwind(guest_memory_tests)
     add_test(NAME guest_memory COMMAND guest_memory_tests)
     if(WIN32)
+        add_library(guest_memory_module SHARED ${LIBS_INCLUDE_DIR}/tests/GuestMemoryModule.cpp)
+        set_target_properties(guest_memory_module PROPERTIES PREFIX "" SUFFIX ".guest.prx")
+        add_dependencies(guest_memory_tests guest_memory_module)
+        target_compile_definitions(guest_memory_tests PRIVATE GUEST_MEMORY_MODULE="$<TARGET_FILE:guest_memory_module>")
         set_tests_properties(guest_memory PROPERTIES ENVIRONMENT_MODIFICATION
             "PATH=path_list_prepend:$<TARGET_FILE_DIR:libkernel>")
     endif()

core/libs/prx/libkernel/DirectMemory/DirectMemory.cpp

@@ -11,9 +11,12 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <cwctype>
 #include <memory>
 #include <limits>
 #include <stdexcept>
+#include <string>
+#include <string_view>
 #include <system_error>
 #include <vector>
 
@@ -777,6 +780,24 @@ int DoMapAnon(void** addr, size_t len, int prot, int flags) {
     return 0;
 }
 
+#ifdef _WIN32
+static bool IsGuestModuleImage(const void* base) {
+    std::wstring path(MAX_PATH, L'\0');
+    for (;;) {
+        const auto length = GetModuleFileNameW(static_cast<HMODULE>(const_cast<void*>(base)), path.data(), static_cast<DWORD>(path.size()));
+        if (length == 0) return false;
+        if (length < path.size()) {
+            path.resize(length);
+            break;
+        }
+        path.resize(path.size() * 2);
+    }
+    constexpr std::wstring_view suffix = L".guest.prx";
+    if (path.size() < suffix.size()) return false;
+    return std::equal(suffix.begin(), suffix.end(), path.end() - static_cast<std::ptrdiff_t>(suffix.size()), [](wchar_t expected, wchar_t actual) { return std::towlower(actual) == expected; });
+}
+#endif
+
 int DoMprotect(const void* addr, size_t len, int prot) {
     Trace("protect %p+0x%zx prot=0x%x", addr, len, prot);
     const auto address = reinterpret_cast<std::uintptr_t>(addr);
@@ -793,8 +814,9 @@ int DoMprotect(const void* addr, size_t len, int prot) {
     MEMORY_BASIC_INFORMATION memory{};
     if (VirtualQuery(pointer, &memory, sizeof(memory)) != sizeof(memory)) throw std::runtime_error("Cannot query guest memory protection range");
     if (memory.Type == MEM_IMAGE) {
-        if (memory.AllocationBase != GetModuleHandleW(nullptr)) throw std::invalid_argument("Memory protection of a foreign image is not supported");
-        mutation.RegisterMainImage();
+        if (memory.AllocationBase == GetModuleHandleW(nullptr)) mutation.RegisterMainImage();
+        else if (IsGuestModuleImage(memory.AllocationBase)) mutation.RegisterImage(memory.AllocationBase);
+        else throw std::invalid_argument("Memory protection of a foreign image is not supported");
     }
 #else
     mutation.RegisterMainImage();

core/libs/prx/libkernel/Module/src/DynamicLoader.cpp

@@ -1,4 +1,5 @@
 #include "prx/libc/include/General.hpp"
+#include "prx/libc/include/GuestAllocations.hpp"
 #include <nid/NidCompute.hpp>
 #include <array>
 #include <filesystem>
@@ -171,12 +172,32 @@ void* APS5_VABI dlsym_nid_postfix(void* handle, const char* name) {
     } catch (const std::exception& error) { Error(error.what()); return nullptr; }
 }
 int APS5_VABI dlclose_nid_postfix(void* handle) {
+    const auto key = reinterpret_cast<std::uintptr_t>(handle);
     std::shared_ptr<Module> module;
+#ifdef _WIN32
+    bool lastReference = true;
+#endif
+    {
+        std::lock_guard lock(modulesMutex);
+        auto found = modules.find(key);
+        if (found == modules.end()) { Error("dlclose: invalid module handle"); return -1; }
+        module = found->second;
+#ifdef _WIN32
+        for (const auto& [other, entry] : modules) {
+            if (other != key && entry->native == module->native) lastReference = false;
+        }
+#endif
+    }
+#ifdef _WIN32
+    if (module->owned && module->native && lastReference) {
+        GuestAllocations::Mutation mutation;
+        mutation.UnregisterImage(module->native);
+    }
+#endif
     {
         std::lock_guard lock(modulesMutex);
-        auto found = modules.find(reinterpret_cast<std::uintptr_t>(handle));
+        auto found = modules.find(key);
         if (found == modules.end()) { Error("dlclose: invalid module handle"); return -1; }
-        module = std::move(found->second);
         modules.erase(found);
     }
     // Unload outside the registry lock: module destructors may call loader APIs.

core/libs/tests/GuestMemory.cpp

@@ -7,6 +7,7 @@
 #include "prx/libkernel/KernelErrors.hpp"
 #include "prx/libkernel/DirectMemory/DirectMemory.hpp"
 #include <array>
+#include <atomic>
 #include "SceTypes.hpp"
 #include <chrono>
 #include <cstring>
@@ -19,6 +20,7 @@
 #include <cstdlib>
 #include <limits>
 #include <source_location>
+#include <stdexcept>
 #include <thread>
 #include <utility>
 #include <vector>
@@ -70,6 +72,9 @@ int APS5_VABI sceKernelAioDeleteRequest(std::int32_t, std::int32_t*);
 int APS5_VABI sceKernelMlock_nid_postfix(void*, std::uint64_t);
 int APS5_VABI sceKernelGetDirectMemoryType(std::int64_t, int*, std::int64_t*, std::int64_t*);
 int APS5_VABI sceKernelBatchMap2(KernelBatchMapEntry*, int, int*, int);
+void* APS5_VABI dlopen_nid_postfix(const char*, int);
+void* APS5_VABI dlsym_nid_postfix(void*, const char*);
+int APS5_VABI dlclose_nid_postfix(void*);
 }
 
 static void Require(bool condition, std::source_location location = std::source_location::current()) {
@@ -534,6 +539,93 @@ static void CheckNoOverwriteRejectsHostOccupiedMapping() {
     Require(sceKernelReleaseDirectMemory(phys, page) == 0);
 }
 
+static DWORD ImageProtection(const void* pointer) {
+    MEMORY_BASIC_INFORMATION memory{};
+    Require(VirtualQuery(pointer, &memory, sizeof(memory)) == sizeof(memory));
+    return memory.Protect & 0xffu;
+}
+
+static unsigned char* GuestModulePage(void* module, std::size_t page) {
+    auto* data = static_cast<unsigned char*>(dlsym_nid_postfix(module, "guestMemoryModuleData"));
+    Require(data != nullptr);
+    auto* target = reinterpret_cast<unsigned char*>((reinterpret_cast<std::uintptr_t>(data) + page - 1) & ~(page - 1));
+    Require(target + page <= data + 0x10000);
+    return target;
+}
+
+static bool RegisteredGuestRange(const void* pointer, std::size_t bytes) {
+    GuestAllocations::Mutation mutation;
+    return mutation.Overlaps(pointer, bytes);
+}
+
+static void CheckGuestModuleImageProtection() {
+    constexpr std::size_t page = 0x4000;
+    const auto path = std::filesystem::relative(GUEST_MEMORY_MODULE).generic_string();
+    void* module = dlopen_nid_postfix(path.c_str(), 2);
+    Require(module != nullptr);
+    auto* target = GuestModulePage(module, page);
+    Require(!RegisteredGuestRange(target, page));
+    Require(sceKernelMprotect(target, page, 1) == 0);
+    Require(RegisteredGuestRange(target, page));
+    Require(ImageProtection(target) == PAGE_READONLY);
+    Require(sceKernelMprotect(target, page, 3) == 0);
+    Require(ImageProtection(target) == PAGE_READWRITE || ImageProtection(target) == PAGE_WRITECOPY);
+    target[1] = 7;
+    Require(target[1] == 7);
+    Require(sceKernelMprotect(target, page, 1) == 0);
+    Require(ImageProtection(target) == PAGE_READONLY);
+
+    {
+        auto held = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
+        bool failed = false;
+        try {
+            dlclose_nid_postfix(module);
+        } catch (const std::runtime_error&) {
+            failed = true;
+        }
+        Require(failed);
+        Require(GuestModulePage(module, page) == target);
+        Require(RegisteredGuestRange(target, page));
+        Require(ImageProtection(target) == PAGE_READONLY);
+    }
+
+    auto lease = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
+    std::atomic<bool> closed{false};
+    std::thread closer([&] {
+        Require(dlclose_nid_postfix(module) == 0);
+        closed = true;
+    });
+    std::this_thread::sleep_for(std::chrono::milliseconds(200));
+    Require(!closed);
+    Require(ImageProtection(target) == PAGE_READONLY);
+    lease.clear();
+    closer.join();
+    Require(!RegisteredGuestRange(target, page));
+    {
+        GuestAllocations::Mutation mutation;
+        mutation.Add(target, page, true, true);
+        mutation.Remove(target);
+    }
+
+    module = dlopen_nid_postfix(path.c_str(), 2);
+    Require(module != nullptr);
+    target = GuestModulePage(module, page);
+    Require(!RegisteredGuestRange(target, page));
+    Require(sceKernelMprotect(target, page, 1) == 0);
+    Require(ImageProtection(target) == PAGE_READONLY);
+    Require(RegisteredGuestRange(target, page));
+    Require(dlclose_nid_postfix(module) == 0);
+    Require(!RegisteredGuestRange(target, page));
+
+    bool rejected = false;
+    try {
+        sceKernelMprotect(GetModuleHandleA("kernel32.dll"), page, 1);
+    } catch (const std::invalid_argument&) {
+        rejected = true;
+    }
+    Require(rejected);
+}
+
 static void CheckFixedMappingsReachTheApplicationAreaEnd() {
     constexpr std::size_t page = 0x4000;
     constexpr std::uintptr_t applicationAreaEnd = 0xFC00000000ull;
@@ -1241,6 +1333,9 @@ static void CheckVirtualQueryPartialMunmap() {
 }
 
 int main() {
+#ifdef _WIN32
+    _putenv_s("APS5_PIN_WAIT_MS", "1000");
+#endif
     CheckVirtualQuerySplitFlexibleRanges();
     CheckVirtualQueryForNonReadableGuestRange();
     CheckReleaseFlexibleMemory();
@@ -1267,6 +1362,7 @@ int main() {
 #ifdef _WIN32
     CheckNoOverwriteRejectsHostOccupiedMapping();
     CheckFixedMappingsReachTheApplicationAreaEnd();
+    CheckGuestModuleImageProtection();
 #endif
     CheckSharedWriteTracking();
     CheckReadsIntoSharedWriteTracking();

core/libs/tests/GuestMemoryModule.cpp

@@ -0,0 +1 @@
+extern "C" unsigned char guestMemoryModuleData[0x10000] = {1};