From 8e2c84ed2030dbcd8fde013b73c3737c101822c6 Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Sat, 20 May 2023 13:23:29 +0300
Subject: [PATCH] Driver: Don't close existing connections on installation

---
 src/driver/fortcout.c | 16 +++++++++-------
 src/driver/fortstat.c | 31 +++++++++++++++++++------------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/driver/fortcout.c b/src/driver/fortcout.c
index ee0094e1..28638ebc 100644
--- a/src/driver/fortcout.c
+++ b/src/driver/fortcout.c
@@ -84,15 +84,17 @@ inline static BOOL fort_callout_ale_associate_flow(PCFORT_CALLOUT_ARG ca,
             group_index, ca->isIPv6, is_tcp, ca->inbound, cx->is_reauth, &log_stat);
 
     if (!NT_SUCCESS(status)) {
-        if (status == FORT_STATUS_FLOW_BLOCK) {
-            cx->blocked = TRUE; /* block (Reauth) */
-            cx->block_reason = FORT_BLOCK_REASON_REAUTH;
-            return TRUE;
+        if (status != FORT_STATUS_FLOW_BLOCK) {
+            LOG("Classify v4: Flow assoc. error: %x\n", status);
+            TRACE(FORT_CALLOUT_FLOW_ASSOC_ERROR, status, 0, 0);
         }
 
-        LOG("Classify v4: Flow assoc. error: %x\n", status);
-        TRACE(FORT_CALLOUT_FLOW_ASSOC_ERROR, status, 0, 0);
-    } else if (!log_stat) {
+        cx->blocked = TRUE; /* block (Error) */
+        cx->block_reason = FORT_BLOCK_REASON_REAUTH;
+        return TRUE;
+    }
+
+    if (!log_stat) {
         fort_buffer_proc_new_write(&fort_device()->buffer, cx->process_id, cx->real_path->Length,
                 cx->real_path->Buffer, &cx->irp, &cx->info);
     }
diff --git a/src/driver/fortstat.c b/src/driver/fortstat.c
index 64db8f58..aa1c1687 100644
--- a/src/driver/fortstat.c
+++ b/src/driver/fortstat.c
@@ -169,7 +169,7 @@ static void fort_flow_context_stream_init(
     }
 }
 
-inline static void fort_flow_context_stream_set(
+inline static NTSTATUS fort_flow_context_stream_set(
         PFORT_STAT stat, UINT64 flow_id, UINT64 flowContext, BOOL isIPv6, BOOL is_tcp)
 {
     UINT16 layerId;
@@ -177,7 +177,7 @@ inline static void fort_flow_context_stream_set(
 
     fort_flow_context_stream_init(stat, isIPv6, is_tcp, &layerId, &calloutId);
 
-    FwpsFlowAssociateContext0(flow_id, layerId, calloutId, flowContext);
+    return FwpsFlowAssociateContext0(flow_id, layerId, calloutId, flowContext);
 }
 
 inline static void fort_flow_context_transport_set(
@@ -196,14 +196,20 @@ inline static void fort_flow_context_transport_set(
     }
 }
 
-static void fort_flow_context_set(
+static NTSTATUS fort_flow_context_set(
         PFORT_STAT stat, PFORT_FLOW flow, BOOL isIPv6, BOOL is_tcp, BOOL inbound)
 {
     const UINT64 flow_id = flow->flow_id;
     const UINT64 flowContext = (UINT64) flow;
 
-    fort_flow_context_stream_set(stat, flow_id, flowContext, isIPv6, is_tcp);
+    const NTSTATUS status =
+            fort_flow_context_stream_set(stat, flow_id, flowContext, isIPv6, is_tcp);
+    if (!NT_SUCCESS(status))
+        return status;
+
     fort_flow_context_transport_set(stat, flow_id, flowContext, isIPv6, inbound);
+
+    return STATUS_SUCCESS;
 }
 
 inline static void fort_flow_context_stream_remove(
@@ -324,8 +330,6 @@ static PFORT_FLOW fort_flow_new(PFORT_STAT stat, UINT64 flow_id, const tommy_key
 
     flow->flow_id = flow_id;
 
-    fort_flow_context_set(stat, flow, isIPv6, is_tcp, inbound);
-
     return flow;
 }
 
@@ -340,16 +344,19 @@ inline static UCHAR fort_stat_group_speed_limit(PFORT_CONF_GROUP conf_group, UCH
 inline static NTSTATUS fort_flow_add_new(PFORT_STAT stat, PFORT_FLOW *flow, UINT64 flow_id,
         tommy_key_t flow_hash, BOOL isIPv6, BOOL is_tcp, BOOL inbound, BOOL is_reauth)
 {
-    if (is_reauth) {
-        /* Can't remove existing context, because of possible deadlock */
-        return FORT_STATUS_FLOW_BLOCK;
-    }
-
     *flow = fort_flow_new(stat, flow_id, flow_hash, isIPv6, is_tcp, inbound);
     if (*flow == NULL)
         return STATUS_INSUFFICIENT_RESOURCES;
 
-    return STATUS_SUCCESS;
+    NTSTATUS status = fort_flow_context_set(stat, *flow, isIPv6, is_tcp, inbound);
+    if (!NT_SUCCESS(status)) {
+        fort_flow_free(stat, *flow);
+
+        /* Can't remove existing context, because of possible deadlock */
+        status = is_reauth ? FORT_STATUS_FLOW_BLOCK : status;
+    }
+
+    return status;
 }
 
 static NTSTATUS fort_flow_add(PFORT_STAT stat, UINT64 flow_id, UCHAR group_index, UINT16 proc_index,