Driver: Add "FORT_FLOW_TCP" flag

This commit is contained in:
Nodir Temirkhodjaev 2022-10-03 16:47:06 +03:00
parent 4d2ca84180
commit 0edd819b7d
2 changed files with 19 additions and 11 deletions

View File

@ -159,20 +159,29 @@ static void fort_flow_context_set(PFORT_STAT stat, PFORT_FLOW flow, BOOL isIPv6,
static void fort_flow_context_remove(PFORT_STAT stat, PFORT_FLOW flow)
{
const UINT64 flow_id = flow->flow_id;
const BOOL is_tcp = (flow->opt.flags & FORT_FLOW_TCP);
const BOOL isIPv6 = (flow->opt.flags & FORT_FLOW_IP6);
if (isIPv6) {
if (!NT_SUCCESS(FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_STREAM_V6, stat->stream6_id))) {
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_DATAGRAM_DATA_V6, stat->datagram6_id);
if (is_tcp) {
if (isIPv6) {
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_STREAM_V6, stat->stream6_id);
FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V6, stat->in_transport6_id);
FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V6, stat->out_transport6_id);
} else {
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_STREAM_V4, stat->stream4_id);
FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V4, stat->in_transport4_id);
FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V4, stat->out_transport4_id);
}
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V6, stat->in_transport6_id);
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V6, stat->out_transport6_id);
} else {
if (!NT_SUCCESS(FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_STREAM_V4, stat->stream4_id))) {
if (isIPv6) {
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_DATAGRAM_DATA_V6, stat->datagram6_id);
} else {
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_DATAGRAM_DATA_V4, stat->datagram4_id);
}
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V4, stat->in_transport4_id);
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V4, stat->out_transport4_id);
}
}
@ -262,8 +271,7 @@ static NTSTATUS fort_flow_add(PFORT_STAT stat, UINT64 flow_id, UCHAR group_index
fort_stat_proc_inc(stat, proc_index);
}
flow->opt.flags = speed_limit | (is_new_flow ? 0 : (flow->opt.flags & FORT_FLOW_XFLAGS))
| (isIPv6 ? FORT_FLOW_IP6 : 0);
flow->opt.flags = speed_limit | (is_tcp ? FORT_FLOW_TCP : 0) | (isIPv6 ? FORT_FLOW_IP6 : 0);
flow->opt.group_index = group_index;
flow->opt.proc_index = proc_index;

View File

@ -49,7 +49,7 @@ typedef struct fort_stat_proc
#define FORT_FLOW_SPEED_LIMIT (FORT_FLOW_SPEED_LIMIT_IN | FORT_FLOW_SPEED_LIMIT_OUT)
#define FORT_FLOW_DEFER_IN 0x04
#define FORT_FLOW_DEFER_OUT 0x08
#define FORT_FLOW_XFLAGS 0
#define FORT_FLOW_TCP 0x40
#define FORT_FLOW_IP6 0x80
typedef struct fort_flow_opt