Use unicode API to get module handle

BasedInc · Sep 30, 2024 · 9cf7f69 · 9cf7f69
1 parent 337ce1e
commit 9cf7f69
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/include/libhat/Process.hpp b/include/libhat/Process.hpp
@@ -2,7 +2,6 @@
 
 #include <optional>
 #include <span>
-#include <string>
 #include <string_view>
 
 namespace hat::process {
@@ -15,7 +14,7 @@ namespace hat::process {
 
     /// Returns an optional containing the module with the given name in the current process
     /// If the module is not found, std::nullopt is returned instead
-    [[nodiscard]] std::optional<module_t> get_module(const std::string& name);
+    [[nodiscard]] std::optional<module_t> get_module(std::string_view name);
 
     /// Returns the module located at the specified base address. Optionally, a size may be provided to indicate
     /// the size of the allocation pointed to by address, preventing a potential out-of-bounds read.

diff --git a/src/os/win32/Process.cpp b/src/os/win32/Process.cpp
@@ -42,17 +42,27 @@ namespace hat::process {
     }
 
     module_t get_process_module() {
-        return module_t{reinterpret_cast<uintptr_t>(GetModuleHandleA(nullptr))};
+        return module_t{reinterpret_cast<uintptr_t>(GetModuleHandleW(nullptr))};
     }
 
-    std::optional<module_t> get_module(const std::string& name) {
-        if (const auto module = GetModuleHandleA(name.c_str()); module) {
+    std::optional<module_t> get_module(const std::string_view name) {
+        const int size = MultiByteToWideChar(CP_UTF8, 0, name.data(), static_cast<int>(name.size()), nullptr, 0);
+        if (!size) {
+            return {};
+        }
+
+        std::wstring str;
+        str.resize(size);
+
+        MultiByteToWideChar(CP_UTF8, 0, name.data(), static_cast<int>(name.size()), str.data(), size);
+
+        if (const auto module = GetModuleHandleW(str.c_str()); module) {
             return module_t{std::bit_cast<uintptr_t>(module)};
         }
         return {};
     }
 
-    std::optional<module_t> module_at(void* address, std::optional<size_t> size) {
+    std::optional<module_t> module_at(void* address, const std::optional<size_t> size) {
         if (isValidModule(address, size)) {
             return module_t{std::bit_cast<uintptr_t>(address)};
         }
@@ -72,7 +82,7 @@ namespace hat::process {
         const auto maxChars = std::min<size_t>(name.size(), 8);
 
         const auto* sectionHeader = IMAGE_FIRST_SECTION(&ntHeaders);
-        for (int i = 0; i < ntHeaders.FileHeader.NumberOfSections; i++, sectionHeader++) {
+        for (DWORD i = 0; i < ntHeaders.OptionalHeader.NumberOfRvaAndSizes; i++, sectionHeader++) {
             if (strncmp(name.data(), reinterpret_cast<const char*>(sectionHeader->Name), maxChars) == 0) {
                 return {
                     bytes + sectionHeader->VirtualAddress,