Driver: Abort bad flows on unload

This commit is contained in:
Nodir Temirkhodjaev 2023-02-18 21:45:54 +03:00
parent b903e5e147
commit b9072120df
3 changed files with 49 additions and 21 deletions

View File

@ -27,7 +27,7 @@
# include <stddef.h> # include <stddef.h>
#else #else
# undef _WIN32_WINNT # undef _WIN32_WINNT
# define _WIN32_WINNT 0x0601 # define _WIN32_WINNT 0x0603
# define WIN32_LEAN_AND_MEAN # define WIN32_LEAN_AND_MEAN
# include <windows.h> # include <windows.h>

View File

@ -192,44 +192,52 @@ static void fort_flow_context_set(
fort_flow_context_transport_set(stat, flow_id, flowContext, isIPv6, inbound); fort_flow_context_transport_set(stat, flow_id, flowContext, isIPv6, inbound);
} }
inline static BOOL fort_flow_context_stream_remove( inline static void fort_flow_context_stream_remove(
PFORT_STAT stat, UINT64 flow_id, BOOL isIPv6, BOOL is_tcp) PFORT_STAT stat, UINT64 flow_id, BOOL isIPv6, BOOL is_tcp, NTSTATUS *status)
{ {
UINT16 layerId; UINT16 layerId;
UINT32 calloutId; UINT32 calloutId;
fort_flow_context_stream_init(stat, isIPv6, is_tcp, &layerId, &calloutId); fort_flow_context_stream_init(stat, isIPv6, is_tcp, &layerId, &calloutId);
return FwpsFlowRemoveContext0(flow_id, layerId, calloutId) != STATUS_PENDING; *status = FwpsFlowRemoveContext0(flow_id, layerId, calloutId);
} }
inline static BOOL fort_flow_context_transport_remove(PFORT_STAT stat, UINT64 flow_id, BOOL isIPv6) inline static void fort_flow_context_transport_remove(
PFORT_STAT stat, UINT64 flow_id, BOOL isIPv6, NTSTATUS *in_status, NTSTATUS *out_status)
{ {
NTSTATUS in_status;
NTSTATUS out_status;
if (isIPv6) { if (isIPv6) {
in_status = FwpsFlowRemoveContext0( *in_status = FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V6, stat->in_transport6_id); flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V6, stat->in_transport6_id);
out_status = FwpsFlowRemoveContext0( *out_status = FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V6, stat->out_transport6_id); flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V6, stat->out_transport6_id);
} else { } else {
in_status = FwpsFlowRemoveContext0( *in_status = FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V4, stat->in_transport4_id); flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V4, stat->in_transport4_id);
out_status = FwpsFlowRemoveContext0( *out_status = FwpsFlowRemoveContext0(
flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V4, stat->out_transport4_id); flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V4, stat->out_transport4_id);
} }
return in_status != STATUS_PENDING && out_status != STATUS_PENDING;
} }
static BOOL fort_flow_context_remove_id(PFORT_STAT stat, UINT64 flow_id, BOOL isIPv6, BOOL is_tcp) static BOOL fort_flow_context_remove_id(
PFORT_STAT stat, UINT64 flow_id, BOOL isIPv6, BOOL is_tcp, BOOL *pending)
{ {
const BOOL stream_res = fort_flow_context_stream_remove(stat, flow_id, isIPv6, is_tcp); NTSTATUS stream_status;
fort_flow_context_stream_remove(stat, flow_id, isIPv6, is_tcp, &stream_status);
const BOOL transport_res = fort_flow_context_transport_remove(stat, flow_id, isIPv6); NTSTATUS transport_in_status;
NTSTATUS transport_out_status;
fort_flow_context_transport_remove(
stat, flow_id, isIPv6, &transport_in_status, &transport_out_status);
return stream_res && transport_res; if (stream_status == STATUS_PENDING || transport_in_status == STATUS_PENDING
|| transport_out_status == STATUS_PENDING) {
*pending = TRUE;
return FALSE;
}
return NT_SUCCESS(stream_status) && NT_SUCCESS(transport_in_status)
&& NT_SUCCESS(transport_out_status);
} }
static void fort_flow_context_remove(PFORT_STAT stat, PFORT_FLOW flow) static void fort_flow_context_remove(PFORT_STAT stat, PFORT_FLOW flow)
@ -240,8 +248,18 @@ static void fort_flow_context_remove(PFORT_STAT stat, PFORT_FLOW flow)
const BOOL is_tcp = (flow_flags & FORT_FLOW_TCP); const BOOL is_tcp = (flow_flags & FORT_FLOW_TCP);
const BOOL isIPv6 = (flow_flags & FORT_FLOW_IP6); const BOOL isIPv6 = (flow_flags & FORT_FLOW_IP6);
if (!fort_flow_context_remove_id(stat, flow_id, isIPv6, is_tcp)) { BOOL pending = FALSE;
fort_stat_flags_set(stat, FORT_STAT_FLOW_PENDING, TRUE);
if (!fort_flow_context_remove_id(stat, flow_id, isIPv6, is_tcp, &pending)) {
if (pending) {
fort_stat_flags_set(stat, FORT_STAT_FLOW_PENDING, TRUE);
} else {
/* The flow has associated context, but FwpsFlowRemoveContext0()
* returns that there is no context as STATUS_UNSUCCESSFUL. */
#if !defined(FORT_WIN7_COMPAT)
FwpsFlowAbort0(flow_id);
#endif
}
} }
} }
@ -320,8 +338,10 @@ inline static NTSTATUS fort_flow_add_new(PFORT_STAT stat, PFORT_FLOW *flow, UINT
tommy_key_t flow_hash, BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth) tommy_key_t flow_hash, BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth)
{ {
if (is_reauth) { if (is_reauth) {
BOOL pending = FALSE;
/* Remove existing flow context after reauth. to be able to associate a flow-context */ /* Remove existing flow context after reauth. to be able to associate a flow-context */
if (!fort_flow_context_remove_id(stat, flow_id, isIPv6, is_tcp)) if (!fort_flow_context_remove_id(stat, flow_id, isIPv6, is_tcp, &pending))
return FORT_STATUS_FLOW_BLOCK; return FORT_STATUS_FLOW_BLOCK;
} }
@ -515,9 +535,11 @@ FORT_API void fort_flow_delete(PFORT_STAT stat, UINT64 flowContext)
KLOCK_QUEUE_HANDLE lock_queue; KLOCK_QUEUE_HANDLE lock_queue;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue); KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
if ((fort_stat_flags(stat) & FORT_STAT_CLOSED) == 0) { if ((fort_stat_flags(stat) & FORT_STAT_CLOSED) == 0) {
fort_flow_free(stat, flow); fort_flow_free(stat, flow);
} }
KeReleaseInStackQueuedSpinLock(&lock_queue); KeReleaseInStackQueuedSpinLock(&lock_queue);
} }

View File

@ -185,3 +185,9 @@ NTSTATUS NTAPI FwpsFlowRemoveContext0(UINT64 flowId, UINT16 layerId, UINT32 call
UNUSED(calloutId); UNUSED(calloutId);
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
NTSTATUS NTAPI FwpsFlowAbort0(UINT64 flowId)
{
UNUSED(flowId);
return STATUS_SUCCESS;
}