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

Driver: Simplify fort_flow_associate().

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/driver/fortcb.c
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)

56
src/driver/fortstat.c
View File

@ -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;

2
src/driver/forttds.c
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"