From e880ea9ae5a00e4b41d57c5131b9adc393cd368a Mon Sep 17 00:00:00 2001 From: Nodir Temirkhodjaev Date: Fri, 16 Feb 2024 19:46:41 +0300 Subject: [PATCH] Driver: Shaper: Rework packets processing --- src/driver/FortFirewallDriver.pro | 2 + src/driver/fortdbg.h | 1 - src/driver/fortdrv_amalg.c | 1 + src/driver/fortpkt.c | 116 ++++++++++++++++++++---------- src/driver/fortpkt.h | 11 ++- src/driver/fortthr.c | 36 ++++++++++ src/driver/fortthr.h | 23 ++++++ src/driver/wdm/um_wdm.c | 14 ++++ src/driver/wdm/um_wdm.h | 7 ++ src/ui/util/conf/confutil.cpp | 2 +- 10 files changed, 169 insertions(+), 44 deletions(-) create mode 100644 src/driver/fortthr.c create mode 100644 src/driver/fortthr.h diff --git a/src/driver/FortFirewallDriver.pro b/src/driver/FortFirewallDriver.pro index f20bb263..6c15e925 100644 --- a/src/driver/FortFirewallDriver.pro +++ b/src/driver/FortFirewallDriver.pro @@ -19,6 +19,7 @@ SOURCES += \ fortscb.c \ fortstat.c \ forttds.c \ + fortthr.c \ forttlsf.c \ forttmr.c \ forttrace.c \ @@ -57,6 +58,7 @@ HEADERS += \ fortscb.h \ fortstat.h \ forttds.h \ + fortthr.h \ forttlsf.h \ forttmr.h \ forttrace.h \ diff --git a/src/driver/fortdbg.h b/src/driver/fortdbg.h index e6527701..e28e6963 100644 --- a/src/driver/fortdbg.h +++ b/src/driver/fortdbg.h @@ -26,7 +26,6 @@ typedef enum FORT_FUNC_ID { FORT_DEVICE_LOAD, FORT_DEVICE_UNLOAD, FORT_PACKET_INJECT_COMPLETE, - FORT_SHAPER_TIMER_PROCESS, FORT_SYSCB_POWER, FORT_SYSCB_TIME, FORT_TIMER_CALLBACK, diff --git a/src/driver/fortdrv_amalg.c b/src/driver/fortdrv_amalg.c index 03856062..757bcd1d 100644 --- a/src/driver/fortdrv_amalg.c +++ b/src/driver/fortdrv_amalg.c @@ -29,6 +29,7 @@ #include "fortps.c" #include "fortstat.c" #include "fortscb.c" +#include "fortthr.c" #include "forttmr.c" #include "forttrace.c" #include "fortutl.c" diff --git a/src/driver/fortpkt.c b/src/driver/fortpkt.c index b5852233..bae33c00 100644 --- a/src/driver/fortpkt.c +++ b/src/driver/fortpkt.c @@ -17,6 +17,16 @@ typedef void FORT_SHAPER_PACKET_FOREACH_FUNC(PFORT_SHAPER, PFORT_FLOW_PACKET); +static UCHAR fort_shaper_flags_set(PFORT_SHAPER shaper, UCHAR flags, BOOL on) +{ + return on ? InterlockedOr8(&shaper->flags, flags) : InterlockedAnd8(&shaper->flags, ~flags); +} + +static UCHAR fort_shaper_flags(PFORT_SHAPER shaper) +{ + return fort_shaper_flags_set(shaper, 0, TRUE); +} + static LONG fort_shaper_io_bits_exchange(volatile LONG *io_bits, LONG v) { return InterlockedExchange(io_bits, v); @@ -83,15 +93,10 @@ inline static BOOL fort_packet_is_ipsec_tunneled(PCFORT_CALLOUT_ARG ca) return info.isTunnelMode && !info.isDeTunneled; } -static ULONG fort_packet_data_length(const PNET_BUFFER_LIST netBufList) +inline static ULONG fort_packet_data_length(PCFORT_CALLOUT_ARG ca) { - ULONG data_length = 0; - - PNET_BUFFER netBuf = NET_BUFFER_LIST_FIRST_NB(netBufList); - while (netBuf != NULL) { - data_length += NET_BUFFER_DATA_LENGTH(netBuf); - netBuf = NET_BUFFER_NEXT_NB(netBuf); - } + PNET_BUFFER netBuf = NET_BUFFER_LIST_FIRST_NB(ca->netBufList); + const ULONG data_length = NET_BUFFER_DATA_LENGTH(netBuf); return data_length; } @@ -578,14 +583,15 @@ static void fort_shaper_queue_process_bandwidth( { const UINT64 bps = queue->limit.bps; - /* Advance the available bytes */ - const UINT64 accumulated = - ((now.QuadPart - queue->last_tick.QuadPart) * bps) / shaper->qpcFrequency.QuadPart; - queue->available_bytes += accumulated; - if (fort_shaper_packet_list_is_empty(&queue->bandwidth_list) && queue->available_bytes > FORT_QUEUE_INITIAL_TOKEN_COUNT) { queue->available_bytes = FORT_QUEUE_INITIAL_TOKEN_COUNT; + } else { + /* Advance the available bytes */ + const UINT64 accumulated = + ((now.QuadPart - queue->last_tick.QuadPart) * bps) / shaper->qpcFrequency.QuadPart; + + queue->available_bytes += accumulated; } queue->last_tick = now; @@ -598,11 +604,13 @@ static void fort_shaper_queue_process_bandwidth( PFORT_FLOW_PACKET pkt_tail = NULL; PFORT_FLOW_PACKET pkt = pkt_chain; do { - if (bps != 0LL && queue->available_bytes < pkt->data_length) + const UINT64 pkt_length = pkt->data_length; + + if (bps != 0LL && queue->available_bytes < pkt_length) break; - queue->available_bytes -= pkt->data_length; - queue->queued_bytes -= pkt->data_length; + queue->available_bytes -= pkt_length; + queue->queued_bytes -= pkt_length; pkt->latency_start = now; @@ -626,6 +634,12 @@ static PFORT_FLOW_PACKET fort_shaper_queue_process_latency( const UINT32 latency_ms = queue->limit.latency_ms; + if (latency_ms == 0) { + fort_shaper_packet_list_cut_chain(&queue->latency_list, queue->latency_list.packet_tail); + + return pkt_chain; + } + const UINT64 qpcFrequency = shaper->qpcFrequency.QuadPart; const UINT64 qpcFrequencyHalfMs = qpcFrequency / 2000LL; @@ -786,17 +800,12 @@ static void fort_shaper_free_queues(PFORT_SHAPER shaper) } } -static void fort_shaper_timer_start(PFORT_SHAPER shaper) +inline static void fort_shaper_thread_set_event(PFORT_SHAPER shaper) { - const ULONG active_io_bits = fort_shaper_io_bits(&shaper->active_io_bits); - - if (active_io_bits == 0) - return; - - fort_timer_set_running(&shaper->timer, /*run=*/TRUE); + KeSetEvent(&shaper->thread_event, IO_NO_INCREMENT, FALSE); } -inline static ULONG fort_shaper_timer_process_queues(PFORT_SHAPER shaper, ULONG active_io_bits) +inline static ULONG fort_shaper_thread_process_queues(PFORT_SHAPER shaper, ULONG active_io_bits) { ULONG new_active_io_bits = 0; @@ -821,25 +830,50 @@ inline static ULONG fort_shaper_timer_process_queues(PFORT_SHAPER shaper, ULONG return new_active_io_bits; } -static void fort_shaper_timer_process(void) +inline static BOOL fort_shaper_thread_process(PFORT_SHAPER shaper) { - FORT_CHECK_STACK(FORT_SHAPER_TIMER_PROCESS); - - PFORT_SHAPER shaper = &fort_device()->shaper; - ULONG active_io_bits = fort_shaper_io_bits_set(&shaper->active_io_bits, FORT_PACKET_FLUSH_ALL, FALSE); if (active_io_bits == 0) - return; + return FALSE; - active_io_bits = fort_shaper_timer_process_queues(shaper, active_io_bits); + active_io_bits = fort_shaper_thread_process_queues(shaper, active_io_bits); if (active_io_bits != 0) { fort_shaper_io_bits_set(&shaper->active_io_bits, active_io_bits, TRUE); - fort_shaper_timer_start(shaper); + return TRUE; } + + return FALSE; +} + +static void fort_shaper_thread_loop(PVOID context) +{ + PFORT_SHAPER shaper = context; + PKEVENT thread_event = &shaper->thread_event; + + LARGE_INTEGER due; + due.QuadPart = 1LL * -10000LL /* ms -> us */; + + PLARGE_INTEGER timeout = NULL; + + do { + KeWaitForSingleObject(thread_event, Executive, KernelMode, FALSE, timeout); + + const BOOL is_active = fort_shaper_thread_process(shaper); + + timeout = is_active ? &due : NULL; + + } while ((fort_shaper_flags(shaper) & FORT_SHAPER_CLOSED) == 0); +} + +static void fort_shaper_thread_close(PFORT_SHAPER shaper) +{ + fort_shaper_thread_set_event(shaper); + + fort_thread_wait(&shaper->thread); } inline static PFORT_FLOW_PACKET fort_shaper_flush_queues(PFORT_SHAPER shaper, UINT32 group_io_bits) @@ -889,13 +923,16 @@ FORT_API void fort_shaper_open(PFORT_SHAPER shaper) KeInitializeSpinLock(&shaper->lock); - fort_timer_open( - &shaper->timer, /*period(ms)=*/1, FORT_TIMER_ONESHOT, &fort_shaper_timer_process); + KeInitializeEvent(&shaper->thread_event, SynchronizationEvent, FALSE); + + fort_thread_run(&shaper->thread, &fort_shaper_thread_loop, shaper); } FORT_API void fort_shaper_close(PFORT_SHAPER shaper) { - fort_timer_close(&shaper->timer); + fort_shaper_flags_set(shaper, FORT_SHAPER_CLOSED, TRUE); + + fort_shaper_thread_close(shaper); fort_shaper_drop_packets(shaper); fort_shaper_free_queues(shaper); @@ -1020,11 +1057,12 @@ inline static NTSTATUS fort_shaper_packet_queue( return STATUS_NO_SUCH_GROUP; /* Calculate the Packets' Data Length */ - const ULONG data_length = fort_packet_data_length(ca->netBufList); + const ULONG data_length = fort_packet_data_length(ca); /* Check the Queue for new Packet */ - if (!fort_shaper_packet_queue_check_packet(queue, data_length)) + if (!fort_shaper_packet_queue_check_packet(queue, data_length)) { return STATUS_SUCCESS; /* drop the packet */ + } /* Create the Packet */ PFORT_FLOW_PACKET pkt = fort_shaper_packet_get(shaper); @@ -1048,8 +1086,8 @@ inline static NTSTATUS fort_shaper_packet_queue( /* Packets in transport layer must be re-injected in DCP due to locking */ fort_shaper_io_bits_set(&shaper->active_io_bits, queue_bit, TRUE); - /* Start the Timer */ - fort_shaper_timer_start(shaper); + /* Process the queued packets in thread */ + fort_shaper_thread_set_event(shaper); return STATUS_SUCCESS; } diff --git a/src/driver/fortpkt.h b/src/driver/fortpkt.h index 168f225c..4f521c4d 100644 --- a/src/driver/fortpkt.h +++ b/src/driver/fortpkt.h @@ -6,7 +6,7 @@ #include "common/fortconf.h" #include "fortcoutarg.h" #include "forttds.h" -#include "forttmr.h" +#include "fortthr.h" #define FORT_PACKET_QUEUE_BAD_INDEX ((UINT16) -1) @@ -97,7 +97,7 @@ typedef struct fort_pending_packet HANDLE completion_context; } FORT_PENDING_PACKET, *PFORT_PENDING_PACKET; -#define FORT_PENDING_PROC_COUNT_MAX 1024 +#define FORT_PENDING_PROC_COUNT_MAX 1024 #define FORT_PENDING_PROC_PACKET_COUNT_MAX 3 typedef struct fort_pending_proc @@ -129,8 +129,12 @@ typedef struct fort_pending KSPIN_LOCK lock; } FORT_PENDING, *PFORT_PENDING; +#define FORT_SHAPER_CLOSED 0x01 + typedef struct fort_shaper { + UCHAR volatile flags; + UINT32 limit_io_bits; LONG volatile group_io_bits; @@ -139,7 +143,8 @@ typedef struct fort_shaper ULONG randomSeed; LARGE_INTEGER qpcFrequency; - FORT_TIMER timer; + KEVENT thread_event; + FORT_THREAD thread; PFORT_FLOW_PACKET packet_free; tommy_arrayof packets; diff --git a/src/driver/fortthr.c b/src/driver/fortthr.c new file mode 100644 index 00000000..21ecc929 --- /dev/null +++ b/src/driver/fortthr.c @@ -0,0 +1,36 @@ +/* Fort Firewall Thread */ + +#include "fortthr.h" + +#include "fortcb.h" +#include "fortdbg.h" + +FORT_API NTSTATUS fort_thread_run(PFORT_THREAD thread, PKSTART_ROUTINE routine, PVOID context) +{ + NTSTATUS status; + + thread->thread_obj = NULL; + + HANDLE hThread; + status = PsCreateSystemThread(&hThread, 0, NULL, NULL, NULL, routine, context); + if (!NT_SUCCESS(status)) + return status; + + status = ObReferenceObjectByHandle( + hThread, THREAD_ALL_ACCESS, NULL, KernelMode, &thread->thread_obj, NULL); + + ZwClose(hThread); + + return status; +} + +FORT_API void fort_thread_wait(PFORT_THREAD thread) +{ + if (thread->thread_obj == NULL) + return; + + KeWaitForSingleObject(thread->thread_obj, Executive, KernelMode, FALSE, NULL); + + ObDereferenceObject(thread->thread_obj); + thread->thread_obj = NULL; +} diff --git a/src/driver/fortthr.h b/src/driver/fortthr.h new file mode 100644 index 00000000..d33784b9 --- /dev/null +++ b/src/driver/fortthr.h @@ -0,0 +1,23 @@ +#ifndef FORTTHR_H +#define FORTTHR_H + +#include "fortdrv.h" + +typedef struct fort_thread +{ + PVOID thread_obj; +} FORT_THREAD, *PFORT_THREAD; + +#if defined(__cplusplus) +extern "C" { +#endif + +FORT_API NTSTATUS fort_thread_run(PFORT_THREAD thread, PKSTART_ROUTINE routine, PVOID context); + +FORT_API void fort_thread_wait(PFORT_THREAD thread); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // FORTTHR_H diff --git a/src/driver/wdm/um_wdm.c b/src/driver/wdm/um_wdm.c index cd22644e..df373f62 100644 --- a/src/driver/wdm/um_wdm.c +++ b/src/driver/wdm/um_wdm.c @@ -315,6 +315,20 @@ void IoQueueWorkItemEx( UNUSED(context); } +NTSTATUS PsCreateSystemThread(PHANDLE threadHandle, ULONG desiredAccess, + POBJECT_ATTRIBUTES objectAttributes, HANDLE processHandle, PVOID clientId, + PKSTART_ROUTINE startRoutine, PVOID startContext) +{ + UNUSED(threadHandle); + UNUSED(desiredAccess); + UNUSED(objectAttributes); + UNUSED(processHandle); + UNUSED(clientId); + UNUSED(startRoutine); + UNUSED(startContext); + return STATUS_SUCCESS; +} + LARGE_INTEGER KeQueryPerformanceCounter(PLARGE_INTEGER performanceFrequency) { UNUSED(performanceFrequency); diff --git a/src/driver/wdm/um_wdm.h b/src/driver/wdm/um_wdm.h index b5601ee1..affb6862 100644 --- a/src/driver/wdm/um_wdm.h +++ b/src/driver/wdm/um_wdm.h @@ -246,6 +246,9 @@ typedef IO_WORKITEM_ROUTINE *PIO_WORKITEM_ROUTINE; typedef VOID IO_WORKITEM_ROUTINE_EX(PVOID IoObject, PVOID Context, PIO_WORKITEM IoWorkItem); typedef IO_WORKITEM_ROUTINE_EX *PIO_WORKITEM_ROUTINE_EX; +typedef VOID KSTART_ROUTINE(PVOID startContext); +typedef KSTART_ROUTINE *PKSTART_ROUTINE; + typedef struct { short Year; // range [1601...] @@ -393,6 +396,10 @@ FORT_API void IoQueueWorkItem( FORT_API void IoQueueWorkItemEx( PIO_WORKITEM workItem, PIO_WORKITEM_ROUTINE_EX routine, int queueType, PVOID context); +FORT_API NTSTATUS PsCreateSystemThread(PHANDLE threadHandle, ULONG desiredAccess, + POBJECT_ATTRIBUTES objectAttributes, HANDLE processHandle, PVOID clientId, + PKSTART_ROUTINE startRoutine, PVOID startContext); + FORT_API LARGE_INTEGER KeQueryPerformanceCounter(PLARGE_INTEGER performanceFrequency); FORT_API void KeQuerySystemTime(PLARGE_INTEGER time); diff --git a/src/ui/util/conf/confutil.cpp b/src/ui/util/conf/confutil.cpp index c05e663a..66d3dbe0 100644 --- a/src/ui/util/conf/confutil.cpp +++ b/src/ui/util/conf/confutil.cpp @@ -664,7 +664,7 @@ void ConfUtil::writeLimit(fort_speed_limit *limit, quint32 kBits, quint32 buffer limit->plr = packetLoss; limit->latency_ms = latencyMsec; limit->buffer_bytes = bufferSize; - limit->bps = quint64(kBits) * (1024 / 8); + limit->bps = quint64(kBits) * (1024LL / 8); /* to bytes per second */ } void ConfUtil::writeAddressRanges(char **data, const addrranges_arr_t &addressRanges)