Driver: Flush flow's deferred packets on disconnect.

This commit is contained in:
Nodir Temirkhodjaev 2019-02-19 13:03:30 +05:00
parent 64aac70a0d
commit 3ada182b43
2 changed files with 116 additions and 78 deletions

View File

@ -317,14 +317,14 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
if (!fort_conf_app_blocked(&conf_ref->conf, app_index)) {
if (conf_flags.log_stat) {
const UINT64 flowId = inMetaValues->flowHandle;
const UINT64 flow_id = inMetaValues->flowHandle;
const UCHAR group_index = fort_conf_app_group_index(
&conf_ref->conf, app_index);
const BOOL is_reauth = (flags & FWP_CONDITION_FLAG_IS_REAUTHORIZE);
BOOL is_new_proc = FALSE;
NTSTATUS status;
status = fort_flow_associate(&g_device->stat, flowId, process_id,
status = fort_flow_associate(&g_device->stat, flow_id, process_id,
group_index, is_reauth, &is_new_proc);
if (!NT_SUCCESS(status)) {
@ -413,6 +413,30 @@ fort_packet_inject_complete (PFORT_PACKET pkt,
fort_defer_packet_free(&g_device->defer, pkt, clonedNetBufList, dispatchLevel);
}
static void
fort_callout_defer_packet_flush (UINT64 flow_id,
UINT32 list_bits,
BOOL dispatchLevel)
{
fort_defer_packet_flush(&g_device->defer, fort_packet_inject_complete,
flow_id, list_bits, dispatchLevel);
}
static void
fort_callout_defer_stream_flush (UINT64 flow_id,
BOOL dispatchLevel)
{
fort_defer_stream_flush(&g_device->defer, fort_packet_inject_complete,
flow_id, FALSE);
}
static void
fort_callout_defer_flush (UINT64 flow_id)
{
fort_callout_defer_packet_flush(flow_id, FORT_DEFER_FLUSH_ALL, FALSE);
fort_callout_defer_stream_flush(flow_id, FALSE);
}
static void
fort_callout_flow_classify_v4 (const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues,
UINT64 flowContext,
@ -444,6 +468,20 @@ fort_callout_stream_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
fort_callout_flow_classify_v4(inMetaValues, flowContext,
classifyOut, dataSize, inbound);
/* Flush flow's deferred TCP packets on FIN */
if (streamFlags & (FWPS_STREAM_FLAG_RECEIVE_DISCONNECT
| FWPS_STREAM_FLAG_SEND_DISCONNECT)) {
PFORT_FLOW flow = (PFORT_FLOW) flowContext;
const UCHAR flow_flags = fort_flow_flags(flow);
if (flow_flags & (FORT_FLOW_SPEED_LIMIT | FORT_FLOW_FRAGMENT)) {
fort_callout_defer_flush(flow->flow_id);
}
goto permit;
}
/* Fragment first TCP packet */
if ((streamFlags & (FWPS_STREAM_FLAG_SEND
| FWPS_STREAM_FLAG_SEND_EXPEDITED
@ -477,7 +515,7 @@ fort_callout_stream_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
}
}
/* permit: */
permit:
fort_callout_classify_permit(filter, classifyOut);
return;
@ -777,22 +815,6 @@ fort_callout_remove (void)
}
}
static void
fort_callout_defer_flush (UINT32 list_bits, BOOL dispatchLevel)
{
fort_defer_packet_flush(&g_device->defer, fort_packet_inject_complete,
list_bits, dispatchLevel);
}
static void
fort_callout_defer_flush_all (void)
{
fort_callout_defer_flush(FORT_DEFER_FLUSH_ALL, FALSE);
fort_defer_stream_flush(&g_device->defer, fort_packet_inject_complete,
FORT_DEFER_STREAM_ALL, FALSE);
}
static NTSTATUS
fort_callout_force_reauth (const FORT_CONF_FLAGS old_conf_flags,
const FORT_CONF_FLAGS conf_flags,
@ -815,7 +837,7 @@ fort_callout_force_reauth (const FORT_CONF_FLAGS old_conf_flags,
}
if (defer_flush_bits != 0) {
fort_callout_defer_flush(defer_flush_bits, FALSE);
fort_callout_defer_packet_flush(FORT_DEFER_STREAM_ALL, defer_flush_bits, FALSE);
}
if ((status = fort_prov_open(&engine)))
@ -942,7 +964,7 @@ fort_callout_timer (void)
}
/* Flush deferred packets */
fort_callout_defer_flush(defer_flush_bits, TRUE);
fort_callout_defer_packet_flush(FORT_DEFER_STREAM_ALL, defer_flush_bits, TRUE);
}
static void
@ -1135,7 +1157,7 @@ fort_power_callback (PVOID context, PVOID event, PVOID specifics)
g_device->power_off = power_off;
if (power_off) {
fort_callout_defer_flush_all();
fort_callout_defer_flush(FORT_DEFER_STREAM_ALL);
}
}
@ -1225,7 +1247,7 @@ fort_driver_unload (PDRIVER_OBJECT driver)
UNICODE_STRING device_link;
if (g_device != NULL) {
fort_callout_defer_flush_all();
fort_callout_defer_flush(FORT_DEFER_STREAM_ALL);
fort_timer_close(&g_device->app_timer);
fort_timer_close(&g_device->log_timer);

View File

@ -60,8 +60,6 @@ typedef struct fort_packet_stream {
UINT32 streamFlags;
UINT32 calloutId;
UINT64 flowId;
} FORT_PACKET_STREAM, *PFORT_PACKET_STREAM;
typedef struct fort_packet {
@ -70,6 +68,8 @@ typedef struct fort_packet {
UINT32 dataOffset : 12;
UINT32 dataSize : 18;
UINT64 flow_id;
PNET_BUFFER_LIST netBufList;
struct fort_packet *next;
@ -200,7 +200,7 @@ fort_packet_inject_stream (PFORT_DEFER defer,
if (NT_SUCCESS(status)) {
status = FwpsStreamInjectAsync0(
defer->stream_injection4_id, NULL, 0,
pkt_stream->flowId, pkt_stream->calloutId,
pkt->flow_id, pkt_stream->calloutId,
pkt_stream->layerId, pkt_stream->streamFlags,
*clonedNetBufList, pkt->dataSize,
complete_func, pkt);
@ -286,6 +286,65 @@ fort_defer_close (PFORT_DEFER defer)
tommy_arrayof_done(&defer->packets);
}
static void
fort_defer_list_add (PFORT_DEFER_LIST defer_list,
PFORT_PACKET pkt)
{
if (defer_list->packet_tail == NULL) {
defer_list->packet_head = defer_list->packet_tail = pkt;
} else {
defer_list->packet_tail->next = pkt;
defer_list->packet_tail = pkt;
}
}
static PFORT_PACKET
fort_defer_list_get (PFORT_DEFER_LIST defer_list,
PFORT_PACKET pkt_chain,
UINT64 flow_id)
{
if (defer_list->packet_head != NULL) {
if (flow_id == FORT_DEFER_STREAM_ALL) {
defer_list->packet_tail->next = pkt_chain;
pkt_chain = defer_list->packet_head;
defer_list->packet_head = defer_list->packet_tail = NULL;
} else {
PFORT_PACKET pkt_tail = pkt_chain;
PFORT_PACKET pkt_prev = NULL;
PFORT_PACKET pkt = defer_list->packet_head;
do {
PFORT_PACKET pkt_next = pkt->next;
if (flow_id == pkt->flow_id) {
if (pkt_prev == NULL) {
pkt_chain = pkt;
} else {
pkt_prev->next = pkt;
}
pkt_prev = pkt;
pkt->next = pkt_tail;
if (pkt == defer_list->packet_head) {
defer_list->packet_head = pkt_next;
}
if (pkt == defer_list->packet_tail) {
defer_list->packet_tail = NULL;
break;
}
}
pkt = pkt_next;
} while (pkt != NULL);
}
}
return pkt_chain;
}
static NTSTATUS
fort_defer_packet_add (PFORT_DEFER defer,
const FWPS_INCOMING_VALUES0 *inFixedValues,
@ -333,17 +392,13 @@ fort_defer_packet_add (PFORT_DEFER defer,
goto end;
}
if (defer_list->packet_tail == NULL) {
defer_list->packet_head = defer_list->packet_tail = pkt;
} else {
defer_list->packet_tail->next = pkt;
defer_list->packet_tail = pkt;
}
fort_defer_list_add(defer_list, pkt);
pkt->inbound = inbound;
pkt->is_stream = FALSE;
pkt->dataOffset = 0;
pkt->dataSize = 0; /* not used */
pkt->flow_id = inMetaValues->flowHandle;
pkt->netBufList = netBufList;
pkt->next = NULL;
@ -432,17 +487,13 @@ fort_defer_stream_add (PFORT_DEFER defer,
goto end;
}
if (defer->stream_list.packet_tail == NULL) {
defer->stream_list.packet_head = defer->stream_list.packet_tail = pkt;
} else {
defer->stream_list.packet_tail->next = pkt;
defer->stream_list.packet_tail = pkt;
}
fort_defer_list_add(&defer->stream_list, pkt);
pkt->inbound = inbound;
pkt->is_stream = TRUE;
pkt->dataOffset = (UINT32) streamData->dataOffset.streamDataOffset;
pkt->dataSize = (UINT32) streamData->dataLength;
pkt->flow_id = inMetaValues->flowHandle;
pkt->netBufList = streamData->netBufferListChain;
pkt->next = NULL;
@ -452,7 +503,6 @@ fort_defer_stream_add (PFORT_DEFER defer,
pkt_stream->layerId = inFixedValues->layerId;
pkt_stream->streamFlags = streamData->flags;
pkt_stream->calloutId = filter->action.calloutId;
pkt_stream->flowId = inMetaValues->flowHandle;
}
FwpsReferenceNetBufferList0(pkt->netBufList, TRUE);
@ -537,6 +587,7 @@ fort_defer_packet_inject (PFORT_DEFER defer,
static void
fort_defer_packet_flush (PFORT_DEFER defer,
FORT_INJECT_COMPLETE_FUNC complete_func,
UINT64 flow_id,
UINT32 list_bits,
BOOL dispatchLevel)
{
@ -566,12 +617,7 @@ fort_defer_packet_flush (PFORT_DEFER defer,
defer_list = &defer->lists[i];
if (defer_list->packet_head != NULL) {
defer_list->packet_tail->next = pkt_chain;
pkt_chain = defer_list->packet_head;
defer_list->packet_head = defer_list->packet_tail = NULL;
}
pkt_chain = fort_defer_list_get(defer_list, pkt_chain, flow_id);
}
/* Clear flushed bits */
@ -589,11 +635,11 @@ fort_defer_packet_flush (PFORT_DEFER defer,
static void
fort_defer_stream_flush (PFORT_DEFER defer,
FORT_INJECT_COMPLETE_FUNC complete_func,
UINT64 flowId,
UINT64 flow_id,
BOOL dispatchLevel)
{
KLOCK_QUEUE_HANDLE lock_queue;
PFORT_PACKET pkt_chain = NULL;
PFORT_PACKET pkt_chain;
if (dispatchLevel) {
KeAcquireInStackQueuedSpinLockAtDpcLevel(&defer->lock, &lock_queue);
@ -601,37 +647,7 @@ fort_defer_stream_flush (PFORT_DEFER defer,
KeAcquireInStackQueuedSpinLock(&defer->lock, &lock_queue);
}
if (flowId == FORT_DEFER_STREAM_ALL) {
pkt_chain = defer->stream_list.packet_head;
defer->stream_list.packet_head = defer->stream_list.packet_tail = NULL;
}
else if (defer->stream_list.packet_head != NULL) {
PFORT_PACKET pkt = defer->stream_list.packet_head;
do {
PFORT_PACKET pkt_next = pkt->next;
if (flowId == pkt->stream.flowId) {
if (pkt_chain != NULL) {
pkt_chain->next = pkt;
}
pkt->next = NULL;
pkt_chain = pkt;
if (pkt == defer->stream_list.packet_head) {
defer->stream_list.packet_head = pkt_next;
}
if (pkt == defer->stream_list.packet_tail) {
defer->stream_list.packet_tail = NULL;
}
}
pkt = pkt_next;
} while (pkt != NULL);
}
pkt_chain = fort_defer_list_get(&defer->stream_list, NULL, flow_id);
if (dispatchLevel) {
KeReleaseInStackQueuedSpinLockFromDpcLevel(&lock_queue);