refactor: make buffer sizes more precise

This commit is contained in:
2024-12-23 22:16:04 +01:00
parent 69d3489bcb
commit eb70f34ef2
12 changed files with 43 additions and 55 deletions
+4 -7
View File
@@ -82,14 +82,11 @@ DWORD WINAPI DllReaderThread(LPVOID param) {
void DllReaderClientThread(PipeServer* pipeServer) {
// send config as first packet
// this is the only write for this pipe
wchar_t config[WCHAR_BUFFER_SIZE];
swprintf_s(config, WCHAR_BUFFER_SIZE, L"callstack:%d;", g_config.do_dllinjection_ucallstack);
wchar_t config[DLL_CONFIG_LEN];
swprintf_s(config, DLL_CONFIG_LEN, L"callstack:%d;", g_config.do_dllinjection_ucallstack);
pipeServer->Send(config);
// Now receive only
char buffer[DATA_BUFFER_SIZE] = {0};
char* buf_ptr = buffer; // buf_ptr and rest_len are synchronized
int rest_len = 0;
while (!DllReaderThreadStop) {
std::vector<std::wstring> result = pipeServer->ReceiveBatch();
if (result.empty()) {
@@ -112,9 +109,9 @@ void DllReaderShutdown() {
// Disconnect server pipe
// Send some stuff so the ReadFile() in the reader thread returns
PipeClient pipeClient;
wchar_t buf[WCHAR_BUFFER_SIZE] = { 0 };
wchar_t buf[DLL_CONFIG_LEN] = { 0 };
pipeClient.Connect(DLL_PIPE_NAME);
pipeClient.Receive(buf, WCHAR_BUFFER_SIZE);
pipeClient.Receive(buf, DLL_CONFIG_LEN);
pipeClient.Send((wchar_t *) L"");
pipeClient.Disconnect();
}
+2 -2
View File
@@ -40,14 +40,14 @@ BOOL EnableKernelDriver(int enable, wchar_t* target) {
else {
dataToSend.enable = 0;
}
char buffer_incoming[128] = { 0 };
char buffer_incoming[128] = { 0 }; // Answer will be "OK" or "FAIL" so this is enough
DWORD bytesReturned = 0;
BOOL success = DeviceIoControl(hDevice,
IOCTL_MY_IOCTL_CODE,
(LPVOID)&dataToSend,
(DWORD)sizeof(dataToSend),
buffer_incoming,
sizeof(buffer_incoming),
sizeof(buffer_incoming), // this should get the correct size
&bytesReturned,
NULL);
if (!success) {
+2 -2
View File
@@ -120,9 +120,9 @@ void KernelReaderShutdown() {
if (! kernelPipeServer->IsConnected()) {
PipeClient pipeClient;
wchar_t buf[WCHAR_BUFFER_SIZE] = { 0 };
wchar_t buf[KRN_CONFIG_LEN] = { 0 };
pipeClient.Connect(KERNEL_PIPE_NAME);
pipeClient.Receive(buf, WCHAR_BUFFER_SIZE);
pipeClient.Receive(buf, KRN_CONFIG_LEN);
pipeClient.Send((wchar_t*)L"");
pipeClient.Disconnect();
}
+4 -4
View File
@@ -22,7 +22,7 @@ BOOL EnablePplProducer(BOOL e, wchar_t* target_name);
BOOL EnablePplProducer(BOOL e, wchar_t* target_name) {
wchar_t buffer[WCHAR_BUFFER_SIZE] = { 0 };
wchar_t buffer[PPL_CONFIG_LEN] = { 0 };
if (!pipeClient.Connect(PPL_SERVICE_PIPE_NAME)) {
LOG_A(LOG_ERROR, "ETW-TI: Error connecting to RedEdrPplService pipe: error code %ld", GetLastError());
@@ -37,7 +37,7 @@ BOOL EnablePplProducer(BOOL e, wchar_t* target_name) {
LOG_A(LOG_ERROR, "ETW-TI: Enable, but no target name given. Abort.");
return FALSE;
}
swprintf_s(buffer, WCHAR_BUFFER_SIZE, L"start:%s", target_name);
swprintf_s(buffer, PPL_CONFIG_LEN, L"start:%s", target_name);
if (!pipeClient.Send(buffer)) {
LOG_A(LOG_INFO, "ETW-TI: Error sending: %s to ppl service", buffer);
return FALSE;
@@ -45,8 +45,8 @@ BOOL EnablePplProducer(BOOL e, wchar_t* target_name) {
LOG_A(LOG_INFO, "ETW-TI: ppl reader: Enabled");
}
else {
//wcscpy_s(buffer, DATA_BUFFER_SIZE, L"stop");
swprintf_s(buffer, WCHAR_BUFFER_SIZE, L"stop");
//wcscpy_s(buffer, PPL_CONFIG_LEN, L"stop");
swprintf_s(buffer, PPL_CONFIG_LEN, L"%s", L"stop");
if (!pipeClient.Send(buffer)) {
return FALSE;
}
+7 -22
View File
@@ -44,8 +44,8 @@ void InitDllPipe() {
// Retrieve config (first packet)
// this is the only time we read from this pipe
LOG_A(LOG_INFO, "Waiting for config...");
wchar_t buffer[WCHAR_BUFFER_SIZE];
if (pipeClient.Receive(buffer, WCHAR_BUFFER_SIZE)) {
wchar_t buffer[DLL_CONFIG_LEN];
if (pipeClient.Receive(buffer, DLL_CONFIG_LEN)) {
if (wcsstr(buffer, L"callstack:1") != NULL) {
Config.do_stacktrace = true;
LOG_W(LOG_INFO, L"Config: Callstack Enabled");
@@ -133,14 +133,14 @@ size_t LogMyStackTrace(wchar_t* buf, size_t buf_size) {
// dont go too deep
break;
}
if (buf_size > DATA_BUFFER_SIZE - 2) { // -2 for ending ]
/*if (buf_size > DATA_BUFFER_SIZE - 2) { // -2 for ending ]
// as buf_size is size_t, it will underflow when too much callstack is appended
LOG_A(LOG_WARNING, "StackWalk: Not enough space for whole stack, stopped at %i", n);
break;
}
}*/
ProcessAddrInfoRet processAddrInfoRet = ProcessAddrInfo(hProcess, address);
w = swprintf_s(buf, WCHAR_BUFFER_SIZE, L"{\"idx\":%i,\"addr\":%llu,\"page_addr\":%llu,\"size\":%zu,\"state\":%lu,\"protect\":\"%s\",\"type\":\"%s\"},",
w = swprintf_s(buf, buf_size, L"{\"idx\":%i,\"addr\":%llu,\"page_addr\":%llu,\"size\":%zu,\"state\":%lu,\"protect\":\"%s\",\"type\":\"%s\"},",
n,
address,
processAddrInfoRet.base_addr,
@@ -149,27 +149,12 @@ size_t LogMyStackTrace(wchar_t* buf, size_t buf_size) {
processAddrInfoRet.protectStr.c_str(),
processAddrInfoRet.typeStr.c_str());
if (w == 0) {
LOG_A(LOG_ERROR, "Error");
LOG_A(LOG_ERROR, "Error writing callstack entry, not enough space? %d", buf_size);
break;
}
buf_size -= w;
buf += w;
written += w;
// Resolve the symbol at this address
/*char symbolBuffer[sizeof(SYMBOL_INFO) + MAX_SYM_NAME * sizeof(TCHAR)];
PSYMBOL_INFO pSymbol = (PSYMBOL_INFO)symbolBuffer;
pSymbol->SizeOfStruct = sizeof(SYMBOL_INFO);
pSymbol->MaxNameLen = MAX_SYM_NAME;
if (SymFromAddr(hProcess, address, 0, pSymbol))
{
printf(" %s - 0x%0llX\n", pSymbol->Name, pSymbol->Address);
}
else
{
printf(" [Unknown symbol] - 0x%0llX\n", address);
}*/
n += 1;
}
+2 -2
View File
@@ -54,7 +54,7 @@ NTSTATUS MyDriverDeviceControl(PDEVICE_OBJECT DeviceObject, PIRP Irp) {
g_config.enable_kapc_injection = 1;
}
g_config.enable_logging = 1;
wcscpy_s(g_config.target, sizeof(g_config.target), data->filename);
wcscpy_s(g_config.target, TARGET_WSTR_LEN, data->filename);
if (!IsUserspacePipeConnected()) {
int ret = ConnectUserspacePipe();
@@ -75,7 +75,7 @@ NTSTATUS MyDriverDeviceControl(PDEVICE_OBJECT DeviceObject, PIRP Irp) {
LOG_A(LOG_INFO, "[IOCTL] Stop\n");
g_config.enable_kapc_injection = 0;
g_config.enable_logging = 0;
wcscpy_s(g_config.target, sizeof(g_config.target), data->filename); // should be zero
wcscpy_s(g_config.target, TARGET_WSTR_LEN, data->filename); // should be zero
DisconnectUserspacePipe();
answer = "OK";
}
+3 -1
View File
@@ -3,6 +3,8 @@
#include <string.h>
#include <stdio.h>
#include "../Shared/common.h"
typedef struct _config {
int init_processnotify;
int init_threadnotify;
@@ -14,7 +16,7 @@ typedef struct _config {
HANDLE trace_pid;
int trace_children;
WCHAR target[256]; // zero length means disabled
WCHAR target[TARGET_WSTR_LEN]; // zero length means disabled
} Config;
+3 -3
View File
@@ -137,18 +137,18 @@ void LoadImageNotifyRoutine(PUNICODE_STRING FullImageName, HANDLE ProcessId, PIM
UNREFERENCED_PARAMETER(ImageInfo);
wchar_t line[DATA_BUFFER_SIZE] = { 0 };
wchar_t ImageName[256] = { 0 };
wchar_t ImageName[PATH_LEN] = { 0 };
// We may only have KAPC injection, and no logging
if (g_config.enable_logging) {
PROCESS_INFO* procInfo = LookupProcessInfo(ProcessId);
if (procInfo != NULL && procInfo->observe) {
UnicodeStringToWChar(FullImageName, ImageName, 256);
UnicodeStringToWChar(FullImageName, ImageName, PATH_LEN);
swprintf(line, L"{\"type\":\"kernel\",\"time\":%llu,\"callback\":\"image_load\",\"krn_pid\":%llu,\"pid\":%llu,\"image\":\"%s\"}",
systemTime,
(unsigned __int64)PsGetCurrentProcessId(),
(unsigned __int64)ProcessId,
JsonEscape(ImageName, 256));
JsonEscape(ImageName, PATH_LEN));
LogEvent(line);
}
}
+2 -2
View File
@@ -20,7 +20,7 @@ PipeServer pipeServer(L"EtwTi");
DWORD WINAPI ServiceControlPipeThread(LPVOID param) {
wchar_t buffer[WCHAR_BUFFER_SIZE];
wchar_t buffer[PPL_CONFIG_LEN];
while (keep_running) {
LOG_W(LOG_INFO, L"Control: Waiting for client (RedEdr.exe) to connect...");
@@ -30,7 +30,7 @@ DWORD WINAPI ServiceControlPipeThread(LPVOID param) {
}
while (keep_running) {
memset(buffer, 0, sizeof(buffer));
if (!pipeServer.Receive(buffer, WCHAR_BUFFER_SIZE)) {
if (!pipeServer.Receive(buffer, PPL_CONFIG_LEN)) {
//LOG_A(LOG_ERROR, "Error waiting for RedEdr.exe config");
break;
}
+2 -2
View File
@@ -22,8 +22,8 @@ BOOL ConnectEmitterPipe() {
// Retrieve config (first packet)
// this is the only read for this pipe
wchar_t buffer[WCHAR_BUFFER_SIZE];
if (pipeClient.Receive(buffer, WCHAR_BUFFER_SIZE)) {
wchar_t buffer[PPL_CONFIG_LEN];
if (pipeClient.Receive(buffer, PPL_CONFIG_LEN)) {
// Ignore config atm
}
+4 -4
View File
@@ -185,12 +185,12 @@ void LOG_A(int verbosity, const char* format, ...)
void LOG_W(int verbosity, const wchar_t* format, ...)
{
WCHAR message[WCHAR_BUFFER_SIZE] = L"[RedEdr PPL] ";
WCHAR message[DATA_BUFFER_SIZE] = L"[RedEdr PPL] ";
size_t offset = wcslen(message);
va_list arg_ptr;
va_start(arg_ptr, format);
int ret = vswprintf(&message[offset], WCHAR_BUFFER_SIZE - offset, format, arg_ptr);
int ret = vswprintf(&message[offset], DATA_BUFFER_SIZE - offset, format, arg_ptr);
va_end(arg_ptr);
Microsoft::VisualStudio::CppUnitTestFramework::Logger::WriteMessage(message);
@@ -226,12 +226,12 @@ void LOG_A(int verbosity, const char* format, ...)
void LOG_W(int verbosity, const wchar_t* format, ...)
{
WCHAR message[WCHAR_BUFFER_SIZE] = L"[RedEdr PPL] ";
WCHAR message[DATA_BUFFER_SIZE] = L"[RedEdr PPL] ";
size_t offset = wcslen(message);
va_list arg_ptr;
va_start(arg_ptr, format);
int ret = vswprintf(&message[offset], WCHAR_BUFFER_SIZE - offset, format, arg_ptr);
int ret = vswprintf(&message[offset], DATA_BUFFER_SIZE - offset, format, arg_ptr);
va_end(arg_ptr);
printf("%s", message);
+8 -4
View File
@@ -4,14 +4,18 @@
#define COMMON_H
#define PIPE_BUFFER_SIZE 8192 // thats the pipe buffer (default 4096)
#define DATA_BUFFER_SIZE 4096 // all buffers for strings
#define WCHAR_BUFFER_SIZE 2048 // Just a bit smaller, used for many things
#define DATA_BUFFER_SIZE 4096 // events, most important
#define IOCTL_MY_IOCTL_CODE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define TARGET_WSTR_LEN 256
#define PATH_LEN 1024
#define DLL_CONFIG_LEN 128 // for DLL pipe
#define PPL_CONFIG_LEN 128 // for PPL pipe
#define KRN_CONFIG_LEN 128 // for Kernel pipe
typedef struct _MY_DRIVER_DATA {
wchar_t filename[256];
wchar_t filename[TARGET_WSTR_LEN];
int enable;
int dll_inject;
} MY_DRIVER_DATA, * PMY_DRIVER_DATA;