init
This commit is contained in:
@@ -0,0 +1,657 @@
|
||||
#include "client_pipe.h"
|
||||
|
||||
|
||||
static LONG g_started = 0;
|
||||
static HANDLE g_stop_event = nullptr;
|
||||
static HANDLE g_queue_event = nullptr;
|
||||
static HANDLE g_connected_event = nullptr;
|
||||
static HANDLE g_worker_thread = nullptr;
|
||||
static HANDLE g_pipe = INVALID_HANDLE_VALUE;
|
||||
static std::atomic<DWORD> g_client_id{ 0 };
|
||||
static std::atomic<uint64_t> g_session_id{ 0 };
|
||||
static std::atomic<uint64_t> g_next_request_id{ 1 };
|
||||
static std::mutex g_queue_mutex;
|
||||
static std::deque<std::shared_ptr<PendingCommand>> g_queue;
|
||||
|
||||
|
||||
static void LogA(const char* fmt, ...)
|
||||
{
|
||||
char buf[512]{};
|
||||
va_list ap;
|
||||
va_start(ap, fmt);
|
||||
vsnprintf_s(buf, sizeof(buf), _TRUNCATE, fmt, ap);
|
||||
va_end(ap);
|
||||
OutputDebugStringA(buf);
|
||||
}
|
||||
|
||||
static void ExitInjectedProcess(UINT exit_code)
|
||||
{
|
||||
LogA("[dll-client] ExitProcess code=%u\n", exit_code);
|
||||
ExitProcess(exit_code);
|
||||
}
|
||||
|
||||
static bool ReadExact(HANDLE pipe, void* buffer, DWORD size)
|
||||
{
|
||||
BYTE* ptr = static_cast<BYTE*>(buffer);
|
||||
DWORD total = 0;
|
||||
while (total < size)
|
||||
{
|
||||
DWORD got = 0;
|
||||
if (!ReadFile(pipe, ptr + total, size - total, &got, nullptr))
|
||||
return false;
|
||||
if (got == 0)
|
||||
return false;
|
||||
total += got;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool WriteExact(HANDLE pipe, const void* buffer, DWORD size)
|
||||
{
|
||||
const BYTE* ptr = static_cast<const BYTE*>(buffer);
|
||||
DWORD total = 0;
|
||||
while (total < size)
|
||||
{
|
||||
DWORD wrote = 0;
|
||||
if (!WriteFile(pipe, ptr + total, size - total, &wrote, nullptr))
|
||||
return false;
|
||||
if (wrote == 0)
|
||||
return false;
|
||||
total += wrote;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
bool WideToUtf8(const std::wstring& input, char* output, size_t output_size)
|
||||
{
|
||||
if (!output || output_size == 0)
|
||||
return false;
|
||||
|
||||
output[0] = '\0';
|
||||
if (input.empty())
|
||||
return true;
|
||||
|
||||
int written = WideCharToMultiByte(
|
||||
CP_UTF8,
|
||||
0,
|
||||
input.c_str(),
|
||||
static_cast<int>(input.size()),
|
||||
output,
|
||||
static_cast<int>(output_size - 1),
|
||||
nullptr,
|
||||
nullptr);
|
||||
|
||||
if (written <= 0)
|
||||
return false;
|
||||
|
||||
output[written] = '\0';
|
||||
return true;
|
||||
}
|
||||
|
||||
static void QueuePush(const std::shared_ptr<PendingCommand>& cmd)
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(g_queue_mutex);
|
||||
g_queue.push_back(cmd);
|
||||
SetEvent(g_queue_event);
|
||||
}
|
||||
|
||||
static std::shared_ptr<PendingCommand> QueuePop()
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(g_queue_mutex);
|
||||
if (g_queue.empty())
|
||||
{
|
||||
ResetEvent(g_queue_event);
|
||||
return {};
|
||||
}
|
||||
|
||||
auto cmd = g_queue.front();
|
||||
g_queue.pop_front();
|
||||
if (g_queue.empty())
|
||||
ResetEvent(g_queue_event);
|
||||
return cmd;
|
||||
}
|
||||
|
||||
static void FailPendingCommands(DWORD status, const char* reply)
|
||||
{
|
||||
(void)reply;
|
||||
|
||||
for (;;)
|
||||
{
|
||||
auto cmd = QueuePop();
|
||||
if (!cmd)
|
||||
break;
|
||||
|
||||
cmd->status = status;
|
||||
if (cmd->done_event)
|
||||
SetEvent(cmd->done_event);
|
||||
}
|
||||
}
|
||||
|
||||
static void RequestStopNoWait()
|
||||
{
|
||||
if (g_stop_event)
|
||||
SetEvent(g_stop_event);
|
||||
}
|
||||
|
||||
static void CleanupClientStateNoWait()
|
||||
{
|
||||
RequestStopNoWait();
|
||||
|
||||
if (g_worker_thread)
|
||||
{
|
||||
CloseHandle(g_worker_thread);
|
||||
g_worker_thread = nullptr;
|
||||
}
|
||||
|
||||
if (g_stop_event)
|
||||
{
|
||||
CloseHandle(g_stop_event);
|
||||
g_stop_event = nullptr;
|
||||
}
|
||||
|
||||
if (g_queue_event)
|
||||
{
|
||||
CloseHandle(g_queue_event);
|
||||
g_queue_event = nullptr;
|
||||
}
|
||||
|
||||
if (g_connected_event)
|
||||
{
|
||||
CloseHandle(g_connected_event);
|
||||
g_connected_event = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
static bool ConnectToServer(HANDLE& pipe)
|
||||
{
|
||||
for (;;)
|
||||
{
|
||||
if (WaitForSingleObject(g_stop_event, 0) == WAIT_OBJECT_0)
|
||||
return false;
|
||||
|
||||
pipe = CreateFileW(
|
||||
kPipeName,
|
||||
GENERIC_READ | GENERIC_WRITE,
|
||||
0,
|
||||
nullptr,
|
||||
OPEN_EXISTING,
|
||||
0,
|
||||
nullptr);
|
||||
|
||||
if (pipe != INVALID_HANDLE_VALUE)
|
||||
{
|
||||
DWORD mode = PIPE_READMODE_MESSAGE;
|
||||
SetNamedPipeHandleState(pipe, &mode, nullptr, nullptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
DWORD err = GetLastError();
|
||||
if (err == ERROR_PIPE_BUSY)
|
||||
{
|
||||
if (!WaitNamedPipeW(kPipeName, 1000))
|
||||
Sleep(250);
|
||||
continue;
|
||||
}
|
||||
|
||||
Sleep(500);
|
||||
}
|
||||
}
|
||||
|
||||
static bool SendHello(HANDLE pipe)
|
||||
{
|
||||
Msg hello{};
|
||||
hello.type = static_cast<uint32_t>(MsgType::Hello);
|
||||
hello.pid = GetCurrentProcessId();
|
||||
hello.version = kProtocolVersion;
|
||||
hello.tick = static_cast<uint64_t>(NowMs());
|
||||
hello.session_id = 0;
|
||||
|
||||
if (!WriteExact(pipe, &hello, sizeof(hello)))
|
||||
{
|
||||
LogA("[dll-client] hello write failed: %lu\n", GetLastError());
|
||||
return false;
|
||||
}
|
||||
|
||||
Msg ack{};
|
||||
if (!ReadExact(pipe, &ack, sizeof(ack)) || ack.type != static_cast<uint32_t>(MsgType::HelloAck))
|
||||
{
|
||||
LogA("[dll-client] hello ack failed: %lu\n", GetLastError());
|
||||
return false;
|
||||
}
|
||||
|
||||
g_client_id.store(ack.client_id);
|
||||
g_session_id.store(ack.session_id);
|
||||
SetEvent(g_connected_event);
|
||||
|
||||
LogA("[dll-client] connected client_id=%lu session=%llu\n",
|
||||
g_client_id.load(),
|
||||
static_cast<unsigned long long>(g_session_id.load()));
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool SendHeartbeat(HANDLE pipe)
|
||||
{
|
||||
Msg heartbeat{};
|
||||
heartbeat.type = static_cast<uint32_t>(MsgType::Heartbeat);
|
||||
heartbeat.client_id = g_client_id.load();
|
||||
heartbeat.pid = GetCurrentProcessId();
|
||||
heartbeat.version = kProtocolVersion;
|
||||
heartbeat.session_id = g_session_id.load();
|
||||
heartbeat.tick = static_cast<uint64_t>(NowMs());
|
||||
|
||||
if (!WriteExact(pipe, &heartbeat, sizeof(heartbeat)))
|
||||
{
|
||||
LogA("[dll-client] heartbeat write failed: %lu\n", GetLastError());
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool PumpServerMessages(HANDLE pipe)
|
||||
{
|
||||
for (;;)
|
||||
{
|
||||
DWORD available = 0;
|
||||
if (!PeekNamedPipe(pipe, nullptr, 0, nullptr, &available, nullptr))
|
||||
{
|
||||
LogA("[dll-client] server disconnected: %lu\n", GetLastError());
|
||||
ExitInjectedProcess(0);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (available < sizeof(Msg))
|
||||
return true;
|
||||
|
||||
Msg msg{};
|
||||
if (!ReadExact(pipe, &msg, sizeof(msg)))
|
||||
return false;
|
||||
|
||||
if (msg.type == static_cast<uint32_t>(MsgType::Exit))
|
||||
{
|
||||
LogA("[dll-client] exit requested by server\n");
|
||||
SetEvent(g_stop_event);
|
||||
ExitInjectedProcess(0);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static bool SendCommand(HANDLE pipe, PendingCommand& cmd)
|
||||
{
|
||||
Msg req{};
|
||||
req.type = static_cast<uint32_t>(MsgType::CommandRequest);
|
||||
req.client_id = g_client_id.load();
|
||||
req.pid = GetCurrentProcessId();
|
||||
req.version = kProtocolVersion;
|
||||
req.command = cmd.command;
|
||||
req.status = 0;
|
||||
req.input_size = static_cast<uint32_t>(cmd.input.size());
|
||||
req.output_size = static_cast<uint32_t>(cmd.output.size());
|
||||
req.tick = static_cast<uint64_t>(NowMs());
|
||||
req.session_id = g_session_id.load();
|
||||
req.request_id = cmd.request_id;
|
||||
|
||||
if (cmd.input.size() > kPayloadBytes || cmd.output.size() > kPayloadBytes)
|
||||
return false;
|
||||
|
||||
if (!cmd.input.empty())
|
||||
memcpy(req.payload, cmd.input.data(), cmd.input.size());
|
||||
|
||||
if (!WriteExact(pipe, &req, sizeof(req)))
|
||||
return false;
|
||||
|
||||
Msg resp{};
|
||||
if (!ReadExact(pipe, &resp, sizeof(resp)))
|
||||
return false;
|
||||
|
||||
if (resp.type != static_cast<uint32_t>(MsgType::CommandReply))
|
||||
return false;
|
||||
|
||||
if (resp.request_id != cmd.request_id || resp.command != cmd.command)
|
||||
return false;
|
||||
|
||||
cmd.status = resp.status;
|
||||
if (resp.output_size && resp.output_size <= cmd.output.size())
|
||||
memcpy(cmd.output.data(), resp.payload, resp.output_size);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static DWORD WINAPI WorkerThreadProc(LPVOID)
|
||||
{
|
||||
for (;;)
|
||||
{
|
||||
if (WaitForSingleObject(g_stop_event, 0) == WAIT_OBJECT_0)
|
||||
break;
|
||||
|
||||
HANDLE pipe = INVALID_HANDLE_VALUE;
|
||||
if (!ConnectToServer(pipe))
|
||||
break;
|
||||
|
||||
g_pipe = pipe;
|
||||
ResetEvent(g_connected_event);
|
||||
|
||||
if (!SendHello(pipe))
|
||||
{
|
||||
CloseHandle(pipe);
|
||||
g_pipe = INVALID_HANDLE_VALUE;
|
||||
Sleep(500);
|
||||
continue;
|
||||
}
|
||||
|
||||
DWORD last_heartbeat = static_cast<DWORD>(NowMs());
|
||||
|
||||
for (;;)
|
||||
{
|
||||
if (WaitForSingleObject(g_stop_event, 0) == WAIT_OBJECT_0)
|
||||
break;
|
||||
|
||||
if (WaitForSingleObject(g_queue_event, 0) == WAIT_OBJECT_0)
|
||||
{
|
||||
for (;;)
|
||||
{
|
||||
auto cmd = QueuePop();
|
||||
if (!cmd)
|
||||
break;
|
||||
|
||||
if (!SendCommand(pipe, *cmd))
|
||||
{
|
||||
cmd->status = GetLastError();
|
||||
}
|
||||
|
||||
if (cmd->done_event)
|
||||
SetEvent(cmd->done_event);
|
||||
|
||||
if (WaitForSingleObject(g_stop_event, 0) == WAIT_OBJECT_0)
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const DWORD now = static_cast<DWORD>(NowMs());
|
||||
if (now - last_heartbeat >= kHeartbeatIntervalMs)
|
||||
{
|
||||
if (!SendHeartbeat(pipe))
|
||||
{
|
||||
LogA("[dll-client] server heartbeat failed, exit process\n");
|
||||
ExitInjectedProcess(0);
|
||||
break;
|
||||
}
|
||||
|
||||
last_heartbeat = now;
|
||||
}
|
||||
|
||||
if (!PumpServerMessages(pipe))
|
||||
break;
|
||||
|
||||
Sleep(10);
|
||||
}
|
||||
|
||||
ResetEvent(g_connected_event);
|
||||
|
||||
if (pipe != INVALID_HANDLE_VALUE)
|
||||
CloseHandle(pipe);
|
||||
|
||||
g_pipe = INVALID_HANDLE_VALUE;
|
||||
|
||||
if (WaitForSingleObject(g_stop_event, 0) == WAIT_OBJECT_0)
|
||||
break;
|
||||
|
||||
Sleep(250);
|
||||
}
|
||||
|
||||
FailPendingCommands(ERROR_CANCELLED, "client stopped");
|
||||
return 0;
|
||||
}
|
||||
|
||||
BOOL DllClientStart()
|
||||
{
|
||||
if (InterlockedCompareExchange(&g_started, 1, 0) != 0)
|
||||
return TRUE;
|
||||
|
||||
g_stop_event = CreateEventW(nullptr, TRUE, FALSE, nullptr);
|
||||
g_queue_event = CreateEventW(nullptr, TRUE, FALSE, nullptr);
|
||||
g_connected_event = CreateEventW(nullptr, TRUE, FALSE, nullptr);
|
||||
if (!g_stop_event || !g_queue_event || !g_connected_event)
|
||||
{
|
||||
CleanupClientStateNoWait();
|
||||
return FALSE;
|
||||
}
|
||||
|
||||
g_worker_thread = CreateThread(nullptr, 0, WorkerThreadProc, nullptr, 0, nullptr);
|
||||
if (!g_worker_thread)
|
||||
{
|
||||
CleanupClientStateNoWait();
|
||||
return FALSE;
|
||||
}
|
||||
|
||||
return TRUE;
|
||||
}
|
||||
|
||||
BOOL HBClientEnsureStarted()
|
||||
{
|
||||
return DllClientStart();
|
||||
}
|
||||
|
||||
BOOL HBClientCall(
|
||||
uint32_t command,
|
||||
const void* input,
|
||||
uint32_t input_size,
|
||||
void* output,
|
||||
uint32_t output_size,
|
||||
uint32_t* bytes_returned)
|
||||
{
|
||||
if (bytes_returned)
|
||||
*bytes_returned = 0;
|
||||
|
||||
if (input_size > kPayloadBytes || output_size > kPayloadBytes)
|
||||
return FALSE;
|
||||
|
||||
if (input_size && !input)
|
||||
return FALSE;
|
||||
|
||||
if (output_size && !output)
|
||||
return FALSE;
|
||||
|
||||
if (InterlockedCompareExchange(&g_started, 0, 0) == 0)
|
||||
{
|
||||
if (!DllClientStart())
|
||||
return FALSE;
|
||||
}
|
||||
|
||||
auto cmd = std::make_shared<PendingCommand>();
|
||||
cmd->command = command;
|
||||
cmd->request_id = g_next_request_id.fetch_add(1);
|
||||
cmd->done_event = CreateEventW(nullptr, TRUE, FALSE, nullptr);
|
||||
if (!cmd->done_event)
|
||||
return FALSE;
|
||||
|
||||
if (input_size)
|
||||
{
|
||||
const auto* ptr = static_cast<const uint8_t*>(input);
|
||||
cmd->input.assign(ptr, ptr + input_size);
|
||||
}
|
||||
|
||||
if (output_size)
|
||||
cmd->output.resize(output_size);
|
||||
|
||||
QueuePush(cmd);
|
||||
|
||||
HANDLE waits[2] = { cmd->done_event, g_stop_event };
|
||||
DWORD wait = WaitForMultipleObjects(2, waits, FALSE, 10000);
|
||||
BOOL ok = (wait == WAIT_OBJECT_0 && cmd->status == ERROR_SUCCESS);
|
||||
|
||||
if (ok && output && !cmd->output.empty())
|
||||
{
|
||||
memcpy(output, cmd->output.data(), cmd->output.size());
|
||||
if (bytes_returned)
|
||||
*bytes_returned = static_cast<uint32_t>(cmd->output.size());
|
||||
}
|
||||
|
||||
CloseHandle(cmd->done_event);
|
||||
return ok;
|
||||
}
|
||||
|
||||
BOOL HBClientSendPrintfW(const wchar_t* text)
|
||||
{
|
||||
return DllClientPrintfW(text);
|
||||
}
|
||||
|
||||
static bool PipeReadMsgTimeout(HANDLE pipe, Msg& msg, DWORD timeout_ms)
|
||||
{
|
||||
const DWORD begin = GetTickCount();
|
||||
|
||||
for (;;)
|
||||
{
|
||||
DWORD available = 0;
|
||||
if (!PeekNamedPipe(pipe, nullptr, 0, nullptr, &available, nullptr))
|
||||
return false;
|
||||
|
||||
if (available >= sizeof(Msg))
|
||||
{
|
||||
DWORD got = 0;
|
||||
return ReadFile(pipe, &msg, sizeof(msg), &got, nullptr) && got == sizeof(msg);
|
||||
}
|
||||
|
||||
if (GetTickCount() - begin >= timeout_ms)
|
||||
return false;
|
||||
|
||||
Sleep(10);
|
||||
}
|
||||
}
|
||||
|
||||
static bool PipeWriteMsg(HANDLE pipe, const Msg& msg)
|
||||
{
|
||||
DWORD wrote = 0;
|
||||
return WriteFile(pipe, &msg, sizeof(msg), &wrote, nullptr) && wrote == sizeof(msg);
|
||||
}
|
||||
|
||||
BOOL QueryServerAlreadyOpen()
|
||||
{
|
||||
HANDLE pipe = CreateFileW(
|
||||
kPipeName,
|
||||
GENERIC_READ | GENERIC_WRITE,
|
||||
0,
|
||||
nullptr,
|
||||
OPEN_EXISTING,
|
||||
0,
|
||||
nullptr);
|
||||
|
||||
if (pipe == INVALID_HANDLE_VALUE)
|
||||
return FALSE;
|
||||
|
||||
DWORD mode = PIPE_READMODE_MESSAGE;
|
||||
SetNamedPipeHandleState(pipe, &mode, nullptr, nullptr);
|
||||
|
||||
Msg hello{};
|
||||
hello.type = static_cast<uint32_t>(MsgType::Hello);
|
||||
hello.pid = GetCurrentProcessId();
|
||||
hello.version = kProtocolVersion;
|
||||
hello.tick = static_cast<uint64_t>(NowMs());
|
||||
|
||||
Msg ack{};
|
||||
if (!PipeWriteMsg(pipe, hello) ||
|
||||
!PipeReadMsgTimeout(pipe, ack, 500) ||
|
||||
ack.type != static_cast<uint32_t>(MsgType::HelloAck))
|
||||
{
|
||||
CloseHandle(pipe);
|
||||
return FALSE;
|
||||
}
|
||||
|
||||
Msg req{};
|
||||
req.type = static_cast<uint32_t>(MsgType::CommandRequest);
|
||||
req.client_id = ack.client_id;
|
||||
req.pid = GetCurrentProcessId();
|
||||
req.version = kProtocolVersion;
|
||||
req.command = static_cast<uint32_t>(CommandId::QueryStatus);
|
||||
req.session_id = ack.session_id;
|
||||
req.request_id = 1;
|
||||
req.tick = static_cast<uint64_t>(NowMs());
|
||||
|
||||
Msg resp{};
|
||||
BOOL already_open =
|
||||
PipeWriteMsg(pipe, req) &&
|
||||
PipeReadMsgTimeout(pipe, resp, 500) &&
|
||||
resp.type == static_cast<uint32_t>(MsgType::CommandReply) &&
|
||||
resp.command == static_cast<uint32_t>(CommandId::QueryStatus);
|
||||
|
||||
Msg bye{};
|
||||
bye.type = static_cast<uint32_t>(MsgType::Goodbye);
|
||||
bye.client_id = ack.client_id;
|
||||
bye.pid = GetCurrentProcessId();
|
||||
bye.version = kProtocolVersion;
|
||||
bye.session_id = ack.session_id;
|
||||
PipeWriteMsg(pipe, bye);
|
||||
|
||||
CloseHandle(pipe);
|
||||
return already_open;
|
||||
}
|
||||
|
||||
|
||||
BOOL HBClientLoadDriverW()
|
||||
{
|
||||
LoadDriverRequest req{};
|
||||
//wcsncpy_s(req.service_name, service_name, _TRUNCATE);
|
||||
//wcsncpy_s(req.driver_path, driver_path, _TRUNCATE);
|
||||
|
||||
return HBClientCall(
|
||||
static_cast<uint32_t>(CommandId::LoadDriver),
|
||||
&req,
|
||||
sizeof(req),
|
||||
nullptr,
|
||||
0,
|
||||
nullptr);
|
||||
}
|
||||
|
||||
BOOL HBClientUnloadDriverW(const wchar_t* service_name)
|
||||
{
|
||||
if (!service_name)
|
||||
return FALSE;
|
||||
|
||||
UnloadDriverRequest req{};
|
||||
wcsncpy_s(req.service_name, service_name, _TRUNCATE);
|
||||
|
||||
return HBClientCall(
|
||||
static_cast<uint32_t>(CommandId::UnloadDriver),
|
||||
&req,
|
||||
sizeof(req),
|
||||
nullptr,
|
||||
0,
|
||||
nullptr);
|
||||
}
|
||||
|
||||
BOOL HBClientNotifyBreakpointRemoved(uint32_t thread_id, uint64_t dr0, uint64_t dr7)
|
||||
{
|
||||
BreakpointRemovedRequest req{};
|
||||
req.thread_id = thread_id;
|
||||
req.dr0 = dr0;
|
||||
req.dr7 = dr7;
|
||||
|
||||
return HBClientCall(
|
||||
static_cast<uint32_t>(CommandId::BreakpointRemoved),
|
||||
&req,
|
||||
sizeof(req),
|
||||
nullptr,
|
||||
0,
|
||||
nullptr);
|
||||
}
|
||||
|
||||
BOOL DllClientPrintfW(const wchar_t* text)
|
||||
{
|
||||
if (!text)
|
||||
return FALSE;
|
||||
|
||||
PrintfRequest req{};
|
||||
if (!WideToUtf8(text, req.text, sizeof(req.text)))
|
||||
return FALSE;
|
||||
|
||||
return HBClientCall(
|
||||
static_cast<uint32_t>(CommandId::Printf),
|
||||
&req,
|
||||
sizeof(req),
|
||||
nullptr,
|
||||
0,
|
||||
nullptr);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user