Driver: Shaper: Drop flow packets on closing

This commit is contained in:
Nodir Temirkhodjaev 2023-01-09 19:24:38 +03:00
parent 9a8c2ad3cf
commit e45cd80f69
4 changed files with 106 additions and 9 deletions

View File

@ -405,7 +405,14 @@ static void NTAPI fort_callout_flow_delete(UINT16 layerId, UINT32 calloutId, UIN
UNUSED(layerId);
UNUSED(calloutId);
fort_flow_delete(&fort_device()->stat, flowContext);
PFORT_STAT stat = &fort_device()->stat;
if ((fort_stat_flags(stat) & FORT_STAT_CLOSED) != 0)
return;
fort_shaper_drop_flow_packets(&fort_device()->shaper, flowContext);
fort_flow_delete(stat, flowContext);
}
static BOOL fort_callout_transport_is_ipsec_detunneled(

View File

@ -282,6 +282,44 @@ static void fort_shaper_packet_list_cut_chain(PFORT_PACKET_LIST pkt_list, PFORT_
}
}
static void fort_shaper_packet_list_cut_packet(
PFORT_PACKET_LIST pkt_list, PFORT_PACKET pkt, PFORT_PACKET pkt_prev, PFORT_PACKET pkt_next)
{
if (pkt_prev != NULL) {
pkt_prev->next = pkt_next;
} else {
pkt_list->packet_head = pkt_next;
}
if (pkt_next == NULL) {
pkt_list->packet_tail = pkt_prev;
}
}
static PFORT_PACKET fort_shaper_packet_list_get_flow_packets(
PFORT_PACKET_LIST pkt_list, PFORT_FLOW flow, PFORT_PACKET pkt_chain)
{
PFORT_PACKET pkt_prev = NULL;
PFORT_PACKET pkt = pkt_list->packet_head;
while (pkt != NULL) {
PFORT_PACKET pkt_next = pkt->next;
if (pkt->flow == flow) {
fort_shaper_packet_list_cut_packet(pkt_list, pkt, pkt_prev, pkt_next);
pkt->next = pkt_chain;
pkt_chain = pkt;
} else {
pkt_prev = pkt;
}
pkt = pkt_next;
}
return pkt_chain;
}
static void fort_shaper_queue_process_bandwidth(
PFORT_SHAPER shaper, PFORT_PACKET_QUEUE queue, const LARGE_INTEGER now)
{
@ -373,6 +411,20 @@ static PFORT_PACKET fort_shaper_queue_get_packets(PFORT_PACKET_QUEUE queue, PFOR
return pkt;
}
static PFORT_PACKET fort_shaper_queue_get_flow_packets(
PFORT_PACKET_QUEUE queue, PFORT_FLOW flow, PFORT_PACKET pkt)
{
KLOCK_QUEUE_HANDLE lock_queue;
KeAcquireInStackQueuedSpinLock(&queue->lock, &lock_queue);
pkt = fort_shaper_packet_list_get_flow_packets(&queue->bandwidth_list, flow, pkt);
pkt = fort_shaper_packet_list_get_flow_packets(&queue->latency_list, flow, pkt);
KeReleaseInStackQueuedSpinLock(&lock_queue);
return pkt;
}
inline static BOOL fort_shaper_queue_is_empty(PFORT_PACKET_QUEUE queue)
{
return fort_shaper_packet_list_is_empty(&queue->bandwidth_list)
@ -676,9 +728,9 @@ static BOOL fort_shaper_packet_queue_check_packet(PFORT_PACKET_QUEUE queue, ULON
static NTSTATUS fort_shaper_packet_queue(PFORT_SHAPER shaper,
const FWPS_INCOMING_VALUES0 *inFixedValues,
const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PNET_BUFFER_LIST netBufList,
BOOL isIPv6, BOOL inbound, UCHAR group_index)
PFORT_FLOW flow, BOOL isIPv6, BOOL inbound)
{
const UINT16 queue_index = group_index * 2 + (inbound ? 0 : 1);
const UINT16 queue_index = flow->opt.group_index * 2 + (inbound ? 0 : 1);
const UINT32 group_io_bits = fort_shaper_io_bits(&shaper->group_io_bits);
@ -706,6 +758,7 @@ static NTSTATUS fort_shaper_packet_queue(PFORT_SHAPER shaper,
fort_shaper_packet_fill(shaper, inFixedValues, inMetaValues, netBufList, pkt, isIPv6, inbound);
pkt->flow = flow;
pkt->data_length = data_length;
pkt->flags = (inbound ? FORT_PACKET_INBOUND : 0) | (isIPv6 ? FORT_PACKET_IP6 : 0);
@ -741,12 +794,48 @@ FORT_API BOOL fort_shaper_packet_process(PFORT_SHAPER shaper,
if (fort_packet_injected_by_self(injection_id, netBufList))
return FALSE;
const NTSTATUS status = fort_shaper_packet_queue(&fort_device()->shaper, inFixedValues,
inMetaValues, netBufList, isIPv6, inbound, flow->opt.group_index);
const NTSTATUS status = fort_shaper_packet_queue(
&fort_device()->shaper, inFixedValues, inMetaValues, netBufList, flow, isIPv6, inbound);
return NT_SUCCESS(status);
}
FORT_API void fort_shaper_drop_flow_packets(PFORT_SHAPER shaper, UINT64 flowContext)
{
PFORT_FLOW flow = (PFORT_FLOW) flowContext;
const UCHAR flow_flags = fort_flow_flags(flow);
const UCHAR speed_limit = (flow_flags & FORT_FLOW_SPEED_LIMIT_FLAGS);
if (speed_limit == 0)
return;
/* Collect flow's packets from Queues */
PFORT_PACKET pkt_chain = NULL;
UINT32 active_io_bits = fort_shaper_io_bits(&shaper->active_io_bits)
& (speed_limit << (flow->opt.group_index * 2));
for (int i = 0; active_io_bits != 0; ++i) {
const BOOL queue_exists = (active_io_bits & 1) != 0;
active_io_bits >>= 1;
if (!queue_exists)
continue;
PFORT_PACKET_QUEUE queue = shaper->queues[i];
if (queue == NULL)
continue;
pkt_chain = fort_shaper_queue_get_flow_packets(queue, flow, pkt_chain);
}
/* Drop the packets */
if (pkt_chain != NULL) {
fort_shaper_packet_foreach(shaper, pkt_chain, &fort_shaper_packet_drop);
}
}
FORT_API void fort_shaper_drop_packets(PFORT_SHAPER shaper)
{
fort_shaper_flush(shaper, FORT_PACKET_FLUSH_ALL, /*drop=*/TRUE);

View File

@ -31,6 +31,8 @@ typedef struct fort_packet
{
struct fort_packet *next;
PVOID flow; /* to drop on flow deletion */
LARGE_INTEGER latency_start; /* Time it was placed in the latency queue */
UINT32 data_length; /* Size of the packet (in bytes) */
@ -108,6 +110,8 @@ FORT_API BOOL fort_shaper_packet_process(PFORT_SHAPER shaper,
const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues, PNET_BUFFER_LIST netBufList,
UINT64 flowContext, BOOL inbound);
FORT_API void fort_shaper_drop_flow_packets(PFORT_SHAPER shaper, UINT64 flowContext);
FORT_API void fort_shaper_drop_packets(PFORT_SHAPER shaper);
#ifdef __cplusplus

View File

@ -514,12 +514,9 @@ FORT_API void fort_flow_delete(PFORT_STAT stat, UINT64 flowContext)
{
PFORT_FLOW flow = (PFORT_FLOW) flowContext;
if ((fort_stat_flags(stat) & FORT_STAT_CLOSED) != 0)
return; /* double check to avoid deadlock after remove-flow-context */
KLOCK_QUEUE_HANDLE lock_queue;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
if ((fort_stat_flags(stat) & FORT_STAT_CLOSED) == 0) {
{
fort_flow_free(stat, flow);
}
KeReleaseInStackQueuedSpinLock(&lock_queue);