From 3b53e542da6fc09ffed0499865eb39301e377610 Mon Sep 17 00:00:00 2001
From: 87owo <85057800+87owo@users.noreply.github.com>
Date: Tue, 7 Jul 2026 02:31:38 +0800
Subject: [PATCH] Add files via upload

---
 Plugins/Filter/DriverEntry.cpp | 69 +++++++++++++++++++++-------------
 1 file changed, 42 insertions(+), 27 deletions(-)

diff --git a/Plugins/Filter/DriverEntry.cpp b/Plugins/Filter/DriverEntry.cpp
index 01a56e9..95c7582 100644
--- a/Plugins/Filter/DriverEntry.cpp
+++ b/Plugins/Filter/DriverEntry.cpp
@@ -105,15 +105,23 @@ static NTSTATUS UnregisterRuntimeCallbacks(BOOLEAN WaitWithoutTimeout) {
     }
 }
 
-static VOID ReleaseConnectedProcessReference() {
+static PEPROCESS DetachConnectedProcessReference(PEPROCESS ExpectedProcess) {
     PEPROCESS Process = NULL;
 
     KIRQL OldIrql;
     KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
-    Process = GlobalData.PyasProcess;
-    GlobalData.PyasProcess = NULL;
-    GlobalData.ClientPort = NULL;
+
+    if (!ExpectedProcess || GlobalData.PyasProcess == ExpectedProcess) {
+        Process = GlobalData.PyasProcess;
+        GlobalData.PyasProcess = NULL;
+    }
+
     KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
+    return Process;
+}
+
+static VOID ReleaseConnectedProcessReference() {
+    PEPROCESS Process = DetachConnectedProcessReference(NULL);
 
     SafeSetPyasPid(0);
     InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
@@ -126,13 +134,32 @@ static VOID ReleaseConnectedProcessReference() {
 static VOID CloseCurrentClientPort() {
     if (InterlockedCompareExchange(&GlobalData.ClientCloseActive, 1, 0) != 0) return;
 
-    if (GlobalData.FilterHandle && GlobalData.ClientPort) {
-        FltCloseClientPort(GlobalData.FilterHandle, &GlobalData.ClientPort);
+    PFLT_FILTER FilterHandle = NULL;
+    PFLT_PORT ClientPort = NULL;
+
+    KIRQL OldIrql;
+    KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
+    FilterHandle = GlobalData.FilterHandle;
+    ClientPort = GlobalData.ClientPort;
+    GlobalData.ClientPort = NULL;
+    KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
+
+    if (FilterHandle && ClientPort) {
+        FltCloseClientPort(FilterHandle, &ClientPort);
     }
 
     InterlockedExchange(&GlobalData.ClientCloseActive, 0);
 }
 
+static VOID WaitForClientPortClose() {
+    LARGE_INTEGER Delay;
+    Delay.QuadPart = -(10LL * 1000);
+
+    while (InterlockedCompareExchange(&GlobalData.ClientCloseActive, 0, 0) != 0) {
+        KeDelayExecutionThread(KernelMode, FALSE, &Delay);
+    }
+}
+
 static VOID ReleaseFilterResources() {
     InterlockedExchange(&GlobalData.PortStopping, 1);
 
@@ -143,6 +170,7 @@ static VOID ReleaseFilterResources() {
     }
 
     CloseCurrentClientPort();
+    WaitForClientPortClose();
     ExWaitForRundownProtectionRelease(&GlobalData.PortRundown);
 
     if (GlobalData.FilterHandle) {
@@ -473,28 +501,15 @@ static NTSTATUS PortConnect(
 }
 
 static VOID PortDisconnect(PVOID ConnectionCookie) {
+    if (!ConnectionCookie) return;
+
+    PEPROCESS Process = DetachConnectedProcessReference((PEPROCESS)ConnectionCookie);
+    if (!Process) return;
+
     CloseCurrentClientPort();
-
-    BOOLEAN DisconnectedCurrentClient = FALSE;
-
-    KIRQL OldIrql;
-    KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
-
-    if (ConnectionCookie && GlobalData.PyasProcess == (PEPROCESS)ConnectionCookie) {
-        GlobalData.PyasProcess = NULL;
-        DisconnectedCurrentClient = TRUE;
-    }
-
-    KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
-
-    if (DisconnectedCurrentClient) {
-        InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
-        SafeSetPyasPid(0);
-    }
-
-    if (ConnectionCookie) {
-        ObDereferenceObject((PEPROCESS)ConnectionCookie);
-    }
+    InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
+    SafeSetPyasPid(0);
+    ObDereferenceObject(Process);
 }
 
 CONST FLT_OPERATION_REGISTRATION Callbacks[] = {