Process: import pdb::Process from winapi-debug

3 files changed, 129 insertions, 32 deletions

diff --git a/include/winapi/process.hpp b/include/winapi/process.hpp
index 2e1420f..b79d7c2 100644
--- a/include/winapi/process.hpp
+++ b/include/winapi/process.hpp
@@ -69,6 +69,8 @@ inline void swap(ShellParameters& a, ShellParameters& b) BOOST_NOEXCEPT_OR_NOTHR
 
 class Process {
 public:
+    using ID = DWORD;
+
     static Process create(ProcessParameters);
     static Process create(const CommandLine&);
     static Process create(const CommandLine&, process::IO);
@@ -76,28 +78,40 @@ public:
     static Process shell(const ShellParameters&);
     static Process shell(const CommandLine&);
 
+    static Process current();
+    static Process open(ID, DWORD permissions = default_permissions());
+    static Process open_r(ID);
+
+    static DWORD default_permissions();
+    static DWORD read_permissions();
+
     // VS 2013 won't generate these automatically.
     Process(Process&&) BOOST_NOEXCEPT_OR_NOTHROW;
     Process& operator=(Process) BOOST_NOEXCEPT_OR_NOTHROW;
     void swap(Process& other) BOOST_NOEXCEPT_OR_NOTHROW;
     Process(const Process&) = delete;
 
+    ID get_id() const { return m_id; }
+    const Handle& get_handle() const { return m_handle; }
+
     bool is_running() const;
     void wait() const;
     void terminate(int ec = 0) const;
     void shut_down(int ec = 0) const;
     int get_exit_code() const;
 
-    static std::string get_exe_path();
+    std::string get_exe_path() const;
 
     static Resource get_resource(uint32_t id);
     static std::string get_resource_string(uint32_t id);
 
 private:
-    explicit Process(Handle&& handle) : m_handle(std::move(handle)) {}
+    explicit Process(Handle&& handle);
+    Process(ID, Handle&& handle);
 
     static HMODULE get_exe_module();
 
+    ID m_id;
     Handle m_handle;
 };
 
diff --git a/src/process.cpp b/src/process.cpp
index 0d9837b..5e9ac13 100644
--- a/src/process.cpp
+++ b/src/process.cpp
@@ -5,6 +5,7 @@
 
 #include <winapi/cmd_line.hpp>
 #include <winapi/error.hpp>
+#include <winapi/handle.hpp>
 #include <winapi/process.hpp>
 #include <winapi/process_io.hpp>
 #include <winapi/resource.hpp>
@@ -158,6 +159,85 @@ Handle shell_execute(const ShellParameters& params) {
     return Handle{info.hProcess};
 }
 
+Handle open_process(DWORD id, DWORD permissions) {
+    Handle process{OpenProcess(permissions, FALSE, id)};
+    if (!process.is_valid()) {
+        throw error::windows(GetLastError(), "OpenProcess");
+    }
+    return process;
+}
+
+class PathBuffer {
+public:
+    PathBuffer() : m_size{min_size} { m_data.resize(m_size); }
+
+    DWORD get_size() const { return m_size; }
+
+    wchar_t* get_data() { return m_data.data(); }
+
+    void grow() {
+        if (m_size < min_size) {
+            m_size = min_size;
+        } else {
+            // Check if we can still multiply by two.
+            if (std::numeric_limits<decltype(m_size)>::max() - m_size < m_size)
+                throw std::range_error{"Path buffer is too large"};
+            m_size *= 2;
+        }
+        m_data.resize(m_size);
+    }
+
+private:
+    static constexpr DWORD min_size = MAX_PATH;
+
+    DWORD m_size;
+    std::vector<wchar_t> m_data;
+};
+
+std::string get_current_exe_path(PathBuffer& buffer) {
+    SetLastError(ERROR_SUCCESS);
+
+    const auto ec = ::GetModuleFileNameW(NULL, buffer.get_data(), buffer.get_size());
+
+    if (ec == 0) {
+        throw error::windows(GetLastError(), "GetModuleFileNameW");
+    }
+
+    if (ec == buffer.get_size() && GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
+        buffer.grow();
+        return get_current_exe_path(buffer);
+    }
+
+    return narrow(buffer.get_data());
+}
+
+std::string get_current_exe_path() {
+    PathBuffer buffer;
+    return get_current_exe_path(buffer);
+}
+
+std::string get_exe_path(const Handle& process, PathBuffer& buffer) {
+    auto size = buffer.get_size();
+
+    const auto ec = ::QueryFullProcessImageNameW(process.get(), 0, buffer.get_data(), &size);
+
+    if (ec != 0) {
+        return narrow(buffer.get_data());
+    }
+
+    if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
+        buffer.grow();
+        return get_exe_path(process, buffer);
+    }
+
+    throw error::windows(GetLastError(), "QueryFullProcessImageNameW");
+}
+
+std::string get_exe_path(const Handle& process) {
+    PathBuffer buffer;
+    return get_exe_path(process, buffer);
+}
+
 } // namespace
 
 ProcessParameters::ProcessParameters(ProcessParameters&& other) BOOST_NOEXCEPT_OR_NOTHROW
@@ -216,6 +296,26 @@ Process Process::shell(const CommandLine& cmd_line) {
     return shell(params);
 }
 
+Process Process::current() {
+    return Process{::GetCurrentProcessId(), Handle{::GetCurrentProcess()}};
+}
+
+Process Process::open(DWORD id, DWORD permissions) {
+    return Process{id, open_process(id, permissions)};
+}
+
+Process Process::open_r(DWORD id) {
+    return open(id, read_permissions());
+}
+
+DWORD Process::default_permissions() {
+    return PROCESS_QUERY_INFORMATION;
+}
+
+DWORD Process::read_permissions() {
+    return default_permissions() | PROCESS_VM_READ;
+}
+
 Process::Process(Process&& other) BOOST_NOEXCEPT_OR_NOTHROW {
     swap(other);
 }
@@ -287,6 +387,14 @@ int Process::get_exit_code() const {
     return static_cast<int>(ec);
 }
 
+std::string Process::get_exe_path() const {
+    if (m_handle.get() == ::GetCurrentProcess()) {
+        return get_current_exe_path();
+    } else {
+        return winapi::get_exe_path(m_handle);
+    }
+}
+
 HMODULE Process::get_exe_module() {
     const auto module = ::GetModuleHandleW(NULL);
     if (module == NULL) {
@@ -295,35 +403,6 @@ HMODULE Process::get_exe_module() {
     return module;
 }
 
-std::string Process::get_exe_path() {
-    BOOST_STATIC_CONSTEXPR std::size_t init_buffer_size = MAX_PATH;
-    static_assert(init_buffer_size > 0, "init_buffer_size must be positive");
-
-    std::vector<wchar_t> buffer;
-    buffer.resize(init_buffer_size);
-
-    while (true) {
-        SetLastError(ERROR_SUCCESS);
-
-        if (buffer.size() > std::numeric_limits<DWORD>::max())
-            throw std::range_error{"Path buffer is too large"};
-        const auto nch =
-            ::GetModuleFileNameW(NULL, buffer.data(), static_cast<DWORD>(buffer.size()));
-
-        if (nch == 0) {
-            throw error::windows(GetLastError(), "GetModuleFileNameW");
-        }
-
-        if (nch == buffer.size() && GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
-            buffer.resize(2 * buffer.size());
-            continue;
-        }
-
-        buffer.resize(nch);
-        return narrow(buffer);
-    }
-}
-
 std::string Process::get_resource_string(uint32_t id) {
     wchar_t* s = nullptr;
 
@@ -364,4 +443,8 @@ Resource Process::get_resource(uint32_t id) {
     return {data, nb};
 }
 
+Process::Process(Handle&& handle) : Process{::GetProcessId(handle.get()), std::move(handle)} {}
+
+Process::Process(ID id, Handle&& handle) : m_id{id}, m_handle{std::move(handle)} {}
+
 } // namespace winapi
diff --git a/test/unit_tests/process.cpp b/test/unit_tests/process.cpp
index a048c84..56a8337 100644
--- a/test/unit_tests/process.cpp
+++ b/test/unit_tests/process.cpp
@@ -25,7 +25,7 @@ using namespace winapi::process;
 BOOST_AUTO_TEST_SUITE(process_tests)
 
 BOOST_AUTO_TEST_CASE(get_exe_path) {
-    const auto path = Process::get_exe_path();
+    const auto path = Process::current().get_exe_path();
     BOOST_TEST_MESSAGE("Executable path: " << path);
 }