From 7e55663f30dfad11381741cb5ff040137ba89be6 Mon Sep 17 00:00:00 2001 From: Nodir Temirkhodjaev Date: Sat, 9 Nov 2024 14:04:04 +0500 Subject: [PATCH] Driver: Refactor connection reauth handling --- src/driver/common/fortconf.h | 2 ++ src/driver/fortbuf.c | 16 +++++++--------- src/driver/fortbuf.h | 2 +- src/driver/fortcout.c | 14 ++++++-------- src/driver/fortcoutarg.h | 1 - src/driver/fortdbg.h | 1 - src/driver/fortstat.c | 31 ++++++++++++++----------------- src/driver/fortstat.h | 4 ++-- 8 files changed, 32 insertions(+), 39 deletions(-) diff --git a/src/driver/common/fortconf.h b/src/driver/common/fortconf.h index b462d9df..f827ff62 100644 --- a/src/driver/common/fortconf.h +++ b/src/driver/common/fortconf.h @@ -171,8 +171,10 @@ typedef struct fort_conf_rule_flag typedef struct fort_conf_meta_conn { + UCHAR is_reauth : 1; UCHAR inbound : 1; UCHAR isIPv6 : 1; + UCHAR is_tcp : 1; UCHAR is_loopback : 1; UCHAR is_local_net : 1; UCHAR inherited : 1; diff --git a/src/driver/fortbuf.c b/src/driver/fortbuf.c index 2eae2865..90d56927 100644 --- a/src/driver/fortbuf.c +++ b/src/driver/fortbuf.c @@ -9,9 +9,9 @@ #define FORT_BUFFER_POOL_TAG 'BwfF' -static FORT_APP_PATH fort_buffer_adjust_log_path(PCFORT_APP_PATH path) +static FORT_APP_PATH fort_buffer_adjust_log_path(PCFORT_CONF_META_CONN conn) { - FORT_APP_PATH log_path = *path; + FORT_APP_PATH log_path = conn->real_path; if (log_path.len > FORT_LOG_PATH_MAX) { log_path.len = 0; /* drop too long path */ @@ -170,7 +170,7 @@ FORT_API NTSTATUS fort_buffer_blocked_write( { NTSTATUS status; - const FORT_APP_PATH log_path = fort_buffer_adjust_log_path(&conn->real_path); + const FORT_APP_PATH log_path = fort_buffer_adjust_log_path(conn); const UINT32 len = FORT_LOG_BLOCKED_SIZE(log_path.len); @@ -192,11 +192,9 @@ FORT_API NTSTATUS fort_buffer_blocked_write( FORT_API NTSTATUS fort_buffer_blocked_ip_write( PFORT_BUFFER buf, PCFORT_CONF_META_CONN conn, PIRP *irp, ULONG_PTR *info) { - FORT_CHECK_STACK(FORT_BUFFER_BLOCKED_IP_WRITE); - NTSTATUS status; - const FORT_APP_PATH log_path = fort_buffer_adjust_log_path(&conn->real_path); + const FORT_APP_PATH log_path = fort_buffer_adjust_log_path(conn); const UINT32 len = FORT_LOG_BLOCKED_IP_SIZE(log_path.len, conn->isIPv6); @@ -216,11 +214,11 @@ FORT_API NTSTATUS fort_buffer_blocked_ip_write( } FORT_API NTSTATUS fort_buffer_proc_new_write( - PFORT_BUFFER buf, UINT32 pid, PCFORT_APP_PATH path, PIRP *irp, ULONG_PTR *info) + PFORT_BUFFER buf, PCFORT_CONF_META_CONN conn, PIRP *irp, ULONG_PTR *info) { NTSTATUS status; - const FORT_APP_PATH log_path = fort_buffer_adjust_log_path(path); + const FORT_APP_PATH log_path = fort_buffer_adjust_log_path(conn); const UINT32 len = FORT_LOG_PROC_NEW_SIZE(log_path.len); @@ -231,7 +229,7 @@ FORT_API NTSTATUS fort_buffer_proc_new_write( status = fort_buffer_prepare(buf, len, &out, irp, info); if (NT_SUCCESS(status)) { - fort_log_proc_new_write(out, pid, &log_path); + fort_log_proc_new_write(out, conn->process_id, &log_path); } } KeReleaseInStackQueuedSpinLock(&lock_queue); diff --git a/src/driver/fortbuf.h b/src/driver/fortbuf.h index 7e0b4638..5a0a55fa 100644 --- a/src/driver/fortbuf.h +++ b/src/driver/fortbuf.h @@ -48,7 +48,7 @@ FORT_API NTSTATUS fort_buffer_blocked_ip_write( PFORT_BUFFER buf, PCFORT_CONF_META_CONN conn, PIRP *irp, ULONG_PTR *info); FORT_API NTSTATUS fort_buffer_proc_new_write( - PFORT_BUFFER buf, UINT32 pid, PCFORT_APP_PATH path, PIRP *irp, ULONG_PTR *info); + PFORT_BUFFER buf, PCFORT_CONF_META_CONN conn, PIRP *irp, ULONG_PTR *info); FORT_API NTSTATUS fort_buffer_xmove( PFORT_BUFFER buf, PIRP irp, PVOID out, ULONG out_len, ULONG_PTR *info); diff --git a/src/driver/fortcout.c b/src/driver/fortcout.c index 79045dfe..3c2d1092 100644 --- a/src/driver/fortcout.c +++ b/src/driver/fortcout.c @@ -101,6 +101,7 @@ static void fort_callout_ale_fill_meta_conn(PCFORT_CALLOUT_ARG ca, PFORT_CALLOUT PFORT_CONF_META_CONN conn = &cx->conn; conn->ip_proto = ca->inFixedValues->incomingValue[ca->fi->ipProto].value.uint8; + conn->is_tcp = (conn->ip_proto == IPPROTO_TCP); conn->local_port = ca->inFixedValues->incomingValue[ca->fi->localPort].value.uint16; conn->remote_port = ca->inFixedValues->incomingValue[ca->fi->remotePort].value.uint16; @@ -127,16 +128,14 @@ inline static BOOL fort_callout_ale_associate_flow( { const UINT64 flow_id = ca->inMetaValues->flowHandle; - const IPPROTO ip_proto = - (IPPROTO) ca->inFixedValues->incomingValue[ca->fi->ipProto].value.uint8; - const BOOL is_tcp = (ip_proto == IPPROTO_TCP); + fort_callout_ale_fill_meta_conn(ca, cx); const UCHAR group_index = (UCHAR) app_flags.group_index; BOOL log_stat = FALSE; - const NTSTATUS status = fort_flow_associate(&fort_device()->stat, flow_id, cx->conn.process_id, - group_index, ca->isIPv6, is_tcp, ca->inbound, cx->is_reauth, &log_stat); + const NTSTATUS status = + fort_flow_associate(&fort_device()->stat, flow_id, &cx->conn, group_index, &log_stat); if (!NT_SUCCESS(status)) { if (status != FORT_STATUS_FLOW_BLOCK) { @@ -149,8 +148,7 @@ inline static BOOL fort_callout_ale_associate_flow( } if (!log_stat) { - fort_buffer_proc_new_write(&fort_device()->buffer, cx->conn.process_id, &cx->conn.real_path, - &cx->irp, &cx->info); + fort_buffer_proc_new_write(&fort_device()->buffer, &cx->conn, &cx->irp, &cx->info); } return FALSE; @@ -522,10 +520,10 @@ static void fort_callout_ale_classify(PFORT_CALLOUT_ARG ca) } FORT_CALLOUT_ALE_EXTRA cx = { - .is_reauth = is_reauth, .conn = { .inbound = ca->inbound, .isIPv6 = ca->isIPv6, + .is_reauth = is_reauth, .remote_ip = fort_callout_meta_ip(ca, ca->fi->remoteIp), }, }; diff --git a/src/driver/fortcoutarg.h b/src/driver/fortcoutarg.h index 811d8a84..b1bff4e8 100644 --- a/src/driver/fortcoutarg.h +++ b/src/driver/fortcoutarg.h @@ -43,7 +43,6 @@ typedef const FORT_CALLOUT_ARG *PCFORT_CALLOUT_ARG; typedef struct fort_callout_ale_extra { - UCHAR is_reauth : 1; UCHAR app_data_found : 1; UCHAR drop_blocked : 1; UCHAR ignore : 1; diff --git a/src/driver/fortdbg.h b/src/driver/fortdbg.h index 1a366def..102daf13 100644 --- a/src/driver/fortdbg.h +++ b/src/driver/fortdbg.h @@ -7,7 +7,6 @@ typedef enum FORT_FUNC_ID { FORT_FUNC_UNKNOWN = 1, - FORT_BUFFER_BLOCKED_IP_WRITE, FORT_DEVICE_CANCEL_PENDING, FORT_CALLOUT_ALE_CLASSIFY, FORT_CALLOUT_TRANSPORT_CLASSIFY, diff --git a/src/driver/fortstat.c b/src/driver/fortstat.c index e5597b01..f54b8bbb 100644 --- a/src/driver/fortstat.c +++ b/src/driver/fortstat.c @@ -237,8 +237,7 @@ static void fort_flow_free(PFORT_STAT stat, PFORT_FLOW flow) stat->flow_free = flow; } -static PFORT_FLOW fort_flow_new(PFORT_STAT stat, UINT64 flow_id, const tommy_key_t flow_hash, - BOOL isIPv6, BOOL is_tcp, BOOL inbound) +static PFORT_FLOW fort_flow_new(PFORT_STAT stat, UINT64 flow_id, const tommy_key_t flow_hash) { PFORT_FLOW flow; @@ -271,32 +270,31 @@ inline static UCHAR fort_stat_group_speed_limit(PFORT_CONF_GROUP conf_group, UCH } inline static NTSTATUS fort_flow_add_new(PFORT_STAT stat, PFORT_FLOW *flow, UINT64 flow_id, - tommy_key_t flow_hash, BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth) + tommy_key_t flow_hash, PCFORT_CONF_META_CONN conn) { - *flow = fort_flow_new(stat, flow_id, flow_hash, isIPv6, is_tcp, inbound); + *flow = fort_flow_new(stat, flow_id, flow_hash); if (*flow == NULL) return STATUS_INSUFFICIENT_RESOURCES; - NTSTATUS status = fort_flow_context_set(stat, *flow, isIPv6); + NTSTATUS status = fort_flow_context_set(stat, *flow, conn->isIPv6); if (!NT_SUCCESS(status)) { fort_flow_free(stat, *flow); /* Can't remove existing context, because of possible deadlock */ - status = is_reauth ? FORT_STATUS_FLOW_BLOCK : status; + status = conn->is_reauth ? FORT_STATUS_FLOW_BLOCK : status; } return status; } -static NTSTATUS fort_flow_add(PFORT_STAT stat, UINT64 flow_id, UCHAR group_index, UINT16 proc_index, - BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth) +static NTSTATUS fort_flow_add(PFORT_STAT stat, UINT64 flow_id, PCFORT_CONF_META_CONN conn, + UINT16 proc_index, UCHAR group_index) { const tommy_key_t flow_hash = fort_flow_hash(flow_id); PFORT_FLOW flow = fort_flow_get(stat, flow_id, flow_hash); if (flow == NULL) { - const NTSTATUS status = fort_flow_add_new( - stat, &flow, flow_id, flow_hash, isIPv6, is_tcp, inbound, is_reauth); + const NTSTATUS status = fort_flow_add_new(stat, &flow, flow_id, flow_hash, conn); if (!NT_SUCCESS(status)) return status; @@ -306,8 +304,8 @@ static NTSTATUS fort_flow_add(PFORT_STAT stat, UINT64 flow_id, UCHAR group_index const UCHAR speed_limit = fort_stat_group_speed_limit(&stat->conf_group, group_index); - flow->opt.flags = speed_limit | (is_tcp ? FORT_FLOW_TCP : 0) | (isIPv6 ? FORT_FLOW_IP6 : 0) - | (inbound ? FORT_FLOW_INBOUND : 0); + flow->opt.flags = speed_limit | (conn->is_tcp ? FORT_FLOW_TCP : 0) + | (conn->isIPv6 ? FORT_FLOW_IP6 : 0) | (conn->inbound ? FORT_FLOW_INBOUND : 0); flow->opt.group_index = group_index; flow->opt.proc_index = proc_index; @@ -433,8 +431,8 @@ static NTSTATUS fort_flow_associate_proc( return STATUS_SUCCESS; } -FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, UINT32 process_id, - UCHAR group_index, BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth, BOOL *log_stat) +FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, PCFORT_CONF_META_CONN conn, + UCHAR group_index, BOOL *log_stat) { NTSTATUS status; @@ -443,12 +441,11 @@ FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, UINT32 pr BOOL is_new_proc = FALSE; PFORT_STAT_PROC proc = NULL; - status = fort_flow_associate_proc(stat, process_id, &is_new_proc, &proc); + status = fort_flow_associate_proc(stat, conn->process_id, &is_new_proc, &proc); /* Add flow */ if (NT_SUCCESS(status)) { - status = fort_flow_add( - stat, flow_id, group_index, proc->proc_index, isIPv6, is_tcp, inbound, is_reauth); + status = fort_flow_add(stat, flow_id, conn, proc->proc_index, group_index); if (NT_SUCCESS(status)) { *log_stat = proc->log_stat; diff --git a/src/driver/fortstat.h b/src/driver/fortstat.h index 97cae3e8..63f2fbde 100644 --- a/src/driver/fortstat.h +++ b/src/driver/fortstat.h @@ -163,8 +163,8 @@ FORT_API void fort_stat_conf_update(PFORT_STAT stat, const PFORT_CONF_IO conf_io FORT_API void fort_stat_conf_flags_update(PFORT_STAT stat, const FORT_CONF_FLAGS conf_flags); -FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, UINT32 process_id, - UCHAR group_index, BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth, BOOL *log_stat); +FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, PCFORT_CONF_META_CONN conn, + UCHAR group_index, BOOL *log_stat); FORT_API void fort_flow_delete(PFORT_STAT stat, UINT64 flowContext);