From 9f10b73aba51e69bb47d5e2a4b9225253d923fb2 Mon Sep 17 00:00:00 2001 From: Nodir Temirkhodjaev Date: Fri, 28 Apr 2023 20:55:43 +0300 Subject: [PATCH] Driver: Add FORT_CHECK_STACK() to functions --- src/driver/fortbuf.c | 2 ++ src/driver/fortcout.c | 24 +++++++++++++++++++++++- src/driver/fortcout.h | 2 +- src/driver/fortdev.c | 16 +++++++++++++++- src/driver/fortdrv.c | 4 +++- src/driver/fortpkt.c | 6 +++++- src/driver/fortps.c | 2 ++ src/driver/fortscb.c | 4 ++++ src/driver/forttmr.c | 2 ++ src/driver/forttmr.h | 2 +- src/driver/fortwrk.c | 4 +++- 11 files changed, 61 insertions(+), 7 deletions(-) diff --git a/src/driver/fortbuf.c b/src/driver/fortbuf.c index dc8f1c67..7162d01a 100644 --- a/src/driver/fortbuf.c +++ b/src/driver/fortbuf.c @@ -314,6 +314,8 @@ static void fort_device_cancel_pending(PDEVICE_OBJECT device, PIRP irp) { UNUSED(device); + FORT_CHECK_STACK(); + ULONG_PTR info; const NTSTATUS status = fort_buffer_cancel_pending(&fort_device()->buffer, irp, &info); diff --git a/src/driver/fortcout.c b/src/driver/fortcout.c index a2208bb0..16d6cb1e 100644 --- a/src/driver/fortcout.c +++ b/src/driver/fortcout.c @@ -391,6 +391,8 @@ static void NTAPI fort_callout_connect_v4(const FWPS_INCOMING_VALUES0 *inFixedVa const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -418,6 +420,8 @@ static void NTAPI fort_callout_connect_v6(const FWPS_INCOMING_VALUES0 *inFixedVa const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -445,6 +449,8 @@ static void NTAPI fort_callout_accept_v4(const FWPS_INCOMING_VALUES0 *inFixedVal const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -472,6 +478,8 @@ static void NTAPI fort_callout_accept_v6(const FWPS_INCOMING_VALUES0 *inFixedVal const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -518,6 +526,8 @@ static void NTAPI fort_callout_stream_classify(const FWPS_INCOMING_VALUES0 *inFi const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FWPS_STREAM_CALLOUT_IO_PACKET0 *packet = layerData; const FWPS_STREAM_DATA0 *streamData = packet->streamData; @@ -561,6 +571,8 @@ static void NTAPI fort_callout_datagram_classify_v4(const FWPS_INCOMING_VALUES0 const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -581,6 +593,8 @@ static void NTAPI fort_callout_datagram_classify_v6(const FWPS_INCOMING_VALUES0 const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -631,6 +645,8 @@ static void NTAPI fort_callout_transport_classify_in(const FWPS_INCOMING_VALUES0 const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -648,6 +664,8 @@ static void NTAPI fort_callout_transport_classify_out(const FWPS_INCOMING_VALUES const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PVOID layerData, const FWPS_FILTER0 *filter, UINT64 flowContext, FWPS_CLASSIFY_OUT0 *classifyOut) { + FORT_CHECK_STACK(); + FORT_CALLOUT_ARG ca = { .inFixedValues = inFixedValues, .inMetaValues = inMetaValues, @@ -869,6 +887,8 @@ static NTSTATUS fort_callout_force_reauth_prov( FORT_API NTSTATUS fort_callout_force_reauth(const FORT_CONF_FLAGS old_conf_flags) { + FORT_CHECK_STACK(); + NTSTATUS status; /* Check app group periods & update group_bits */ @@ -952,8 +972,10 @@ inline static void fort_callout_flush_stat_traf( } } -FORT_API void NTAPI fort_callout_timer(void) +FORT_API void fort_callout_timer(void) { + FORT_CHECK_STACK(); + PFORT_BUFFER buf = &fort_device()->buffer; PFORT_STAT stat = &fort_device()->stat; diff --git a/src/driver/fortcout.h b/src/driver/fortcout.h index 25f945ea..c4391f85 100644 --- a/src/driver/fortcout.h +++ b/src/driver/fortcout.h @@ -15,7 +15,7 @@ FORT_API void fort_callout_remove(void); FORT_API NTSTATUS fort_callout_force_reauth(const FORT_CONF_FLAGS old_conf_flags); -FORT_API void NTAPI fort_callout_timer(void); +FORT_API void fort_callout_timer(void); #ifdef __cplusplus } // extern "C" diff --git a/src/driver/fortdev.c b/src/driver/fortdev.c index fe63ba92..961ba0e1 100644 --- a/src/driver/fortdev.c +++ b/src/driver/fortdev.c @@ -25,6 +25,8 @@ FORT_API void fort_device_set(PFORT_DEVICE device) static void fort_worker_reauth(void) { + FORT_CHECK_STACK(); + const FORT_CONF_FLAGS conf_flags = fort_device()->conf.conf_flags; const NTSTATUS status = fort_callout_force_reauth(conf_flags); @@ -35,8 +37,10 @@ static void fort_worker_reauth(void) } } -static void NTAPI fort_app_period_timer(void) +static void fort_app_period_timer(void) { + FORT_CHECK_STACK(); + if (fort_conf_ref_period_update(&fort_device()->conf, /*force=*/FALSE, /*periods_n=*/NULL)) { fort_worker_queue(&fort_device()->worker, FORT_WORKER_REAUTH); } @@ -53,6 +57,8 @@ FORT_API NTSTATUS fort_device_create(PDEVICE_OBJECT device, PIRP irp) { UNUSED(device); + FORT_CHECK_STACK(); + NTSTATUS status = STATUS_SUCCESS; /* Device opened */ @@ -75,6 +81,8 @@ FORT_API NTSTATUS fort_device_close(PDEVICE_OBJECT device, PIRP irp) { UNUSED(device); + FORT_CHECK_STACK(); + fort_request_complete(irp, STATUS_SUCCESS); return STATUS_SUCCESS; @@ -84,6 +92,8 @@ FORT_API NTSTATUS fort_device_cleanup(PDEVICE_OBJECT device, PIRP irp) { UNUSED(device); + FORT_CHECK_STACK(); + /* Device closed */ fort_device_flag_set( &fort_device()->conf, (FORT_DEVICE_IS_OPENED | FORT_DEVICE_IS_VALIDATED), FALSE); @@ -291,6 +301,8 @@ FORT_API NTSTATUS fort_device_control(PDEVICE_OBJECT device, PIRP irp) { UNUSED(device); + FORT_CHECK_STACK(); + ULONG_PTR info = 0; const PIO_STACK_LOCATION irp_stack = IoGetCurrentIrpStackLocation(irp); @@ -312,6 +324,8 @@ FORT_API NTSTATUS fort_device_shutdown(PDEVICE_OBJECT device, PIRP irp) { UNUSED(device); + FORT_CHECK_STACK(); + if (fort_device() != NULL) { fort_stat_close_flows(&fort_device()->stat); } diff --git a/src/driver/fortdrv.c b/src/driver/fortdrv.c index 0a03c2df..f647a95a 100644 --- a/src/driver/fortdrv.c +++ b/src/driver/fortdrv.c @@ -9,7 +9,7 @@ #include "forttrace.h" #include "fortutl.h" -#define FORT_STACK_USAGE_FREE_MIN (1 * 1024) +#define FORT_STACK_USAGE_FREE_MIN 400 FORT_API void fort_check_stack_usage(const char *func_name) { @@ -83,6 +83,8 @@ static NTSTATUS fort_driver_create_device(PDRIVER_OBJECT driver, PUNICODE_STRING static void fort_driver_unload(PDRIVER_OBJECT driver) { + FORT_CHECK_STACK(); + if (fort_device() == NULL) return; diff --git a/src/driver/fortpkt.c b/src/driver/fortpkt.c index 12d27a36..e8c123ee 100644 --- a/src/driver/fortpkt.c +++ b/src/driver/fortpkt.c @@ -372,6 +372,8 @@ static void NTAPI fort_packet_inject_complete( { UNUSED(dispatchLevel); + FORT_CHECK_STACK(); + switch (pkt->flags & FORT_PACKET_TYPE_MASK) { case FORT_PACKET_TYPE_FLOW: { fort_shaper_packet_free(&fort_device()->shaper, (PFORT_FLOW_PACKET) pkt, clonedNetBufList); @@ -818,8 +820,10 @@ inline static ULONG fort_shaper_timer_process_queues(PFORT_SHAPER shaper, ULONG return new_active_io_bits; } -static void NTAPI fort_shaper_timer_process(void) +static void fort_shaper_timer_process(void) { + FORT_CHECK_STACK(); + PFORT_SHAPER shaper = &fort_device()->shaper; ULONG active_io_bits = diff --git a/src/driver/fortps.c b/src/driver/fortps.c index aeaadd24..abb6632f 100644 --- a/src/driver/fortps.c +++ b/src/driver/fortps.c @@ -564,6 +564,8 @@ static void NTAPI fort_pstree_notify( { UNUSED(process); + FORT_CHECK_STACK(); + PFORT_PSTREE ps_tree = &fort_device()->ps_tree; PFORT_PSNODE proc = fort_pstree_notify_process(ps_tree, process, processHandle, createInfo); diff --git a/src/driver/fortscb.c b/src/driver/fortscb.c index 17f4e776..c45e3635 100644 --- a/src/driver/fortscb.c +++ b/src/driver/fortscb.c @@ -41,6 +41,8 @@ static void NTAPI fort_syscb_power(PVOID context, PVOID event, PVOID specifics) { UNUSED(context); + FORT_CHECK_STACK(); + if (event != (PVOID) PO_CB_SYSTEM_STATE_LOCK) return; @@ -71,6 +73,8 @@ static void NTAPI fort_syscb_time(PVOID context, PVOID event, PVOID specifics) UNUSED(event); UNUSED(specifics); + FORT_CHECK_STACK(); + fort_stat_flags_set(&fort_device()->stat, FORT_STAT_SYSTEM_TIME_CHANGED, TRUE); fort_device_on_system_time(); diff --git a/src/driver/forttmr.c b/src/driver/forttmr.c index 93b3fb9d..0003bed1 100644 --- a/src/driver/forttmr.c +++ b/src/driver/forttmr.c @@ -25,6 +25,8 @@ static void NTAPI fort_timer_callback(PKDPC dpc, PFORT_TIMER timer, PVOID arg1, UNUSED(arg1); UNUSED(arg2); + FORT_CHECK_STACK(); + const UCHAR flags = fort_timer_flags(timer); if ((flags & FORT_TIMER_ONESHOT) != 0) { fort_timer_flags_set(timer, FORT_TIMER_RUNNING, FALSE); diff --git a/src/driver/forttmr.h b/src/driver/forttmr.h index 10916917..e05fea98 100644 --- a/src/driver/forttmr.h +++ b/src/driver/forttmr.h @@ -3,7 +3,7 @@ #include "fortdrv.h" -typedef void(NTAPI *FORT_TIMER_FUNC)(void); +typedef void (*FORT_TIMER_FUNC)(void); #define FORT_TIMER_RUNNING 0x01 #define FORT_TIMER_ONESHOT 0x02 diff --git a/src/driver/fortwrk.c b/src/driver/fortwrk.c index 85553db2..14286f0d 100644 --- a/src/driver/fortwrk.c +++ b/src/driver/fortwrk.c @@ -17,7 +17,7 @@ static void fort_worker_callback_run( static NTSTATUS fort_worker_callback_expand(PVOID context) { - PFORT_WORKER worker = (PFORT_WORKER) context; + PFORT_WORKER worker = context; InterlockedDecrement16(&worker->queue_size); @@ -32,6 +32,8 @@ static void NTAPI fort_worker_callback(PDEVICE_OBJECT device, PVOID context) { UNUSED(device); + FORT_CHECK_STACK(); + const NTSTATUS status = fort_expand_stack(&fort_worker_callback_expand, context); UNUSED(status); }