mirrors/fort
1
0
Fork 0
mirror of https://github.com/tnodir/fort synced 2024-11-15 05:25:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Driver: Use flow-contexts to track UDP closure.

Browse Source 
Instead of FWPM_LAYER_ALE_ENDPOINT_CLOSURE_V4.
This commit is contained in:
Nodir Temirkhodjaev 2018-01-13 18:31:30 +05:00
parent ff6ab09908
commit 1d7f107398
7 changed files with 221 additions and 141 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
src/3rdparty/tommyds/tommyhashdyn.c vendored
View File

@ -223,6 +223,42 @@ void tommy_hashdyn_foreach_arg(tommy_hashdyn* hashdyn, tommy_foreach_arg_func* f
}
}
TOMMY_API //!!
void tommy_hashdyn_foreach_node(tommy_hashdyn* hashdyn, tommy_foreach_node_func* func)
{
tommy_count_t bucket_max = hashdyn->bucket_max;
tommy_hashdyn_node** bucket = hashdyn->bucket;
tommy_count_t pos;
for (pos = 0; pos < bucket_max; ++pos) {
tommy_hashdyn_node* node = bucket[pos];
while (node) {
tommy_hashdyn_node* next = node->next;
func(node);
node = next;
}
}
}
TOMMY_API //!!
void tommy_hashdyn_foreach_node_arg(tommy_hashdyn* hashdyn, tommy_foreach_node_arg_func* func, void* arg)
{
tommy_count_t bucket_max = hashdyn->bucket_max;
tommy_hashdyn_node** bucket = hashdyn->bucket;
tommy_count_t pos;
for (pos = 0; pos < bucket_max; ++pos) {
tommy_hashdyn_node* node = bucket[pos];
while (node) {
tommy_hashdyn_node* next = node->next;
func(arg, node);
node = next;
}
}
}
TOMMY_API //!!
tommy_size_t tommy_hashdyn_memory_usage(tommy_hashdyn* hashdyn)
{

12
src/3rdparty/tommyds/tommyhashdyn.h vendored
View File

@ -285,6 +285,18 @@ void tommy_hashdyn_foreach(tommy_hashdyn* hashdyn, tommy_foreach_func* func);
TOMMY_API //!!
void tommy_hashdyn_foreach_arg(tommy_hashdyn* hashdyn, tommy_foreach_arg_func* func, void* arg);
/**
* Calls the specified function for each node in the hashtable.
*/
TOMMY_API //!!
void tommy_hashdyn_foreach_node(tommy_hashdyn* hashdyn, tommy_foreach_node_func* func);
/**
* Calls the specified function with an argument for each node in the hashtable.
*/
TOMMY_API //!!
void tommy_hashdyn_foreach_node_arg(tommy_hashdyn* hashdyn, tommy_foreach_node_arg_func* func, void* arg);
/**
* Gets the number of elements.
*/

13
src/3rdparty/tommyds/tommytypes.h vendored
View File

@ -302,6 +302,19 @@ typedef void tommy_foreach_func(void* obj);
*/
typedef void tommy_foreach_arg_func(void* arg, void* obj);
/** //!!
* Foreach function.
* \param node Pointer to the node to iterate.
*/
typedef void tommy_foreach_node_func(void* node);
/** //!!
* Foreach function with an argument.
* \param arg Pointer to a generic argument.
* \param node Pointer to the node to iterate.
*/
typedef void tommy_foreach_node_arg_func(void* arg, void* node);
/******************************************************************************/
/* bit hacks */

8
src/common/common.h
View File

@ -20,10 +20,6 @@ DEFINE_GUID(FORT_GUID_CALLOUT_CONNECT_V4,
DEFINE_GUID(FORT_GUID_CALLOUT_ACCEPT_V4,
0xced9c0fc, 0xdf22, 0x4f39, 0x91, 0x85, 0xb4, 0x35, 0x72, 0x6c, 0xab, 0xb2);
/* {72201F95-359A-427A-94EF-6C9C6038AEE3} */
DEFINE_GUID(FORT_GUID_CALLOUT_CLOSURE_V4,
0x72201f95, 0x359a, 0x427a, 0x94, 0xef, 0x6c, 0x9c, 0x60, 0x38, 0xae, 0xe3);
/* {1F50005D-CDBC-42A0-A0C0-53E43081FABE} */
DEFINE_GUID(FORT_GUID_CALLOUT_STREAM_V4,
0x1f50005d, 0xcdbc, 0x42a0, 0xa0, 0xc0, 0x53, 0xe4, 0x30, 0x81, 0xfa, 0xbe);
@ -52,10 +48,6 @@ DEFINE_GUID(FORT_GUID_FILTER_CONNECT_V4,
DEFINE_GUID(FORT_GUID_FILTER_ACCEPT_V4,
0x544a3e25, 0x7beb, 0x4970, 0x88, 0xef, 0xb4, 0xbc, 0xa2, 0xce, 0x24, 0x82);
/* {93E3BC97-E338-4304-8BF3-EB21705874C1} */
DEFINE_GUID(FORT_GUID_FILTER_CLOSURE_V4,
0x93e3bc97, 0xe338, 0x4304, 0x8b, 0xf3, 0xeb, 0x21, 0x70, 0x58, 0x74, 0xc1);
/* {ED0F2527-A787-4CA2-9493-C96320422FCF} */
DEFINE_GUID(FORT_GUID_FILTER_STREAM_V4,
0xed0f2527, 0xa787, 0x4ca2, 0x94, 0x93, 0xc9, 0x63, 0x20, 0x42, 0x2f, 0xcf);

28
src/common/fortprov.c
View File

@ -20,7 +20,6 @@ fort_prov_unregister (HANDLE transEngine)
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_CONNECT_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_ACCEPT_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_CLOSURE_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_STREAM_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_DATAGRAM_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_IN_TRANSPORT_V4);
@ -30,7 +29,6 @@ fort_prov_unregister (HANDLE transEngine)
FwpmSubLayerDeleteByKey0(engine, (GUID *) &FORT_GUID_SUBLAYER);
FwpmCalloutDeleteByKey0(engine, (GUID *) &FORT_GUID_CALLOUT_CONNECT_V4);
FwpmCalloutDeleteByKey0(engine, (GUID *) &FORT_GUID_CALLOUT_ACCEPT_V4);
FwpmCalloutDeleteByKey0(engine, (GUID *) &FORT_GUID_CALLOUT_CLOSURE_V4);
FwpmCalloutDeleteByKey0(engine, (GUID *) &FORT_GUID_CALLOUT_STREAM_V4);
FwpmCalloutDeleteByKey0(engine, (GUID *) &FORT_GUID_CALLOUT_DATAGRAM_V4);
FwpmCalloutDeleteByKey0(engine, (GUID *) &FORT_GUID_CALLOUT_IN_TRANSPORT_V4);
@ -55,7 +53,6 @@ fort_prov_flow_unregister (HANDLE transEngine)
fort_prov_trans_begin(engine);
}
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_CLOSURE_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_STREAM_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_DATAGRAM_V4);
FwpmFilterDeleteByKey0(engine, (GUID *) &FORT_GUID_FILTER_IN_TRANSPORT_V4);
@ -72,7 +69,7 @@ fort_prov_register (HANDLE transEngine, BOOL is_boot)
{
FWPM_PROVIDER0 provider;
FWPM_CALLOUT0 ocallout4, icallout4;
FWPM_CALLOUT0 ccallout4, scallout4, dcallout4;
FWPM_CALLOUT0 scallout4, dcallout4;
FWPM_CALLOUT0 itcallout4, otcallout4;
FWPM_SUBLAYER0 sublayer;
FWPM_FILTER0 ofilter4, ifilter4;
@ -108,13 +105,6 @@ fort_prov_register (HANDLE transEngine, BOOL is_boot)
icallout4.providerKey = (GUID *) &FORT_GUID_PROVIDER;
icallout4.applicableLayer = FWPM_LAYER_ALE_AUTH_RECV_ACCEPT_V4;
RtlZeroMemory(&ccallout4, sizeof(FWPM_CALLOUT0));
ccallout4.calloutKey = FORT_GUID_CALLOUT_CLOSURE_V4;
ccallout4.displayData.name = (PWCHAR) L"FortCalloutClosure4";
ccallout4.displayData.description = (PWCHAR) L"Fort Firewall Callout Closure V4";
ccallout4.providerKey = (GUID *) &FORT_GUID_PROVIDER;
ccallout4.applicableLayer = FWPM_LAYER_ALE_ENDPOINT_CLOSURE_V4;
RtlZeroMemory(&scallout4, sizeof(FWPM_CALLOUT0));
scallout4.calloutKey = FORT_GUID_CALLOUT_STREAM_V4;
scallout4.displayData.name = (PWCHAR) L"FortCalloutStream4";
@ -172,7 +162,6 @@ fort_prov_register (HANDLE transEngine, BOOL is_boot)
if ((status = FwpmProviderAdd0(engine, &provider, NULL))
|| (status = FwpmCalloutAdd0(engine, &ocallout4, NULL, NULL))
|| (status = FwpmCalloutAdd0(engine, &icallout4, NULL, NULL))
|| (status = FwpmCalloutAdd0(engine, &ccallout4, NULL, NULL))
|| (status = FwpmCalloutAdd0(engine, &scallout4, NULL, NULL))
|| (status = FwpmCalloutAdd0(engine, &dcallout4, NULL, NULL))
|| (status = FwpmCalloutAdd0(engine, &itcallout4, NULL, NULL))
@ -198,7 +187,7 @@ fort_prov_register (HANDLE transEngine, BOOL is_boot)
static DWORD
fort_prov_flow_register (HANDLE transEngine, BOOL speed_limit)
{
FWPM_FILTER0 cfilter4, sfilter4, dfilter4;
FWPM_FILTER0 sfilter4, dfilter4;
FWPM_FILTER0 itfilter4, otfilter4;
HANDLE engine = transEngine;
const UINT32 filter_flags = FWPM_FILTER_FLAG_PERMIT_IF_CALLOUT_UNREGISTERED
@ -212,16 +201,6 @@ fort_prov_flow_register (HANDLE transEngine, BOOL speed_limit)
fort_prov_trans_begin(engine);
}
RtlZeroMemory(&cfilter4, sizeof(FWPM_FILTER0));
cfilter4.flags = 0;
cfilter4.filterKey = FORT_GUID_FILTER_CLOSURE_V4;
cfilter4.layerKey = FWPM_LAYER_ALE_ENDPOINT_CLOSURE_V4;
cfilter4.subLayerKey = FORT_GUID_SUBLAYER;
cfilter4.displayData.name = (PWCHAR) L"FortFilterClosure4";
cfilter4.displayData.description = (PWCHAR) L"Fort Firewall Filter Closure V4";
cfilter4.action.type = FWP_ACTION_CALLOUT_INSPECTION;
cfilter4.action.calloutKey = FORT_GUID_CALLOUT_CLOSURE_V4;
RtlZeroMemory(&sfilter4, sizeof(FWPM_FILTER0));
sfilter4.flags = filter_flags;
sfilter4.filterKey = FORT_GUID_FILTER_STREAM_V4;
@ -262,8 +241,7 @@ fort_prov_flow_register (HANDLE transEngine, BOOL speed_limit)
otfilter4.action.type = FWP_ACTION_CALLOUT_UNKNOWN;
otfilter4.action.calloutKey = FORT_GUID_CALLOUT_OUT_TRANSPORT_V4;
if ((status = FwpmFilterAdd0(engine, &cfilter4, NULL, NULL))
|| (status = FwpmFilterAdd0(engine, &sfilter4, NULL, NULL))
if ((status = FwpmFilterAdd0(engine, &sfilter4, NULL, NULL))
|| (status = FwpmFilterAdd0(engine, &dfilter4, NULL, NULL))
#if 0
|| (speed_limit && ((status = FwpmFilterAdd0(engine, &itfilter4, NULL, NULL))

118
src/driver/fortdrv.c
View File

@ -56,14 +56,10 @@ typedef struct fort_conf_ref {
typedef struct fort_device {
UINT32 prov_boot : 1;
UINT32 is_opened : 1;
UINT32 was_conf : 1;
UINT32 connect4_id;
UINT32 accept4_id;
UINT32 closure4_id;
UINT32 stream4_id;
UINT32 datagram4_id;
UINT32 in_transport4_id;
UINT32 out_transport4_id;
FORT_BUFFER buffer;
FORT_STAT stat;
@ -145,6 +141,7 @@ fort_conf_ref_set (PFORT_CONF_REF conf_ref)
if (conf_ref != NULL) {
g_device->prov_boot = conf_ref->conf.flags.prov_boot;
g_device->was_conf = TRUE;
}
}
KeReleaseInStackQueuedSpinLock(&lock_queue);
@ -241,7 +238,8 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
conf_ref = fort_conf_ref_take();
if (conf_ref == NULL) {
if (g_device->prov_boot) {
if (g_device->prov_boot || ((flags & FWP_CONDITION_FLAG_IS_REAUTHORIZE)
&& !g_device->was_conf)) { // Block existing flows after driver installation to use flow-contexts
fort_callout_classify_block(classifyOut);
} else {
fort_callout_classify_continue(classifyOut);
@ -275,16 +273,16 @@ fort_callout_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
const UINT64 flowId = inMetaValues->flowHandle;
const UCHAR group_index = fort_conf_app_group_index(
&conf_ref->conf, app_index);
BOOL is_new = FALSE;
BOOL is_new_proc = FALSE;
NTSTATUS status;
status = fort_stat_flow_associate(&g_device->stat,
flowId, process_id, group_index, &is_new);
flowId, process_id, group_index, &is_new_proc);
if (!NT_SUCCESS(status)) {
DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
"FORT: Classify v4: Flow assoc. error: %d\n", status);
} else if (is_new) {
} else if (is_new_proc) {
fort_buffer_proc_new_write(&g_device->buffer,
process_id, path_len, path, &irp, &info);
}
@ -344,25 +342,6 @@ fort_callout_accept_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
FWPS_FIELD_ALE_AUTH_RECV_ACCEPT_V4_IP_REMOTE_ADDRESS);
}
static void
fort_callout_closure_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues,
void *layerData,
const FWPS_FILTER0 *filter,
UINT64 flowContext,
FWPS_CLASSIFY_OUT0 *classifyOut)
{
const UINT64 flowId = inMetaValues->flowHandle;
UNUSED(inFixedValues);
UNUSED(layerData);
UNUSED(filter);
UNUSED(flowContext);
UNUSED(classifyOut);
fort_stat_flow_delete(&g_device->stat, flowId);
}
static NTSTATUS NTAPI
fort_callout_notify (FWPS_CALLOUT_NOTIFY_TYPE notifyType,
const GUID *filterKey, const FWPS_FILTER0 *filter)
@ -449,6 +428,15 @@ fort_callout_datagram_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
fort_callout_flow_classify_v4(inMetaValues, classifyOut, dataSize, inbound);
}
static void
fort_callout_flow_delete_v4 (UINT16 layerId, UINT32 calloutId, UINT64 flowId)
{
UNUSED(layerId);
UNUSED(calloutId);
fort_stat_flow_delete(&g_device->stat, flowId);
}
static void
fort_callout_transport_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValues,
const FWPS_INCOMING_METADATA_VALUES0 *inMetaValues,
@ -522,6 +510,14 @@ fort_callout_out_transport_classify_v4 (const FWPS_INCOMING_VALUES0 *inFixedValu
FALSE);
}
static void
fort_callout_transport_delete_v4 (UINT16 layerId, UINT32 calloutId, UINT64 flowId)
{
UNUSED(layerId);
UNUSED(calloutId);
UNUSED(flowId);
}
static NTSTATUS
fort_callout_install (PDEVICE_OBJECT device)
{
@ -554,22 +550,14 @@ fort_callout_install (PDEVICE_OBJECT device)
return status;
}
/* IPv4 closure callout */
c.calloutKey = FORT_GUID_CALLOUT_CLOSURE_V4;
c.classifyFn = fort_callout_closure_v4;
status = FwpsCalloutRegister0(device, &c, &g_device->closure4_id);
if (!NT_SUCCESS(status)) {
DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
"FORT: Register Closure V4: Error: %d\n", status);
return status;
}
/* IPv4 stream callout */
c.calloutKey = FORT_GUID_CALLOUT_STREAM_V4;
c.classifyFn = fort_callout_stream_classify_v4;
status = FwpsCalloutRegister0(device, &c, &g_device->stream4_id);
c.flowDeleteFn = fort_callout_flow_delete_v4;
c.flags = FWP_CALLOUT_FLAG_CONDITIONAL_ON_FLOW;
status = FwpsCalloutRegister0(device, &c, &g_device->stat.stream4_id);
if (!NT_SUCCESS(status)) {
DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
"FORT: Register Stream V4: Error: %d\n", status);
@ -580,7 +568,7 @@ fort_callout_install (PDEVICE_OBJECT device)
c.calloutKey = FORT_GUID_CALLOUT_DATAGRAM_V4;
c.classifyFn = fort_callout_datagram_classify_v4;
status = FwpsCalloutRegister0(device, &c, &g_device->datagram4_id);
status = FwpsCalloutRegister0(device, &c, &g_device->stat.datagram4_id);
if (!NT_SUCCESS(status)) {
DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
"FORT: Register Datagram V4: Error: %d\n", status);
@ -591,7 +579,9 @@ fort_callout_install (PDEVICE_OBJECT device)
c.calloutKey = FORT_GUID_CALLOUT_IN_TRANSPORT_V4;
c.classifyFn = fort_callout_in_transport_classify_v4;
status = FwpsCalloutRegister0(device, &c, &g_device->in_transport4_id);
c.flowDeleteFn = fort_callout_transport_delete_v4;
status = FwpsCalloutRegister0(device, &c, &g_device->stat.in_transport4_id);
if (!NT_SUCCESS(status)) {
DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
"FORT: Register Inbound Transport V4: Error: %d\n", status);
@ -602,7 +592,7 @@ fort_callout_install (PDEVICE_OBJECT device)
c.calloutKey = FORT_GUID_CALLOUT_OUT_TRANSPORT_V4;
c.classifyFn = fort_callout_out_transport_classify_v4;
status = FwpsCalloutRegister0(device, &c, &g_device->out_transport4_id);
status = FwpsCalloutRegister0(device, &c, &g_device->stat.out_transport4_id);
if (!NT_SUCCESS(status)) {
DbgPrintEx(DPFLTR_IHVNETWORK_ID, DPFLTR_ERROR_LEVEL,
"FORT: Register Outbound Transport V4: Error: %d\n", status);
@ -625,29 +615,24 @@ fort_callout_remove (void)
g_device->accept4_id = 0;
}
if (g_device->closure4_id) {
FwpsCalloutUnregisterById0(g_device->closure4_id);
g_device->closure4_id = 0;
if (g_device->stat.stream4_id) {
FwpsCalloutUnregisterById0(g_device->stat.stream4_id);
g_device->stat.stream4_id = 0;
}
if (g_device->stream4_id) {
FwpsCalloutUnregisterById0(g_device->stream4_id);
g_device->stream4_id = 0;
if (g_device->stat.datagram4_id) {
FwpsCalloutUnregisterById0(g_device->stat.datagram4_id);
g_device->stat.datagram4_id = 0;
}
if (g_device->datagram4_id) {
FwpsCalloutUnregisterById0(g_device->datagram4_id);
g_device->datagram4_id = 0;
if (g_device->stat.in_transport4_id) {
FwpsCalloutUnregisterById0(g_device->stat.in_transport4_id);
g_device->stat.in_transport4_id = 0;
}
if (g_device->in_transport4_id) {
FwpsCalloutUnregisterById0(g_device->in_transport4_id);
g_device->in_transport4_id = 0;
}
if (g_device->out_transport4_id) {
FwpsCalloutUnregisterById0(g_device->out_transport4_id);
g_device->out_transport4_id = 0;
if (g_device->stat.out_transport4_id) {
FwpsCalloutUnregisterById0(g_device->stat.out_transport4_id);
g_device->stat.out_transport4_id = 0;
}
}
@ -656,14 +641,14 @@ fort_callout_force_reauth (PDEVICE_OBJECT device,
const FORT_CONF_FLAGS old_conf_flags,
const FORT_CONF_FLAGS conf_flags)
{
PFORT_STAT stat = &g_device->stat;
HANDLE engine;
NTSTATUS status;
UNUSED(device);
if (old_conf_flags.log_stat) {
fort_stat_clear(&g_device->stat);
}
fort_stat_update(stat, conf_flags.log_stat);
if ((status = fort_prov_open(&engine)))
goto end;
@ -689,7 +674,7 @@ fort_callout_force_reauth (PDEVICE_OBJECT device,
stat:
if (conf_flags.log_stat) {
if ((status = fort_prov_flow_register(engine,
(&g_device->stat.limit_bits != 0))))
(stat->limit_bits != 0))))
goto cleanup;
}
}
@ -928,15 +913,14 @@ fort_driver_unload (PDRIVER_OBJECT driver)
if (g_device != NULL) {
fort_timer_close(&g_device->timer);
fort_stat_close(&g_device->stat);
fort_buffer_close(&g_device->buffer);
if (!g_device->prov_boot) {
fort_prov_unregister(0);
}
fort_callout_remove();
fort_stat_close(&g_device->stat);
fort_buffer_close(&g_device->buffer);
}
RtlInitUnicodeString(&device_link, DOS_DEVICE_NAME);

147
src/driver/fortstat.c
View File

@ -58,7 +58,10 @@ typedef struct fort_stat_flow {
} FORT_STAT_FLOW, *PFORT_STAT_FLOW;
typedef struct fort_stat {
UCHAR is_dirty;
UCHAR volatile closing;
UCHAR is_dirty : 1;
UCHAR log_stat : 1;
UINT16 limit_bits;
@ -69,6 +72,11 @@ typedef struct fort_stat {
UINT16 proc_head_index;
UINT16 proc_free_index;
UINT32 stream4_id;
UINT32 datagram4_id;
UINT32 in_transport4_id;
UINT32 out_transport4_id;
PFORT_STAT_PROC procs;
PFORT_STAT_FLOW flow_free;
@ -238,6 +246,32 @@ fort_stat_proc_add (PFORT_STAT stat, UINT32 process_id)
return proc_index;
}
static void
fort_stat_flow_context_set (PFORT_STAT stat, UINT64 flow_id)
{
FwpsFlowAssociateContext0(flow_id, FWPS_LAYER_STREAM_V4, stat->stream4_id, flow_id);
FwpsFlowAssociateContext0(flow_id, FWPS_LAYER_DATAGRAM_DATA_V4, stat->datagram4_id, flow_id);
FwpsFlowAssociateContext0(flow_id, FWPS_LAYER_INBOUND_TRANSPORT_V4, stat->in_transport4_id, flow_id);
FwpsFlowAssociateContext0(flow_id, FWPS_LAYER_OUTBOUND_TRANSPORT_V4, stat->out_transport4_id, flow_id);
}
static void
fort_stat_flow_context_remove (PFORT_STAT stat, PFORT_STAT_FLOW flow)
{
const UINT64 flow_id = flow->flow_id;
FwpsFlowRemoveContext0(flow_id, FWPS_LAYER_STREAM_V4, stat->stream4_id);
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);
}
static void
fort_stat_flow_close (PFORT_STAT_FLOW flow)
{
flow->opt.proc_index = FORT_PROC_BAD_INDEX;
}
static PFORT_STAT_FLOW
fort_stat_flow_get (PFORT_STAT stat, UINT64 flow_id, tommy_key_t flow_hash)
{
@ -256,7 +290,11 @@ fort_stat_flow_get (PFORT_STAT stat, UINT64 flow_id, tommy_key_t flow_hash)
static void
fort_stat_flow_free (PFORT_STAT stat, PFORT_STAT_FLOW flow)
{
fort_stat_proc_dec(stat, flow->opt.proc_index);
const UINT16 proc_index = flow->opt.proc_index;
if (proc_index != FORT_PROC_BAD_INDEX) {
fort_stat_proc_dec(stat, proc_index);
}
tommy_hashdyn_remove_existing(&stat->flows_map, (tommy_hashdyn_node *) flow);
@ -265,13 +303,14 @@ fort_stat_flow_free (PFORT_STAT stat, PFORT_STAT_FLOW flow)
stat->flow_free = flow;
}
static PFORT_STAT_FLOW
static NTSTATUS
fort_stat_flow_add (PFORT_STAT stat, UINT64 flow_id,
UCHAR group_index, UINT16 proc_index,
BOOL speed_limit)
{
const tommy_key_t flow_hash = fort_stat_flow_hash(flow_id);
PFORT_STAT_FLOW flow = fort_stat_flow_get(stat, flow_id, flow_hash);
BOOL is_new_flow = FALSE;
if (flow == NULL) {
if (stat->flow_free != NULL) {
@ -282,45 +321,35 @@ fort_stat_flow_add (PFORT_STAT stat, UINT64 flow_id,
// TODO: tommy_arrayof_grow(): check calloc()'s result for NULL
if (tommy_arrayof_grow(&stat->flows, size + 1), 0)
return NULL;
return STATUS_INSUFFICIENT_RESOURCES;
flow = tommy_arrayof_ref(&stat->flows, size);
}
fort_stat_proc_inc(stat, proc_index);
tommy_hashdyn_insert(&stat->flows_map, (tommy_hashdyn_node *) flow, 0, flow_hash);
flow->flow_id = flow_id;
fort_stat_flow_context_set(stat, flow_id);
is_new_flow = TRUE;
} else {
is_new_flow = (flow->opt.proc_index == FORT_PROC_BAD_INDEX);
}
if (is_new_flow) {
fort_stat_proc_inc(stat, proc_index);
}
flow->opt.speed_limit = (UCHAR) speed_limit;
flow->opt.group_index = group_index;
flow->opt.proc_index = proc_index;
return flow;
return STATUS_SUCCESS;
}
static void
fort_stat_init (PFORT_STAT stat)
{
stat->proc_head_index = FORT_PROC_BAD_INDEX;
stat->proc_free_index = FORT_PROC_BAD_INDEX;
tommy_arrayof_init(&stat->flows, sizeof(FORT_STAT_FLOW));
tommy_hashdyn_init(&stat->flows_map);
}
static void
fort_stat_open (PFORT_STAT stat)
{
fort_stat_init(stat);
KeInitializeSpinLock(&stat->lock);
}
static void
fort_stat_close (PFORT_STAT stat)
{
stat->is_dirty = FALSE;
@ -333,20 +362,48 @@ fort_stat_close (PFORT_STAT stat)
stat->proc_end = 0;
}
stat->flow_free = NULL;
stat->proc_head_index = FORT_PROC_BAD_INDEX;
stat->proc_free_index = FORT_PROC_BAD_INDEX;
}
static void
fort_stat_open (PFORT_STAT stat)
{
fort_stat_init(stat);
tommy_arrayof_init(&stat->flows, sizeof(FORT_STAT_FLOW));
tommy_hashdyn_init(&stat->flows_map);
KeInitializeSpinLock(&stat->lock);
}
static void
fort_stat_close (PFORT_STAT stat)
{
stat->closing = TRUE;
tommy_hashdyn_foreach_node_arg(&stat->flows_map,
fort_stat_flow_context_remove, stat);
tommy_arrayof_done(&stat->flows);
tommy_hashdyn_done(&stat->flows_map);
}
static void
fort_stat_clear (PFORT_STAT stat)
fort_stat_update (PFORT_STAT stat, BOOL log_stat)
{
KLOCK_QUEUE_HANDLE lock_queue;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
fort_stat_close(stat);
fort_stat_init(stat);
if (stat->log_stat) {
tommy_hashdyn_foreach_node(&stat->flows_map, fort_stat_flow_close);
fort_stat_init(stat);
}
stat->log_stat = (UCHAR) log_stat;
KeReleaseInStackQueuedSpinLock(&lock_queue);
}
@ -374,13 +431,17 @@ fort_stat_flow_associate (PFORT_STAT stat, UINT64 flow_id,
BOOL *is_new_proc)
{
KLOCK_QUEUE_HANDLE lock_queue;
PFORT_STAT_FLOW flow;
UINT16 proc_index;
BOOL speed_limit;
NTSTATUS status;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
if (!stat->log_stat) {
status = STATUS_RESOURCE_REQUIREMENTS_CHANGED;
goto end;
}
proc_index = fort_stat_proc_index(stat, process_id);
if (proc_index == FORT_PROC_BAD_INDEX) {
@ -396,19 +457,13 @@ fort_stat_flow_associate (PFORT_STAT stat, UINT64 flow_id,
speed_limit = fort_stat_group_speed_limit(stat, group_index) != 0;
flow = fort_stat_flow_add(stat, flow_id,
status = fort_stat_flow_add(stat, flow_id,
group_index, proc_index, speed_limit);
if (flow == NULL) {
if (*is_new_proc) {
PFORT_STAT_PROC proc = &stat->procs[proc_index];
if (!NT_SUCCESS(status) && *is_new_proc) {
PFORT_STAT_PROC proc = &stat->procs[proc_index];
fort_stat_proc_free(stat, proc, proc_index, NULL);
*is_new_proc = FALSE;
}
status = STATUS_INSUFFICIENT_RESOURCES;
fort_stat_proc_free(stat, proc, proc_index, NULL);
}
end:
@ -424,6 +479,9 @@ fort_stat_flow_delete (PFORT_STAT stat, UINT64 flow_id)
PFORT_STAT_FLOW flow;
const tommy_key_t flow_hash = fort_stat_flow_hash(flow_id);
if (stat->closing)
return;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
flow = fort_stat_flow_get(stat, flow_id, flow_hash);
if (flow != NULL) {
@ -442,8 +500,13 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flow_id,
BOOL limited = FALSE;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
if (!stat->log_stat)
goto end;
flow = fort_stat_flow_get(stat, flow_id, flow_hash);
if (flow != NULL) {
if (flow != NULL && flow->opt.proc_index != FORT_PROC_BAD_INDEX) {
PFORT_STAT_PROC proc = &stat->procs[flow->opt.proc_index];
UINT32 *proc_bytes = inbound ? &proc->traf.in_bytes
: &proc->traf.out_bytes;
@ -473,6 +536,8 @@ fort_stat_flow_classify (PFORT_STAT stat, UINT64 flow_id,
stat->is_dirty = TRUE;
}
end:
KeReleaseInStackQueuedSpinLock(&lock_queue);
return limited;