555 lines
15 KiB
C++
555 lines
15 KiB
C++
#include "server_pipe.h"
|
|
#include "Driver.h"
|
|
|
|
struct ClientInfo
|
|
{
|
|
DWORD id = 0;
|
|
DWORD pid = 0;
|
|
LONG64 last_tick = 0;
|
|
LONG active = 1;
|
|
HANDLE pipe = nullptr;
|
|
std::wstring label;
|
|
};
|
|
|
|
static HANDLE g_exit_event = nullptr;
|
|
static LONG g_shutdown_requested = 0;
|
|
static LONG g_next_client_id = 0;
|
|
static LONG g_seen_client = 0;
|
|
static LONG64 g_session_id = 0;
|
|
static std::mutex g_clients_mutex;
|
|
static std::vector<std::shared_ptr<ClientInfo>> g_clients;
|
|
static constexpr DWORD kFirstClientTimeoutMs = 3000;
|
|
|
|
|
|
static void ExitServerProcess(UINT exit_code)
|
|
{
|
|
printf("[server] ExitProcess code=%u\n", exit_code);
|
|
fflush(stdout);
|
|
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;
|
|
}
|
|
|
|
static std::vector<std::shared_ptr<ClientInfo>> SnapshotClients()
|
|
{
|
|
std::lock_guard<std::mutex> lock(g_clients_mutex);
|
|
return g_clients;
|
|
}
|
|
|
|
static void RemoveClient(DWORD id)
|
|
{
|
|
std::lock_guard<std::mutex> lock(g_clients_mutex);
|
|
for (auto& client : g_clients)
|
|
{
|
|
if (client && client->id == id)
|
|
InterlockedExchange(&client->active, 0);
|
|
}
|
|
|
|
g_clients.erase(
|
|
std::remove_if(g_clients.begin(), g_clients.end(),
|
|
[id](const std::shared_ptr<ClientInfo>& c) { return c && c->id == id; }),
|
|
g_clients.end());
|
|
|
|
if (InterlockedCompareExchange(&g_seen_client, 0, 0) != 0 && g_clients.empty())
|
|
{
|
|
printf("[server] all clients exited, server shutdown\n");
|
|
fflush(stdout);
|
|
SetEvent(g_exit_event);
|
|
}
|
|
}
|
|
|
|
static std::shared_ptr<ClientInfo> RegisterClient(HANDLE pipe, const Msg& hello)
|
|
{
|
|
auto client = std::make_shared<ClientInfo>();
|
|
client->id = static_cast<DWORD>(InterlockedIncrement(&g_next_client_id));
|
|
client->pid = hello.pid;
|
|
client->last_tick = static_cast<LONG64>(hello.tick);
|
|
client->pipe = pipe;
|
|
InterlockedExchange(&g_seen_client, 1);
|
|
|
|
std::lock_guard<std::mutex> lock(g_clients_mutex);
|
|
g_clients.push_back(client);
|
|
return client;
|
|
}
|
|
|
|
static void BroadcastExit(const char* reason)
|
|
{
|
|
auto clients = SnapshotClients();
|
|
Msg msg{};
|
|
msg.type = static_cast<uint32_t>(MsgType::Exit);
|
|
msg.session_id = static_cast<uint64_t>(g_session_id);
|
|
|
|
printf("[server] broadcast exit: %s\n", reason ? reason : "unknown");
|
|
fflush(stdout);
|
|
|
|
for (auto& client : clients)
|
|
{
|
|
if (client && client->pipe)
|
|
WriteExact(client->pipe, &msg, sizeof(msg));
|
|
}
|
|
}
|
|
|
|
void RequestShutdown(const char* reason)
|
|
{
|
|
if (InterlockedCompareExchange(&g_shutdown_requested, 1, 0) == 0)
|
|
{
|
|
SetEvent(g_exit_event);
|
|
BroadcastExit(reason);
|
|
}
|
|
}
|
|
|
|
static bool ReadMessage(HANDLE pipe, Msg& msg)
|
|
{
|
|
return ReadExact(pipe, &msg, sizeof(msg));
|
|
}
|
|
|
|
static bool WriteMessage(HANDLE pipe, const Msg& msg)
|
|
{
|
|
return WriteExact(pipe, &msg, sizeof(msg));
|
|
}
|
|
|
|
static bool SendCommandReply(
|
|
HANDLE pipe,
|
|
const Msg& req,
|
|
uint32_t status,
|
|
const void* output,
|
|
uint32_t output_size)
|
|
{
|
|
Msg reply{};
|
|
reply.type = static_cast<uint32_t>(MsgType::CommandReply);
|
|
reply.client_id = req.client_id;
|
|
reply.pid = GetCurrentProcessId();
|
|
reply.version = kProtocolVersion;
|
|
reply.command = req.command;
|
|
reply.status = status;
|
|
reply.input_size = 0;
|
|
reply.output_size = 0;
|
|
reply.tick = static_cast<uint64_t>(NowMs());
|
|
reply.session_id = static_cast<uint64_t>(g_session_id);
|
|
reply.request_id = req.request_id;
|
|
|
|
if (output && output_size)
|
|
{
|
|
if (output_size > kPayloadBytes)
|
|
return false;
|
|
|
|
memcpy(reply.payload, output, output_size);
|
|
reply.output_size = output_size;
|
|
}
|
|
|
|
return WriteMessage(pipe, reply);
|
|
}
|
|
|
|
static bool SendCommandStatus(HANDLE pipe, const Msg& req, uint32_t status)
|
|
{
|
|
return SendCommandReply(pipe, req, status, nullptr, 0);
|
|
}
|
|
|
|
static uint32_t HandlePrintf(const std::shared_ptr<ClientInfo>& client, const Msg& req)
|
|
{
|
|
if (req.input_size < sizeof(PrintfRequest))
|
|
return ERROR_INVALID_PARAMETER;
|
|
|
|
const auto* in = reinterpret_cast<const PrintfRequest*>(req.payload);
|
|
printf("[server][printf][client=%lu pid=%lu] %s\n", client->id, client->pid, in->text);
|
|
fflush(stdout);
|
|
return ERROR_SUCCESS;
|
|
}
|
|
|
|
static uint32_t HandleLoadDriver(const Msg& req)
|
|
{
|
|
return Driver::fn_get_instance()->driver_install() ? ERROR_SUCCESS : ERROR_INVALID_PARAMETER;
|
|
}
|
|
|
|
static uint32_t HandleUnloadDriver(const Msg& req)
|
|
{
|
|
if (req.input_size < sizeof(UnloadDriverRequest))
|
|
return ERROR_INVALID_PARAMETER;
|
|
|
|
const auto* in = reinterpret_cast<const UnloadDriverRequest*>(req.payload);
|
|
if (!in->service_name[0])
|
|
return ERROR_INVALID_PARAMETER;
|
|
|
|
SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_CONNECT);
|
|
if (!scm)
|
|
return GetLastError();
|
|
|
|
SC_HANDLE service = OpenServiceW(scm, in->service_name, SERVICE_STOP | DELETE | SERVICE_QUERY_STATUS);
|
|
if (!service)
|
|
{
|
|
DWORD err = GetLastError();
|
|
CloseServiceHandle(scm);
|
|
return err;
|
|
}
|
|
|
|
SERVICE_STATUS status{};
|
|
ControlService(service, SERVICE_CONTROL_STOP, &status);
|
|
|
|
DWORD err = ERROR_SUCCESS;
|
|
if (!DeleteService(service))
|
|
{
|
|
err = GetLastError();
|
|
if (err == ERROR_SERVICE_MARKED_FOR_DELETE)
|
|
err = ERROR_SUCCESS;
|
|
}
|
|
|
|
CloseServiceHandle(service);
|
|
CloseServiceHandle(scm);
|
|
return err;
|
|
}
|
|
|
|
static uint32_t HandleBreakpointRemoved(const std::shared_ptr<ClientInfo>& client, const Msg& req)
|
|
{
|
|
if (req.input_size < sizeof(BreakpointRemovedRequest))
|
|
return ERROR_INVALID_PARAMETER;
|
|
|
|
const auto* in = reinterpret_cast<const BreakpointRemovedRequest*>(req.payload);
|
|
printf("[server][breakpoint-removed][client=%lu pid=%lu] tid=%u dr0=0x%llx dr7=0x%llx\n",
|
|
client->id,
|
|
client->pid,
|
|
in->thread_id,
|
|
static_cast<unsigned long long>(in->dr0),
|
|
static_cast<unsigned long long>(in->dr7));
|
|
fflush(stdout);
|
|
return ERROR_SUCCESS;
|
|
}
|
|
|
|
static bool DispatchCommand(HANDLE pipe, const std::shared_ptr<ClientInfo>& client, const Msg& req)
|
|
{
|
|
if (!client)
|
|
return SendCommandStatus(pipe, req, ERROR_INVALID_PARAMETER);
|
|
|
|
if (req.input_size > kPayloadBytes || req.output_size > kPayloadBytes)
|
|
return SendCommandStatus(pipe, req, ERROR_INVALID_PARAMETER);
|
|
|
|
uint32_t status = ERROR_INVALID_FUNCTION;
|
|
|
|
switch (static_cast<CommandId>(req.command))
|
|
{
|
|
case CommandId::Printf:
|
|
status = HandlePrintf(client, req);
|
|
break;
|
|
|
|
case CommandId::LoadDriver:
|
|
status = HandleLoadDriver(req);
|
|
break;
|
|
|
|
case CommandId::UnloadDriver:
|
|
status = HandleUnloadDriver(req);
|
|
break;
|
|
|
|
case CommandId::QueryStatus:
|
|
status = ERROR_SUCCESS;
|
|
break;
|
|
case CommandId::BreakpointRemoved:
|
|
status = HandleBreakpointRemoved(client, req);
|
|
break;
|
|
default:
|
|
printf("[server] unknown command=%u from client=%lu pid=%lu\n",
|
|
req.command, client->id, client->pid);
|
|
fflush(stdout);
|
|
status = ERROR_INVALID_FUNCTION;
|
|
break;
|
|
}
|
|
|
|
return SendCommandStatus(pipe, req, status);
|
|
}
|
|
|
|
static DWORD WINAPI FirstClientTimeoutProc(LPVOID)
|
|
{
|
|
DWORD wait = WaitForSingleObject(g_exit_event, kFirstClientTimeoutMs);
|
|
if (wait == WAIT_TIMEOUT &&
|
|
InterlockedCompareExchange(&g_seen_client, 0, 0) == 0)
|
|
{
|
|
printf("[server] no client connected in %lu ms, server shutdown\n", kFirstClientTimeoutMs);
|
|
fflush(stdout);
|
|
RequestShutdown("first client timeout");
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static DWORD WINAPI MonitorThreadProc(LPVOID)
|
|
{
|
|
while (WaitForSingleObject(g_exit_event, kHeartbeatIntervalMs) == WAIT_TIMEOUT)
|
|
{
|
|
const LONG64 now = NowMs();
|
|
auto clients = SnapshotClients();
|
|
|
|
for (const auto& client : clients)
|
|
{
|
|
if (!client)
|
|
continue;
|
|
|
|
if (InterlockedCompareExchange(&client->active, 0, 0) == 0)
|
|
continue;
|
|
|
|
const LONG64 last = InterlockedCompareExchange64(&client->last_tick, 0, 0);
|
|
if (last != 0 && now - last > kHeartbeatTimeoutMs)
|
|
{
|
|
printf("[server] timeout pid=%lu id=%lu\n", client->pid, client->id);
|
|
fflush(stdout);
|
|
InterlockedExchange64(&client->last_tick, now);
|
|
}
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
static DWORD WINAPI ClientSessionProc(LPVOID param)
|
|
{
|
|
HANDLE pipe = reinterpret_cast<HANDLE>(param);
|
|
|
|
Msg hello{};
|
|
if (!ReadMessage(pipe, hello) || hello.type != static_cast<uint32_t>(MsgType::Hello))
|
|
{
|
|
CloseHandle(pipe);
|
|
return 0;
|
|
}
|
|
|
|
auto client = RegisterClient(pipe, hello);
|
|
printf("[server] client connected id=%lu pid=%lu session=%llu\n",
|
|
client->id, client->pid, static_cast<unsigned long long>(hello.session_id));
|
|
fflush(stdout);
|
|
|
|
Msg ack{};
|
|
ack.type = static_cast<uint32_t>(MsgType::HelloAck);
|
|
ack.client_id = client->id;
|
|
ack.pid = GetCurrentProcessId();
|
|
ack.version = kProtocolVersion;
|
|
ack.tick = static_cast<uint64_t>(NowMs());
|
|
ack.session_id = static_cast<uint64_t>(g_session_id);
|
|
if (!WriteMessage(pipe, ack))
|
|
{
|
|
RemoveClient(client->id);
|
|
CloseHandle(pipe);
|
|
return 0;
|
|
}
|
|
|
|
for (;;)
|
|
{
|
|
if (WaitForSingleObject(g_exit_event, 0) == WAIT_OBJECT_0)
|
|
break;
|
|
|
|
Msg msg{};
|
|
if (!ReadMessage(pipe, msg))
|
|
{
|
|
printf("[server] client disconnected id=%lu pid=%lu\n", client->id, client->pid);
|
|
fflush(stdout);
|
|
RemoveClient(client->id);
|
|
CloseHandle(pipe);
|
|
break;
|
|
}
|
|
|
|
if (msg.session_id != static_cast<uint64_t>(g_session_id))
|
|
{
|
|
printf("[server] client session mismatch id=%lu pid=%lu msg_session=%llu server_session=%llu\n",
|
|
client->id,
|
|
client->pid,
|
|
static_cast<unsigned long long>(msg.session_id),
|
|
static_cast<unsigned long long>(g_session_id));
|
|
fflush(stdout);
|
|
RemoveClient(client->id);
|
|
CloseHandle(pipe);
|
|
break;
|
|
}
|
|
|
|
if (msg.type == static_cast<uint32_t>(MsgType::Heartbeat))
|
|
{
|
|
InterlockedExchange64(&client->last_tick, static_cast<LONG64>(msg.tick));
|
|
client->pid = msg.pid;
|
|
}
|
|
else if (msg.type == static_cast<uint32_t>(MsgType::Goodbye))
|
|
{
|
|
printf("[server] client goodbye id=%lu pid=%lu\n", client->id, client->pid);
|
|
fflush(stdout);
|
|
RemoveClient(client->id);
|
|
CloseHandle(pipe);
|
|
break;
|
|
}
|
|
else if (msg.type == static_cast<uint32_t>(MsgType::Exit))
|
|
{
|
|
printf("[server] client exit id=%lu pid=%lu, remove only\n", client->id, client->pid);
|
|
fflush(stdout);
|
|
RemoveClient(client->id);
|
|
CloseHandle(pipe);
|
|
break;
|
|
}
|
|
else if (msg.type == static_cast<uint32_t>(MsgType::CommandRequest))
|
|
{
|
|
if (!DispatchCommand(pipe, client, msg))
|
|
{
|
|
printf("[server] command reply failed id=%lu pid=%lu err=%lu\n",
|
|
client->id, client->pid, GetLastError());
|
|
fflush(stdout);
|
|
RemoveClient(client->id);
|
|
CloseHandle(pipe);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static DWORD WINAPI AcceptThreadProc(LPVOID)
|
|
{
|
|
for (;;)
|
|
{
|
|
if (WaitForSingleObject(g_exit_event, 0) == WAIT_OBJECT_0)
|
|
return 0;
|
|
|
|
HANDLE pipe = CreateNamedPipeW(
|
|
kPipeName,
|
|
PIPE_ACCESS_DUPLEX,
|
|
PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT,
|
|
PIPE_UNLIMITED_INSTANCES,
|
|
sizeof(Msg),
|
|
sizeof(Msg),
|
|
0,
|
|
nullptr);
|
|
|
|
if (pipe == INVALID_HANDLE_VALUE)
|
|
{
|
|
printf("[server] CreateNamedPipeW failed: %lu\n", GetLastError());
|
|
Sleep(1000);
|
|
continue;
|
|
}
|
|
|
|
BOOL connected = ConnectNamedPipe(pipe, nullptr);
|
|
if (!connected)
|
|
{
|
|
DWORD err = GetLastError();
|
|
if (err != ERROR_PIPE_CONNECTED)
|
|
{
|
|
CloseHandle(pipe);
|
|
if (WaitForSingleObject(g_exit_event, 0) == WAIT_OBJECT_0)
|
|
return 0;
|
|
continue;
|
|
}
|
|
}
|
|
|
|
HANDLE session_thread = CreateThread(nullptr, 0, ClientSessionProc, pipe, 0, nullptr);
|
|
if (!session_thread)
|
|
{
|
|
printf("[server] ClientSession thread failed: %lu\n", GetLastError());
|
|
CloseHandle(pipe);
|
|
continue;
|
|
}
|
|
|
|
CloseHandle(session_thread);
|
|
}
|
|
}
|
|
|
|
static BOOL WINAPI CtrlHandler(DWORD type)
|
|
{
|
|
switch (type)
|
|
{
|
|
case CTRL_C_EVENT:
|
|
case CTRL_BREAK_EVENT:
|
|
case CTRL_CLOSE_EVENT:
|
|
case CTRL_LOGOFF_EVENT:
|
|
case CTRL_SHUTDOWN_EVENT:
|
|
RequestShutdown("console ctrl");
|
|
return TRUE;
|
|
default:
|
|
return FALSE;
|
|
}
|
|
}
|
|
|
|
static bool InitializeServer()
|
|
{
|
|
g_exit_event = CreateEventW(nullptr, TRUE, FALSE, nullptr);
|
|
if (!g_exit_event)
|
|
{
|
|
printf("[server] CreateEventW failed: %lu\n", GetLastError());
|
|
return false;
|
|
}
|
|
|
|
g_session_id = (static_cast<LONG64>(GetCurrentProcessId()) << 32) ^ NowMs();
|
|
|
|
printf("[server] started protocol=%u session=%llu pipe=%ws\n",
|
|
kProtocolVersion,
|
|
static_cast<unsigned long long>(g_session_id),
|
|
kPipeName);
|
|
fflush(stdout);
|
|
|
|
return true;
|
|
}
|
|
|
|
auto RunServerCore() -> BOOL
|
|
{
|
|
HANDLE first_client_timeout_thread = CreateThread(nullptr, 0, FirstClientTimeoutProc, nullptr, 0, nullptr);
|
|
if (!first_client_timeout_thread)
|
|
{
|
|
printf("[server] first client timeout thread failed: %lu\n", GetLastError());
|
|
return FALSE;
|
|
}
|
|
|
|
HANDLE monitor_thread = CreateThread(nullptr, 0, MonitorThreadProc, nullptr, 0, nullptr);
|
|
if (!monitor_thread)
|
|
{
|
|
printf("[server] monitor thread failed: %lu\n", GetLastError());
|
|
return 1;
|
|
}
|
|
|
|
HANDLE accept_thread = CreateThread(nullptr, 0, AcceptThreadProc, nullptr, 0, nullptr);
|
|
if (!accept_thread)
|
|
{
|
|
printf("[server] accept thread failed: %lu\n", GetLastError());
|
|
return 1;
|
|
}
|
|
|
|
WaitForSingleObject(g_exit_event, INFINITE);
|
|
|
|
BroadcastExit("server shutdown");
|
|
Sleep(400);
|
|
|
|
|
|
CloseHandle(accept_thread);
|
|
CloseHandle(monitor_thread);
|
|
CloseHandle(first_client_timeout_thread);
|
|
CloseHandle(g_exit_event);
|
|
ExitServerProcess(0);
|
|
|
|
return TRUE;
|
|
}
|
|
|
|
auto StartServer() -> BOOL
|
|
{
|
|
if (!InitializeServer())
|
|
return FALSE;
|
|
return RunServerCore();
|
|
}
|
|
|