Driver: Simplify fort_flow_associate().

This commit is contained in:
Nodir Temirkhodjaev 2021-05-10 12:28:50 +03:00
parent 92fc8f2d98
commit bbf41624d1
4 changed files with 38 additions and 34 deletions

View File

@ -2,8 +2,8 @@
#include "fortdrv.h"
#include "fortdev.h"
#include "fortcout.h"
#include "fortdev.h"
static NTSTATUS fort_callback_register(
PCWSTR sourcePath, PCALLBACK_OBJECT *cb_obj, PVOID *cb_reg, PCALLBACK_FUNCTION cb_func)

View File

@ -156,7 +156,7 @@ static PFORT_CONF_EXE_NODE fort_conf_ref_exe_find_node(
FORT_API FORT_APP_FLAGS fort_conf_exe_find(const PFORT_CONF conf, const char *path, UINT32 path_len)
{
PFORT_CONF_REF conf_ref = (PFORT_CONF_REF)((char *) conf - offsetof(FORT_CONF_REF, conf));
PFORT_CONF_REF conf_ref = (PFORT_CONF_REF) ((char *) conf - offsetof(FORT_CONF_REF, conf));
const tommy_key_t path_hash = (tommy_key_t) tommy_hash_u64(0, path, path_len);
FORT_APP_FLAGS app_flags;
@ -542,7 +542,7 @@ FORT_API BOOL fort_conf_zones_ip_included(
while (zones_mask != 0) {
const int zone_index = bit_scan_forward(zones_mask);
PFORT_CONF_ADDR_LIST addr_list =
(PFORT_CONF_ADDR_LIST)(zones->data + zones->addr_off[zone_index]);
(PFORT_CONF_ADDR_LIST) (zones->data + zones->addr_off[zone_index]);
if (fort_conf_ip_inlist(remote_ip, addr_list)) {
res = TRUE;

View File

@ -9,8 +9,8 @@
#define FORT_PROC_BAD_INDEX ((UINT16) -1)
#define FORT_PROC_COUNT_MAX 0x7FFF
#define fort_stat_proc_hash(process_id) tommy_inthash_u32((UINT32)(process_id))
#define fort_flow_hash(flow_id) tommy_inthash_u32((UINT32)(flow_id))
#define fort_stat_proc_hash(process_id) tommy_inthash_u32((UINT32) (process_id))
#define fort_flow_hash(flow_id) tommy_inthash_u32((UINT32) (flow_id))
#define fort_stat_group_fragment(stat, group_index) \
((((stat)->conf_group.fragment_bits >> (group_index)) & 1) != 0 ? FORT_FLOW_FRAGMENT : 0)
@ -306,41 +306,46 @@ FORT_API void fort_stat_conf_update(PFORT_STAT stat, PFORT_CONF_IO conf_io)
KeReleaseInStackQueuedSpinLock(&lock_queue);
}
static NTSTATUS fort_flow_associate_proc(PFORT_STAT stat, UINT32 process_id, BOOL is_reauth,
BOOL *is_new_proc, PFORT_STAT_PROC *proc)
{
if (!stat->log_stat)
return STATUS_DEVICE_DATA_ERROR;
const tommy_key_t proc_hash = fort_stat_proc_hash(process_id);
*proc = fort_stat_proc_get(stat, process_id, proc_hash);
if (*proc == NULL) {
if (is_reauth) {
/* Block existing flow after reauth. to be able to associate a flow-context */
return FORT_STATUS_FLOW_BLOCK;
}
*proc = fort_stat_proc_add(stat, process_id);
if (*proc == NULL)
return STATUS_INSUFFICIENT_RESOURCES;
*is_new_proc = TRUE;
}
return STATUS_SUCCESS;
}
FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, UINT32 process_id,
UCHAR group_index, BOOL is_tcp, BOOL is_reauth, BOOL *is_new_proc)
{
const tommy_key_t proc_hash = fort_stat_proc_hash(process_id);
KLOCK_QUEUE_HANDLE lock_queue;
NTSTATUS status;
KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
if (!stat->log_stat) {
status = STATUS_DEVICE_DATA_ERROR;
goto end;
}
PFORT_STAT_PROC proc = fort_stat_proc_get(stat, process_id, proc_hash);
if (proc == NULL) {
if (is_reauth) {
/* Block existing flow after reauth. to be able to associate a flow-context */
status = FORT_STATUS_FLOW_BLOCK;
goto end;
}
proc = fort_stat_proc_add(stat, process_id);
if (proc == NULL) {
status = STATUS_INSUFFICIENT_RESOURCES;
goto end;
}
*is_new_proc = TRUE;
}
PFORT_STAT_PROC proc = NULL;
status = fort_flow_associate_proc(stat, process_id, is_reauth, is_new_proc, &proc);
/* Add flow */
{
if (NT_SUCCESS(status)) {
const UCHAR fragment = fort_stat_group_fragment(stat, group_index);
const UCHAR speed_limit = fort_stat_group_speed_limit(stat, group_index);
@ -352,7 +357,6 @@ FORT_API NTSTATUS fort_flow_associate(PFORT_STAT stat, UINT64 flow_id, UINT32 pr
}
}
end:
KeReleaseInStackQueuedSpinLock(&lock_queue);
return status;
@ -453,9 +457,9 @@ FORT_API void fort_stat_dpc_traf_flush(PFORT_STAT stat, UINT16 proc_count, PCHAR
while (proc != NULL && proc_count-- != 0) {
PFORT_STAT_PROC proc_next = proc->next_active;
UINT32 *out_proc = (UINT32 *) out;
PFORT_TRAF out_traf = (PFORT_TRAF)(out_proc + 1);
PFORT_TRAF out_traf = (PFORT_TRAF) (out_proc + 1);
out = (PCHAR)(out_traf + 1);
out = (PCHAR) (out_traf + 1);
/* Write bytes */
*out_traf = proc->traf;

View File

@ -36,6 +36,6 @@ FORT_API PVOID fort_tommy_realloc(PVOID p, SIZE_T new_size)
}
#include "../3rdparty/tommyds/tommyarrayof.c"
#include "../3rdparty/tommyds/tommylist.c"
#include "../3rdparty/tommyds/tommyhash.c"
#include "../3rdparty/tommyds/tommyhashdyn.c"
#include "../3rdparty/tommyds/tommylist.c"