OrcLib: WinApi: add GetModuleFileNameExApi

DFIR-ORC · Jan 26, 2021 · 988fe17 · 988fe17
1 parent 8bfa2ff
commit 988fe17
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/src/OrcLib/WinApiHelper.cpp b/src/OrcLib/WinApiHelper.cpp
@@ -10,6 +10,7 @@
 
 #include <windows.h>
 #include <safeint.h>
+#include <psapi.h>
 
 #include "WinApiHelper.h"
 
@@ -207,4 +208,38 @@ std::wstring GetModuleFileNameApi(HMODULE hModule, std::error_code& ec) noexcept
     return GetModuleFileNameApi(hModule, 32767 * sizeof(wchar_t), ec);
 }
 
+std::wstring GetModuleFileNameExApi(HANDLE hProcess, HMODULE hModule, size_t cbMaxOutput, std::error_code& ec) noexcept
+{
+    std::wstring path;
+
+    try
+    {
+        path.resize(cbMaxOutput);
+
+        DWORD cch = ::GetModuleFileNameExW(hProcess, hModule, path.data(), path.size());
+
+        DWORD lastError = GetLastError();
+        if (cch == 0 || lastError != ERROR_SUCCESS)
+        {
+            ec.assign(lastError, std::system_category());
+            return {};
+        }
+
+        path.resize(cch);
+        path.shrink_to_fit();
+    }
+    catch (const std::length_error& e)
+    {
+        ec = std::make_error_code(std::errc::not_enough_memory);
+        return {};
+    }
+
+    return path;
+}
+
+std::wstring GetModuleFileNameExApi(HANDLE hProcess, HMODULE hModule, std::error_code& ec) noexcept
+{
+    return GetModuleFileNameExApi(hProcess, hModule, 32767 * sizeof(wchar_t), ec);
+}
+
 }  // namespace Orc
diff --git a/src/OrcLib/WinApiHelper.h b/src/OrcLib/WinApiHelper.h
@@ -50,6 +50,12 @@ std::wstring GetModuleFileNameApi(HMODULE hModule, size_t cbMaxOutput, std::erro
 // GetModuleFileName wrapper with default maximum buffer size of 32767 characters
 std::wstring GetModuleFileNameApi(HMODULE hModule, std::error_code& ec) noexcept;
 
+// GetModuleFileNameEx wrapper with custom maximum buffer size
+std::wstring GetModuleFileNameExApi(HANDLE hProcess, HMODULE hModule, size_t cbMaxOutput, std::error_code& ec) noexcept;
+
+// GetModuleFileNameEx wrapper with default maximum buffer size of 32767 characters
+std::wstring GetModuleFileNameExApi(HANDLE hProcess, HMODULE hModule, std::error_code& ec) noexcept;
+
 }  // namespace Orc
 
 #pragma managed(pop)