100 lines
4.2 KiB
C++
100 lines
4.2 KiB
C++
// SPDX-License-Identifier: LGPL-2.1-or-later
|
|
#pragma once
|
|
#include <windows.h>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <memory>
|
|
|
|
namespace native_control {
|
|
inline bool query(HWND window, UINT message, WPARAM wp, LPARAM lp, DWORD_PTR *result) {
|
|
return SendMessageTimeoutW(window, message, wp, lp, SMTO_ABORTIFHUNG,
|
|
100, result) != 0;
|
|
}
|
|
|
|
// Common-control messages above WM_USER do not marshal embedded pointers.
|
|
// Own the target allocation and process handle together, including failures.
|
|
class RemoteBuffer {
|
|
struct Allocation {
|
|
HANDLE process = nullptr;
|
|
void *memory = nullptr;
|
|
BOOL wow64 = FALSE;
|
|
bool pending = false;
|
|
LRESULT result = 0;
|
|
~Allocation() {
|
|
if (memory) VirtualFreeEx(process, memory, 0, MEM_RELEASE);
|
|
if (process) CloseHandle(process);
|
|
}
|
|
};
|
|
inline static thread_local std::weak_ptr<Allocation> outstanding_;
|
|
static void CALLBACK completed(HWND, UINT, ULONG_PTR data, LRESULT result) {
|
|
std::unique_ptr<std::shared_ptr<Allocation>> retained(
|
|
reinterpret_cast<std::shared_ptr<Allocation> *>(data));
|
|
(*retained)->result = result;
|
|
(*retained)->pending = false;
|
|
}
|
|
public:
|
|
RemoteBuffer(HWND window, size_t size) {
|
|
// One late application request may retain its allocation. Do not create
|
|
// another on every polling tick while that application is unresponsive.
|
|
if (auto previous = outstanding_.lock(); previous && previous->pending) return;
|
|
allocation_ = std::make_shared<Allocation>();
|
|
DWORD processId = 0;
|
|
GetWindowThreadProcessId(window, &processId);
|
|
allocation_->process = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ |
|
|
PROCESS_VM_WRITE | PROCESS_QUERY_LIMITED_INFORMATION,
|
|
FALSE, processId);
|
|
if (!allocation_->process || !IsWow64Process(allocation_->process, &allocation_->wow64)) return;
|
|
allocation_->memory = VirtualAllocEx(allocation_->process, nullptr, size,
|
|
MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
|
|
// Every byte of the allocation must be addressable by the target.
|
|
if (allocation_->memory && allocation_->wow64 &&
|
|
(address() > UINT32_MAX || size - 1 > UINT32_MAX - address())) {
|
|
VirtualFreeEx(allocation_->process, allocation_->memory, 0, MEM_RELEASE);
|
|
allocation_->memory = nullptr;
|
|
}
|
|
}
|
|
RemoteBuffer(const RemoteBuffer &) = delete;
|
|
RemoteBuffer &operator=(const RemoteBuffer &) = delete;
|
|
|
|
explicit operator bool() const { return allocation_ && allocation_->memory; }
|
|
bool is32Bit() const { return allocation_->wow64 != FALSE; }
|
|
uintptr_t address(size_t offset = 0) const {
|
|
return reinterpret_cast<uintptr_t>(allocation_->memory) + offset;
|
|
}
|
|
bool write(const void *buffer, size_t size) const {
|
|
return !allocation_->pending && WriteProcessMemory(allocation_->process, allocation_->memory, buffer, size, nullptr) != 0;
|
|
}
|
|
bool read(void *buffer, size_t size, size_t offset = 0) const {
|
|
return !allocation_->pending && ReadProcessMemory(allocation_->process, reinterpret_cast<const void *>(address(offset)),
|
|
buffer, size, nullptr) != 0;
|
|
}
|
|
|
|
bool query(HWND window, UINT message, WPARAM wp, DWORD_PTR *result) {
|
|
if (!*this || allocation_->pending) return false;
|
|
allocation_->pending = true;
|
|
outstanding_ = allocation_;
|
|
auto retained = new std::shared_ptr<Allocation>(allocation_);
|
|
if (!SendMessageCallbackW(window, message, wp, static_cast<LPARAM>(address()),
|
|
completed, reinterpret_cast<ULONG_PTR>(retained))) {
|
|
allocation_->pending = false;
|
|
delete retained;
|
|
return false;
|
|
}
|
|
const ULONGLONG deadline = GetTickCount64() + 100;
|
|
while (allocation_->pending && GetTickCount64() < deadline) {
|
|
// Dispatch sent-message completions without consuming queued UI work.
|
|
// Reentrant polls cannot allocate while outstanding_ remains pending.
|
|
MSG message{};
|
|
PeekMessageW(&message, nullptr, 0, 0, PM_NOREMOVE);
|
|
if (allocation_->pending) Sleep(1);
|
|
}
|
|
if (allocation_->pending) return false;
|
|
*result = static_cast<DWORD_PTR>(allocation_->result);
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
std::shared_ptr<Allocation> allocation_;
|
|
};
|
|
} // namespace native_control
|