From a3697289d0e01f7e65a8e56fb367e733b2da7396 Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Fri, 14 Apr 2023 12:08:56 +0300
Subject: [PATCH] Driver: fortstat: Don't remove process from map when
 iterating it

---
 src/driver/fortstat.c | 21 +++++++++++++++++----
 src/driver/fortstat.h | 10 +++++-----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/driver/fortstat.c b/src/driver/fortstat.c
index 6580b4f0..83038620 100644
--- a/src/driver/fortstat.c
+++ b/src/driver/fortstat.c
@@ -61,13 +61,18 @@ static PFORT_STAT_PROC fort_stat_proc_get(PFORT_STAT stat, UINT32 process_id, to
 
 static void fort_stat_proc_free(PFORT_STAT stat, PFORT_STAT_PROC proc)
 {
-    tommy_hashdyn_remove_existing(&stat->procs_map, (tommy_hashdyn_node *) proc);
-
     /* Add to free chain */
     proc->next = stat->proc_free;
     stat->proc_free = proc;
 }
 
+static void fort_stat_proc_del(PFORT_STAT stat, PFORT_STAT_PROC proc)
+{
+    tommy_hashdyn_remove_existing(&stat->procs_map, (tommy_hashdyn_node *) proc);
+
+    fort_stat_proc_free(stat, proc);
+}
+
 static PFORT_STAT_PROC fort_stat_proc_add(PFORT_STAT stat, UINT32 process_id)
 {
     const tommy_key_t proc_hash = fort_stat_proc_hash(process_id);
@@ -445,9 +450,17 @@ static void fort_stat_clear(PFORT_STAT stat)
     KLOCK_QUEUE_HANDLE lock_queue;
     KeAcquireInStackQueuedSpinLock(&stat->lock, &lock_queue);
 
+    /* Clear processes active list */
     fort_stat_proc_active_clear(stat);
 
-    tommy_hashdyn_foreach_node_arg(&stat->procs_map, &fort_stat_proc_free, stat);
+    /* Clear processes map */
+    if (tommy_hashdyn_count(&stat->procs_map) > 0) {
+        tommy_hashdyn_foreach_node_arg(&stat->procs_map, &fort_stat_proc_free, stat);
+        tommy_hashdyn_done(&stat->procs_map);
+        tommy_hashdyn_init(&stat->procs_map);
+    }
+
+    /* Close flows */
     tommy_hashdyn_foreach_node(&stat->flows_map, &fort_flow_close);
 
     KeReleaseInStackQueuedSpinLock(&lock_queue);
@@ -609,7 +622,7 @@ FORT_API void fort_stat_dpc_traf_flush(PFORT_STAT stat, UINT16 proc_count, PCHAR
             /* The process is inactive */
             *out_proc |= 1;
 
-            fort_stat_proc_free(stat, proc);
+            fort_stat_proc_del(stat, proc);
         } else {
             proc->active = FALSE;
 
diff --git a/src/driver/fortstat.h b/src/driver/fortstat.h
index 2573135d..ffd1a1bb 100644
--- a/src/driver/fortstat.h
+++ b/src/driver/fortstat.h
@@ -25,17 +25,17 @@ typedef struct fort_stat_proc
 
     tommy_key_t proc_hash; /* tommy_hashdyn_node::index */
 
-    UINT16 proc_index : 15; /* Synchronize with FORT_PROC_COUNT_MAX! */
-    UINT16 active : 1;
-
-    UINT16 refcount;
-
 #if defined(_WIN64)
     UINT32 process_id;
 #else
     FORT_TRAF traf;
 #endif
 
+    UINT16 proc_index : 15; /* Synchronize with FORT_PROC_COUNT_MAX! */
+    UINT16 active : 1;
+
+    UINT16 refcount;
+
     struct fort_stat_proc *next_active;
 } FORT_STAT_PROC, *PFORT_STAT_PROC;