From b1a1212227883a54b4c521c8de07edc564bc4cec Mon Sep 17 00:00:00 2001 From: Nodir Temirkhodjaev Date: Thu, 20 Apr 2023 12:03:34 +0300 Subject: [PATCH] Driver: Expand kernel stack for workers --- src/driver/evt/FORTEVT_MSG00001.bin | Bin 828 -> 896 bytes src/driver/evt/fortevt.h | 11 +++++++++++ src/driver/evt/fortevt.mc | 8 ++++++++ src/driver/fortdev.c | 10 ++++++---- src/driver/fortdrv.c | 2 +- src/driver/fortdrv.h | 2 ++ src/driver/fortps.c | 4 +++- src/driver/fortps.h | 2 +- src/driver/fortwrk.c | 12 +++++++++++- src/driver/fortwrk.h | 2 +- 10 files changed, 44 insertions(+), 9 deletions(-) diff --git a/src/driver/evt/FORTEVT_MSG00001.bin b/src/driver/evt/FORTEVT_MSG00001.bin index a08e79d999304cbc7ddcb5380203481a7e6ec1f2..cd4a769ecce615d600dbac614f1f12df9956190d 100644 GIT binary patch delta 164 zcmW-bu?@mN3`PI>E_V?UutFxFj%X1Q6F^OX3M8Z`15mIKJ;fg-6?5>F0P`~b7rfNIzdTw&Un^o$7rH=`3* diff --git a/src/driver/evt/fortevt.h b/src/driver/evt/fortevt.h index fdc7640a..f9ed39b1 100644 --- a/src/driver/evt/fortevt.h +++ b/src/driver/evt/fortevt.h @@ -39,6 +39,7 @@ #define FACILITY_DRIVER 0x4 #define FACILITY_SHAPER 0x5 #define FACILITY_PROCESS_TREE 0x6 +#define FACILITY_WORKER 0x7 // @@ -172,4 +173,14 @@ // #define FORT_PSTREE_ENUM_PROCESSES_ERROR ((NTSTATUS)0xC0060002L) +/* Worker */ +// +// MessageId: FORT_WORKER_CALLBACK_ERROR +// +// MessageText: +// +// Worker: Callback Error. +// +#define FORT_WORKER_CALLBACK_ERROR ((NTSTATUS)0xC0070001L) + #endif // FORTEVT_H diff --git a/src/driver/evt/fortevt.mc b/src/driver/evt/fortevt.mc index 10d17f0b..45da925c 100644 --- a/src/driver/evt/fortevt.mc +++ b/src/driver/evt/fortevt.mc @@ -21,6 +21,7 @@ FacilityNames = ( Driver = 4:FACILITY_DRIVER Shaper = 5:FACILITY_SHAPER ProcessTree = 6:FACILITY_PROCESS_TREE + Worker = 7:FACILITY_WORKER ) @@ -100,4 +101,11 @@ Language=English Enum Processes Error. . + +;/* Worker */ +MessageId=1 Facility=Worker Severity=Error SymbolicName=FORT_WORKER_CALLBACK_ERROR +Language=English +Worker: Callback Error. +. + ;#endif // FORTEVT_H diff --git a/src/driver/fortdev.c b/src/driver/fortdev.c index 6a63668a..84dbdc8a 100644 --- a/src/driver/fortdev.c +++ b/src/driver/fortdev.c @@ -21,8 +21,10 @@ FORT_API void fort_device_set(PFORT_DEVICE device) g_device = device; } -static void NTAPI fort_worker_reauth(void) +static void NTAPI fort_worker_reauth(PVOID worker) { + UNUSED(worker); + const FORT_CONF_FLAGS conf_flags = fort_device()->conf.conf_flags; const NTSTATUS status = fort_callout_force_reauth(conf_flags); @@ -212,7 +214,7 @@ static NTSTATUS fort_device_control_app(const PFORT_APP_ENTRY app_entry, ULONG l fort_conf_ref_put(&fort_device()->conf, conf_ref); if (NT_SUCCESS(status)) { - fort_worker_reauth(); + fort_worker_reauth(NULL); } return status; @@ -232,7 +234,7 @@ static NTSTATUS fort_device_control_setzones(const PFORT_CONF_ZONES zones, ULONG } else { fort_conf_zones_set(&fort_device()->conf, conf_zones); - fort_worker_reauth(); + fort_worker_reauth(NULL); return STATUS_SUCCESS; } @@ -246,7 +248,7 @@ static NTSTATUS fort_device_control_setzoneflag(const PFORT_CONF_ZONE_FLAG zone_ if (len == sizeof(FORT_CONF_ZONE_FLAG)) { fort_conf_zone_flag_set(&fort_device()->conf, zone_flag); - fort_worker_reauth(); + fort_worker_reauth(NULL); return STATUS_SUCCESS; } diff --git a/src/driver/fortdrv.c b/src/driver/fortdrv.c index ef044941..6caaefcd 100644 --- a/src/driver/fortdrv.c +++ b/src/driver/fortdrv.c @@ -92,7 +92,7 @@ static void NTAPI fort_driver_load_device_expand(PVOID param) inline static void fort_driver_load_device(NTSTATUS *status) { const NTSTATUS status_expand = KeExpandKernelStackAndCallout( - &fort_driver_load_device_expand, status, KERNEL_STACK_SIZE); + &fort_driver_load_device_expand, status, FORT_KERNEL_STACK_SIZE); if (!NT_SUCCESS(status_expand)) { *status = status_expand; diff --git a/src/driver/fortdrv.h b/src/driver/fortdrv.h index 1528499c..380de2c7 100644 --- a/src/driver/fortdrv.h +++ b/src/driver/fortdrv.h @@ -38,6 +38,8 @@ #define fort_request_complete(irp, status) fort_request_complete_info((irp), (status), 0) +#define FORT_KERNEL_STACK_SIZE (8 * 1024) + #if defined(FORT_DEBUG_STACK) # define FORT_CHECK_STACK() fort_check_stack_usage(__func__) #else diff --git a/src/driver/fortps.c b/src/driver/fortps.c index b34847ea..a78ca889 100644 --- a/src/driver/fortps.c +++ b/src/driver/fortps.c @@ -813,8 +813,10 @@ static void fort_pstree_enum_processes_loop( } } -FORT_API void NTAPI fort_pstree_enum_processes(void) +FORT_API void NTAPI fort_pstree_enum_processes(PVOID worker) { + UNUSED(worker); + NTSTATUS status; ULONG bufferSize; diff --git a/src/driver/fortps.h b/src/driver/fortps.h index 28709946..ab2f16ab 100644 --- a/src/driver/fortps.h +++ b/src/driver/fortps.h @@ -32,7 +32,7 @@ FORT_API void fort_pstree_open(PFORT_PSTREE ps_tree); FORT_API void fort_pstree_close(PFORT_PSTREE ps_tree); -FORT_API void NTAPI fort_pstree_enum_processes(void); +FORT_API void NTAPI fort_pstree_enum_processes(PVOID worker); FORT_API BOOL fort_pstree_get_proc_name( PFORT_PSTREE ps_tree, DWORD processId, PUNICODE_STRING path, BOOL *inherited); diff --git a/src/driver/fortwrk.c b/src/driver/fortwrk.c index 84fbd37c..9880a96a 100644 --- a/src/driver/fortwrk.c +++ b/src/driver/fortwrk.c @@ -7,11 +7,21 @@ #include "fortcb.h" #include "forttrace.h" +static void fort_worker_callback_run_expand(PFORT_WORKER worker, PEXPAND_STACK_CALLOUT callout) +{ + const NTSTATUS status = KeExpandKernelStackAndCallout(callout, worker, FORT_KERNEL_STACK_SIZE); + + if (!NT_SUCCESS(status)) { + LOG("Worker Callback: Error: %x\n", status); + TRACE(FORT_WORKER_CALLBACK_ERROR, status, 0, 0); + } +} + static void fort_worker_callback_run( PFORT_WORKER worker, enum FORT_WORKER_TYPE worker_type, UCHAR id_bits) { if ((id_bits & (1 << worker_type)) != 0) { - worker->funcs[worker_type](); + fort_worker_callback_run_expand(worker, worker->funcs[worker_type]); } } diff --git a/src/driver/fortwrk.h b/src/driver/fortwrk.h index a56cff67..8a2e27a9 100644 --- a/src/driver/fortwrk.h +++ b/src/driver/fortwrk.h @@ -9,7 +9,7 @@ enum FORT_WORKER_TYPE { FORT_WORKER_FUNC_COUNT, }; -typedef void(NTAPI *FORT_WORKER_FUNC)(void); +typedef void(NTAPI *FORT_WORKER_FUNC)(PVOID worker); typedef struct fort_worker {