Driver: PsTree: Simplify ps_name handling

Process can't dye during classifying.
This commit is contained in:
Nodir Temirkhodjaev 2022-02-06 12:01:14 +03:00
parent 8042e5c0bf
commit 2b063349da
3 changed files with 12 additions and 28 deletions

View File

@ -166,8 +166,7 @@ static void fort_callout_classify_v4_check(const FWPS_INCOMING_VALUES0 *inFixedV
const UINT32 process_id = (UINT32) inMetaValues->processId;
UNICODE_STRING path;
PFORT_PSNAME ps_name =
fort_pstree_acquire_proc_name(&fort_device()->ps_tree, process_id, &path);
PFORT_PSNAME ps_name = fort_pstree_get_proc_name(&fort_device()->ps_tree, process_id, &path);
if (ps_name == NULL) {
const UINT16 path_len = (UINT16) (inMetaValues->processPath->size
- sizeof(WCHAR)); /* chop terminating zero */
@ -207,10 +206,6 @@ static void fort_callout_classify_v4_check(const FWPS_INCOMING_VALUES0 *inFixedV
fort_callout_classify_permit(filter, classifyOut);
}
}
if (ps_name != NULL) {
fort_pstree_release_proc_name(&fort_device()->ps_tree, ps_name);
}
}
static void fort_callout_classify_v4(const FWPS_INCOMING_VALUES0 *inFixedValues,

View File

@ -17,8 +17,7 @@
struct fort_psname
{
UINT8 refcount;
UINT8 size;
UINT16 size;
WCHAR data[1];
};
@ -34,6 +33,9 @@ typedef struct fort_psnode
UINT32 process_id;
UINT32 parent_process_id;
UINT16 flags;
UINT16 conf_chn;
} FORT_PSNODE, *PFORT_PSNODE;
typedef struct _SYSTEM_PROCESSES
@ -117,8 +119,7 @@ static PFORT_PSNAME fort_pstree_add_name(
FORT_PSNAME_DATA_OFF + size + sizeof(WCHAR)); /* include terminating zero */
if (ps_name != NULL) {
ps_name->refcount = 1;
ps_name->size = (UINT8) size;
ps_name->size = size;
PCHAR data = (PCHAR) &ps_name->data;
RtlCopyMemory(data, svchostPrefix.Buffer, svchostPrefix.Length);
@ -137,7 +138,7 @@ static PFORT_PSNAME fort_pstree_add_name(
static void fort_pstree_del_name(PFORT_PSTREE ps_tree, PFORT_PSNAME ps_name)
{
if (ps_name != NULL && --ps_name->refcount == 0) {
if (ps_name != NULL) {
fort_pool_free(&ps_tree->pool_list, ps_name);
}
}
@ -207,8 +208,6 @@ static void fort_pstree_handle_new_proc(PFORT_PSTREE ps_tree, PCUNICODE_STRING p
PCUNICODE_STRING commandLine, tommy_key_t pid_hash, DWORD processId, DWORD parentProcessId)
{
PFORT_PSNAME ps_name = fort_pstree_add_name(ps_tree, path, commandLine);
if (ps_name == NULL)
return;
PFORT_PSNODE proc = fort_pstree_proc_new(ps_tree, ps_name, pid_hash);
if (proc == NULL) {
@ -220,6 +219,9 @@ static void fort_pstree_handle_new_proc(PFORT_PSTREE ps_tree, PCUNICODE_STRING p
proc->process_id = processId;
proc->parent_process_id = parentProcessId;
proc->flags = 0;
proc->conf_chn = 0;
}
static void NTAPI fort_pstree_notify(
@ -312,7 +314,7 @@ FORT_API void fort_pstree_close(PFORT_PSTREE ps_tree)
KeReleaseInStackQueuedSpinLock(&lock_queue);
}
FORT_API PFORT_PSNAME fort_pstree_acquire_proc_name(
FORT_API PFORT_PSNAME fort_pstree_get_proc_name(
PFORT_PSTREE ps_tree, DWORD processId, PUNICODE_STRING path)
{
PFORT_PSNAME ps_name = NULL;
@ -323,7 +325,6 @@ FORT_API PFORT_PSNAME fort_pstree_acquire_proc_name(
PFORT_PSNODE proc = fort_pstree_find_proc(ps_tree, processId);
if (proc != NULL && proc->ps_name != NULL) {
ps_name = proc->ps_name;
++ps_name->refcount;
path->Length = ps_name->size;
path->MaximumLength = ps_name->size;
@ -334,13 +335,3 @@ FORT_API PFORT_PSNAME fort_pstree_acquire_proc_name(
return ps_name;
}
FORT_API void fort_pstree_release_proc_name(PFORT_PSTREE ps_tree, PFORT_PSNAME ps_name)
{
KLOCK_QUEUE_HANDLE lock_queue;
KeAcquireInStackQueuedSpinLock(&ps_tree->lock, &lock_queue);
{
fort_pstree_del_name(ps_tree, ps_name);
}
KeReleaseInStackQueuedSpinLock(&lock_queue);
}

View File

@ -32,11 +32,9 @@ FORT_API void fort_pstree_open(PFORT_PSTREE ps_tree);
FORT_API void fort_pstree_close(PFORT_PSTREE ps_tree);
FORT_API PFORT_PSNAME fort_pstree_acquire_proc_name(
FORT_API PFORT_PSNAME fort_pstree_get_proc_name(
PFORT_PSTREE ps_tree, DWORD processId, PUNICODE_STRING path);
FORT_API void fort_pstree_release_proc_name(PFORT_PSTREE ps_tree, PFORT_PSNAME ps_name);
#ifdef __cplusplus
} // extern "C"
#endif