Driver: Expand kernel stack for workers

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

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

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;
     }

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;

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

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;

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);

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]);
     }
 }
 

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
 {