Add files via upload

Plugins/Filter/DriverCommon.h

@@ -8,6 +8,7 @@
 constexpr auto PYAS_PORT_NAME = L"\\PYAS_Output_Pipe";
 constexpr auto MAX_PATH_LEN = 1024;
 constexpr auto PYAS_POOL_TAG = 'SAYP';
+constexpr auto PYAS_OB_ALTITUDE = L"320000.4101";
 
 #ifndef PROCESS_TERMINATE
 #define PROCESS_TERMINATE                  (0x0001)
@@ -60,6 +61,12 @@ constexpr auto PYAS_POOL_TAG = 'SAYP';
 #ifndef THREAD_SET_THREAD_TOKEN
 #define THREAD_SET_THREAD_TOKEN            (0x0080)
 #endif
+#ifndef THREAD_IMPERSONATE
+#define THREAD_IMPERSONATE                 (0x0100)
+#endif
+#ifndef THREAD_DIRECT_IMPERSONATION
+#define THREAD_DIRECT_IMPERSONATION        (0x0200)
+#endif
 #ifndef POOL_FLAG_NON_PAGED
 #define POOL_FLAG_NON_PAGED                0x0000000000000040UI64
 #endif
@@ -94,6 +101,16 @@ typedef struct _FS_BPIO_OUTPUT {
 } FS_BPIO_OUTPUT, * PFS_BPIO_OUTPUT;
 #endif
 
+
+typedef enum _PYAS_DRIVER_STATE {
+    PyasDriverStateCold = 0,
+    PyasDriverStateStarting = 1,
+    PyasDriverStateRunning = 2,
+    PyasDriverStateStopping = 3,
+    PyasDriverStateStopRetry = 4,
+    PyasDriverStateStopped = 5
+} PYAS_DRIVER_STATE;
+
 typedef struct _PYAS_MESSAGE {
     ULONG MessageCode;
     ULONG ProcessId;
@@ -106,7 +123,8 @@ typedef enum _PYAS_COMMAND {
     PyasCommandLoadRuleFile = 3,
     PyasCommandClearRules = 4,
     PyasCommandAuthorizeUnload = 5,
-    PyasCommandRevokeUnload = 6
+    PyasCommandRevokeUnload = 6,
+    PyasCommandQueryState = 7
 } PYAS_COMMAND;
 
 typedef struct _PYAS_USER_MESSAGE {
@@ -126,8 +144,12 @@ typedef struct _DRIVER_DATA {
     KSPIN_LOCK TrackerMutex;
     KSPIN_LOCK PidLock;
     BOOLEAN Initialized;
+    volatile LONG DriverState;
     volatile LONG UnloadAuthorized;
+    volatile LONG PortStopping;
+    volatile LONG ClientCloseActive;
     EX_RUNDOWN_REF PortRundown;
+    KEVENT CleanupCompleteEvent;
     PVOID ObRegistrationHandle;
 } DRIVER_DATA, * PDRIVER_DATA;
 
@@ -146,40 +168,97 @@ typedef enum _RULE_CATEGORY {
     RuleCategoryUnknown
 } RULE_CATEGORY;
 
-#define OP_WRITE         0x01
-#define OP_DELETE        0x02
-#define OP_CREATE        0x04
-#define OP_EXECUTE       0x08
-#define OP_RENAME        0x10
-#define OP_IOCTL         0x20
-#define OP_VM_READ           0x00000040
-#define OP_VM_WRITE          0x00000080
-#define OP_TERMINATE         0x00000100
-#define OP_SUSPEND_RESUME    0x00000200
-#define OP_DUP_HANDLE        0x00000400
-#define OP_SET_INFORMATION   0x00000800
+#define OP_WRITE                   0x00000001
+#define OP_DELETE                  0x00000002
+#define OP_CREATE                  0x00000004
+#define OP_EXECUTE                 0x00000008
+#define OP_RENAME                  0x00000010
+#define OP_IOCTL                   0x00000020
+#define OP_VM_READ                 0x00000040
+#define OP_VM_WRITE                0x00000080
+#define OP_TERMINATE               0x00000100
+#define OP_SUSPEND_RESUME          0x00000200
+#define OP_DUP_HANDLE              0x00000400
+#define OP_SET_INFORMATION         0x00000800
+#define OP_VM_OPERATION            0x00001000
+#define OP_CREATE_THREAD           0x00002000
+#define OP_THREAD_SET_CONTEXT      0x00004000
+#define OP_THREAD_SET_TOKEN        0x00008000
+#define OP_CREATE_PROCESS          0x00010000
+#define OP_IMAGE_LOAD              0x00020000
+#define OP_IMPERSONATE             0x00040000
 
-BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode);
+#define PYAS_HANDLE_CREATE         0x00000001
+#define PYAS_HANDLE_DUPLICATE      0x00000002
+
+#define PYAS_OBJECT_PROCESS        0x00000001
+#define PYAS_OBJECT_THREAD         0x00000002
+
+#define PYAS_MEMORY_PRIVATE        0x00000001
+#define PYAS_MEMORY_MAPPED         0x00000002
+#define PYAS_MEMORY_IMAGE          0x00000004
+
+#define PYAS_PROTECT_EXECUTE       0x00000001
+#define PYAS_PROTECT_EXECUTE_WRITE 0x00000002
+
+typedef enum _RULE_TRI_STATE {
+    RuleTriAny = 0,
+    RuleTriTrue = 1,
+    RuleTriFalse = 2
+} RULE_TRI_STATE;
+
+typedef enum _RULE_OPERATION_MATCH {
+    RuleOperationAny = 0,
+    RuleOperationAll = 1
+} RULE_OPERATION_MATCH;
 
 typedef struct _DYNAMIC_RULE {
     ULONG Code;
+    BOOLEAN Kill;
+    ULONG Priority;
     RULE_CATEGORY Category;
     ULONG Operations;
+    RULE_OPERATION_MATCH OperationMatch;
+    BOOLEAN Invalid;
 
     PRULE_NODE Initiator;
     PRULE_NODE InitiatorExclude;
+    PRULE_NODE InitiatorParent;
+    PRULE_NODE InitiatorParentExclude;
+    PRULE_NODE InitiatorProcessTree;
+    PRULE_NODE InitiatorProcessTreeExclude;
     PRULE_NODE Target;
     PRULE_NODE TargetExclude;
     PRULE_NODE TargetProcessTree;
+    PRULE_NODE TargetProcessTreeExclude;
+    PRULE_NODE Creator;
+    PRULE_NODE CreatorExclude;
+    PRULE_NODE Parent;
+    PRULE_NODE ParentExclude;
     PRULE_NODE CommandLine;
+    PRULE_NODE CommandLineExclude;
     PRULE_NODE Extensions;
 
+    ULONG HandleTypes;
+    ULONG ObjectTypes;
+    ULONG MinimumRiskScore;
+    ULONG MaximumRiskScore;
+    ULONG ThreadMemoryTypes;
+    ULONG ThreadMemoryProtections;
+    ULONGLONG MinimumRegionSize;
+    ULONGLONG MaximumRegionSize;
+    RULE_TRI_STATE ParentMismatch;
+    RULE_TRI_STATE FileOpenNameAvailable;
+    RULE_TRI_STATE SubsystemProcess;
+
     ULONG Threshold;
     ULONG TimeWindow;
 
     struct _DYNAMIC_RULE* Next;
 } DYNAMIC_RULE, * PDYNAMIC_RULE;
 
+BOOLEAN EvaluateDeviceRule(HANDLE ProcessId, PULONG OutCode, PBOOLEAN OutKill);
+
 extern DRIVER_DATA GlobalData;
 
 FLT_PREOP_CALLBACK_STATUS ProtectFile_PreCreate(PFLT_CALLBACK_DATA Data, PCFLT_RELATED_OBJECTS FltObjects, PVOID* CompletionContext);
@@ -190,15 +269,19 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_SetSecurity(PFLT_CALLBACK_DATA Data, PCFLT
 FLT_PREOP_CALLBACK_STATUS ProtectBoot_PreDeviceControl(PFLT_CALLBACK_DATA Data, PCFLT_RELATED_OBJECTS FltObjects, PVOID* CompletionContext);
 
 NTSTATUS InitializeRegistryProtection(PDRIVER_OBJECT DriverObject);
-VOID UninitializeRegistryProtection();
+NTSTATUS UninitializeRegistryProtection();
+
+NTSTATUS InitializeProcessProtection();
+NTSTATUS UninitializeProcessProtection();
+BOOLEAN GetProcessRelation(HANDLE ProcessId, PHANDLE ParentProcessId, PHANDLE CreatorProcessId);
 
 NTSTATUS InitializeMemoryProtection(PDRIVER_OBJECT DriverObject);
-VOID UninitializeMemoryProtection();
+NTSTATUS UninitializeMemoryProtection(BOOLEAN WaitWithoutTimeout);
 
 NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath);
 VOID UnloadRules();
 
-VOID InitializeRulesEngine();
+NTSTATUS InitializeRulesEngine();
 VOID UninitializeRulesEngine();
 
 VOID AddDynamicWhitelist(PUNICODE_STRING RuleStr);
@@ -216,13 +299,45 @@ VOID SafeSetPyasPid(ULONG Pid);
 BOOLEAN IsProcessTrusted(HANDLE ProcessId);
 BOOLEAN WildcardMatch(PCWSTR Pattern, PCWSTR String, USHORT StringLengthBytes);
 
-BOOLEAN EvaluateProcessRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, PCUNICODE_STRING CommandLine, PULONG OutCode);
-BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Operation, PVOID WriteBuffer, ULONG WriteLength, PULONG OutCode);
-BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG Operation, PULONG OutCode);
-BOOLEAN EvaluateMemoryRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode);
-BOOLEAN EvaluateProcessHandleRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode);
-BOOLEAN EvaluateThreadRule(HANDLE SourcePid, HANDLE TargetPid, PVOID StartAddress, PULONG OutCode);
+BOOLEAN EvaluateProcessCreateRule(
+    HANDLE CreatorPid,
+    HANDLE ParentPid,
+    HANDLE ProcessId,
+    PCUNICODE_STRING TargetPath,
+    PCUNICODE_STRING CommandLine,
+    BOOLEAN FileOpenNameAvailable,
+    BOOLEAN IsSubsystemProcess,
+    PULONG OutCode,
+    PBOOLEAN OutKill
+);
+BOOLEAN EvaluateProcessAccessRule(
+    HANDLE SourcePid,
+    HANDLE TargetPid,
+    ULONG RequestedOperations,
+    ULONG HandleType,
+    ULONG ObjectType,
+    PULONG DeniedOperations,
+    PULONG OutCode,
+    PBOOLEAN OutKill
+);
+BOOLEAN EvaluateThreadRule(
+    HANDLE SourcePid,
+    HANDLE TargetPid,
+    PVOID StartAddress,
+    ULONG MemoryType,
+    ULONG MemoryProtection,
+    SIZE_T RegionSize,
+    PULONG OutCode,
+    PBOOLEAN OutKill
+);
+BOOLEAN EvaluateImageLoadRule(HANDLE ProcessId, PCUNICODE_STRING ImagePath, PIMAGE_INFO ImageInfo, PULONG OutCode, PBOOLEAN OutKill);
+BOOLEAN EvaluateProcessRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, PCUNICODE_STRING CommandLine, PULONG OutCode, PBOOLEAN OutKill);
+BOOLEAN EvaluateFileRule(HANDLE ProcessId, PCUNICODE_STRING TargetPath, ULONG Operation, PVOID WriteBuffer, ULONG WriteLength, PULONG OutCode, PBOOLEAN OutKill);
+BOOLEAN EvaluateRegistryRule(HANDLE ProcessId, PCUNICODE_STRING KeyName, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill);
+BOOLEAN EvaluateMemoryRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill);
+BOOLEAN EvaluateProcessHandleRule(HANDLE SourcePid, HANDLE TargetPid, ULONG Operation, PULONG OutCode, PBOOLEAN OutKill);
 
+VOID QueueRuleProcessTermination(HANDLE ProcessId, BOOLEAN Kill);
 NTSTATUS SendMessageToUser(ULONG Code, ULONG Pid, PWCHAR Path, USHORT PathSize);
 
 FORCEINLINE PVOID PyasAllocate(SIZE_T Size) {

Plugins/Filter/DriverEntry.cpp

@@ -2,6 +2,7 @@
 
 DRIVER_DATA GlobalData;
 static BOOLEAN g_ProcessNotifyRegistered = FALSE;
+static BOOLEAN g_ProcessProtectionInitialized = FALSE;
 static BOOLEAN g_RegistryProtectionInitialized = FALSE;
 static BOOLEAN g_RulesEngineInitialized = FALSE;
 
@@ -43,35 +44,121 @@ static VOID InstanceTeardownComplete(
     UNREFERENCED_PARAMETER(Flags);
 }
 
-static VOID UnregisterRuntimeCallbacks() {
-    UninitializeMemoryProtection();
+static NTSTATUS SelectFailure(NTSTATUS CurrentStatus, NTSTATUS CandidateStatus) {
+    if (!NT_SUCCESS(CurrentStatus)) return CurrentStatus;
+    return CandidateStatus;
+}
+
+static NTSTATUS UnregisterRuntimeCallbacksPass(BOOLEAN WaitWithoutTimeout) {
+    NTSTATUS Status = STATUS_SUCCESS;
+
+    NTSTATUS MemoryStatus = UninitializeMemoryProtection(WaitWithoutTimeout);
+    if (!NT_SUCCESS(MemoryStatus)) {
+        Status = SelectFailure(Status, MemoryStatus);
+    }
 
     if (g_ProcessNotifyRegistered) {
-        PsSetCreateProcessNotifyRoutineEx(ProcessNotifyCallbackEx, TRUE);
-        g_ProcessNotifyRegistered = FALSE;
+        NTSTATUS ProcessNotifyStatus = PsSetCreateProcessNotifyRoutineEx(
+            ProcessNotifyCallbackEx,
+            TRUE
+        );
+
+        if (NT_SUCCESS(ProcessNotifyStatus) || ProcessNotifyStatus == STATUS_INVALID_PARAMETER) {
+            g_ProcessNotifyRegistered = FALSE;
+        }
+        else {
+            Status = SelectFailure(Status, ProcessNotifyStatus);
+        }
+    }
+
+    if (!g_ProcessNotifyRegistered && g_ProcessProtectionInitialized) {
+        NTSTATUS ProcessStatus = UninitializeProcessProtection();
+        if (NT_SUCCESS(ProcessStatus)) {
+            g_ProcessProtectionInitialized = FALSE;
+        }
+        else {
+            Status = SelectFailure(Status, ProcessStatus);
+        }
     }
 
     if (g_RegistryProtectionInitialized) {
-        UninitializeRegistryProtection();
-        g_RegistryProtectionInitialized = FALSE;
+        NTSTATUS RegistryStatus = UninitializeRegistryProtection();
+        if (NT_SUCCESS(RegistryStatus)) {
+            g_RegistryProtectionInitialized = FALSE;
+        }
+        else {
+            Status = SelectFailure(Status, RegistryStatus);
+        }
     }
+
+    return Status;
+}
+
+static NTSTATUS UnregisterRuntimeCallbacks(BOOLEAN WaitWithoutTimeout) {
+    for (;;) {
+        NTSTATUS Status = UnregisterRuntimeCallbacksPass(WaitWithoutTimeout);
+        if (NT_SUCCESS(Status) || !WaitWithoutTimeout) return Status;
+
+        LARGE_INTEGER Delay;
+        Delay.QuadPart = -(10LL * 10 * 1000);
+        KeDelayExecutionThread(KernelMode, FALSE, &Delay);
+    }
+}
+
+static VOID ReleaseConnectedProcessReference() {
+    PEPROCESS Process = NULL;
+
+    KIRQL OldIrql;
+    KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
+    Process = GlobalData.PyasProcess;
+    GlobalData.PyasProcess = NULL;
+    GlobalData.ClientPort = NULL;
+    KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
+
+    SafeSetPyasPid(0);
+    InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
+
+    if (Process) {
+        ObDereferenceObject(Process);
+    }
+}
+
+static VOID CloseCurrentClientPort() {
+    if (InterlockedCompareExchange(&GlobalData.ClientCloseActive, 1, 0) != 0) return;
+
+    if (GlobalData.FilterHandle && GlobalData.ClientPort) {
+        FltCloseClientPort(GlobalData.FilterHandle, &GlobalData.ClientPort);
+    }
+
+    InterlockedExchange(&GlobalData.ClientCloseActive, 0);
 }
 
 static VOID ReleaseFilterResources() {
+    InterlockedExchange(&GlobalData.PortStopping, 1);
+
     if (GlobalData.ServerPort) {
-        FltCloseCommunicationPort(GlobalData.ServerPort);
+        PFLT_PORT ServerPort = GlobalData.ServerPort;
         GlobalData.ServerPort = NULL;
+        FltCloseCommunicationPort(ServerPort);
     }
 
+    CloseCurrentClientPort();
+    ExWaitForRundownProtectionRelease(&GlobalData.PortRundown);
+
     if (GlobalData.FilterHandle) {
-        FltUnregisterFilter(GlobalData.FilterHandle);
+        PFLT_FILTER FilterHandle = GlobalData.FilterHandle;
+        FltUnregisterFilter(FilterHandle);
         GlobalData.FilterHandle = NULL;
     }
+
+    ReleaseConnectedProcessReference();
 }
 
-static VOID CleanupDriverState() {
+static NTSTATUS CleanupDriverState(BOOLEAN WaitWithoutTimeout) {
+    NTSTATUS Status = UnregisterRuntimeCallbacks(WaitWithoutTimeout);
+    if (!NT_SUCCESS(Status)) return Status;
+
     GlobalData.Initialized = FALSE;
-    UnregisterRuntimeCallbacks();
     ReleaseFilterResources();
 
     if (g_RulesEngineInitialized) {
@@ -79,17 +166,94 @@ static VOID CleanupDriverState() {
         UninitializeRulesEngine();
         g_RulesEngineInitialized = FALSE;
     }
+
+    return STATUS_SUCCESS;
+}
+
+static NTSTATUS BeginUnload(BOOLEAN MandatoryUnload) {
+    for (;;) {
+        LONG State = InterlockedCompareExchange(
+            &GlobalData.DriverState,
+            PyasDriverStateCold,
+            PyasDriverStateCold
+        );
+
+        if (State == PyasDriverStateStopped) return STATUS_SUCCESS;
+
+        if (State == PyasDriverStateStopping) {
+            LARGE_INTEGER Timeout;
+            PLARGE_INTEGER TimeoutPointer = NULL;
+
+            if (!MandatoryUnload) {
+                Timeout.QuadPart = -(2LL * 10 * 1000 * 1000);
+                TimeoutPointer = &Timeout;
+            }
+
+            NTSTATUS WaitStatus = KeWaitForSingleObject(
+                &GlobalData.CleanupCompleteEvent,
+                Executive,
+                KernelMode,
+                FALSE,
+                TimeoutPointer
+            );
+
+            if (!MandatoryUnload && WaitStatus == STATUS_TIMEOUT) {
+                return STATUS_DEVICE_BUSY;
+            }
+
+            if (!NT_SUCCESS(WaitStatus)) return WaitStatus;
+            continue;
+        }
+
+        if (State != PyasDriverStateRunning && State != PyasDriverStateStopRetry) {
+            return STATUS_DEVICE_NOT_READY;
+        }
+
+        if (!MandatoryUnload && InterlockedCompareExchange(&GlobalData.UnloadAuthorized, 0, 1) != 1) {
+            return STATUS_ACCESS_DENIED;
+        }
+
+        if (MandatoryUnload) {
+            InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
+        }
+
+        if (InterlockedCompareExchange(
+            &GlobalData.DriverState,
+            PyasDriverStateStopping,
+            State
+        ) != State) {
+            if (!MandatoryUnload) return STATUS_DEVICE_BUSY;
+            continue;
+        }
+
+        KeClearEvent(&GlobalData.CleanupCompleteEvent);
+        return STATUS_SUCCESS;
+    }
 }
 
 static NTSTATUS DriverUnload(FLT_FILTER_UNLOAD_FLAGS Flags) {
-    UNREFERENCED_PARAMETER(Flags);
+    BOOLEAN MandatoryUnload = FlagOn(Flags, FLTFL_FILTER_UNLOAD_MANDATORY) ? TRUE : FALSE;
 
-    if (InterlockedCompareExchange(&GlobalData.UnloadAuthorized, 0, 1) != 1) {
-        return STATUS_ACCESS_DENIED;
-    }
+    NTSTATUS Status = BeginUnload(MandatoryUnload);
+    if (!NT_SUCCESS(Status)) return Status;
 
-    CleanupDriverState();
-    return STATUS_SUCCESS;
+    LONG State = InterlockedCompareExchange(
+        &GlobalData.DriverState,
+        PyasDriverStateCold,
+        PyasDriverStateCold
+    );
+
+    if (State == PyasDriverStateStopped) return STATUS_SUCCESS;
+
+    Status = CleanupDriverState(MandatoryUnload);
+
+    InterlockedExchange(
+        &GlobalData.DriverState,
+        NT_SUCCESS(Status) ? PyasDriverStateStopped : PyasDriverStateStopRetry
+    );
+
+    KeSetEvent(&GlobalData.CleanupCompleteEvent, IO_NO_INCREMENT, FALSE);
+    return Status;
 }
 
 static NTSTATUS PortMessage(
@@ -100,52 +264,138 @@ static NTSTATUS PortMessage(
     ULONG OutputBufferLength,
     PULONG ReturnOutputBufferLength
 ) {
-    UNREFERENCED_PARAMETER(OutputBuffer);
-    UNREFERENCED_PARAMETER(OutputBufferLength);
+    if (ReturnOutputBufferLength) *ReturnOutputBufferLength = 0;
+
+    if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) {
+        return STATUS_DELETE_PENDING;
+    }
+
+    if (!ExAcquireRundownProtection(&GlobalData.PortRundown)) {
+        return STATUS_DELETE_PENDING;
+    }
+
+    NTSTATUS Status = STATUS_SUCCESS;
+    PYAS_USER_MESSAGE Message = { 0 };
+    UNICODE_STRING Path = { 0 };
+    LONG DriverState = PyasDriverStateCold;
+    ULONG StateValue = PyasDriverStateCold;
 
     if (!PortCookie || PortCookie != GlobalData.PyasProcess) {
-        return STATUS_ACCESS_DENIED;
+        Status = STATUS_ACCESS_DENIED;
+        goto Exit;
     }
 
     if (!InputBuffer || InputBufferLength < sizeof(PYAS_USER_MESSAGE)) {
-        return STATUS_INVALID_PARAMETER;
+        Status = STATUS_INVALID_PARAMETER;
+        goto Exit;
     }
 
-    PPYAS_USER_MESSAGE Message = (PPYAS_USER_MESSAGE)InputBuffer;
-    Message->Path[MAX_PATH_LEN - 1] = L'\0';
+    __try {
+        RtlCopyMemory(&Message, InputBuffer, sizeof(Message));
+    }
+    __except (EXCEPTION_EXECUTE_HANDLER) {
+        Status = GetExceptionCode();
+        goto Exit;
+    }
+
+    Message.Path[MAX_PATH_LEN - 1] = L'\0';
+
+    if (Message.Command == PyasCommandQueryState) {
+        if (!OutputBuffer || OutputBufferLength < sizeof(ULONG) || !ReturnOutputBufferLength) {
+            Status = STATUS_BUFFER_TOO_SMALL;
+            goto Exit;
+        }
+
+        StateValue = (ULONG)InterlockedCompareExchange(
+            &GlobalData.DriverState,
+            PyasDriverStateCold,
+            PyasDriverStateCold
+        );
+
+        __try {
+            *(PULONG)OutputBuffer = StateValue;
+            *ReturnOutputBufferLength = sizeof(ULONG);
+        }
+        __except (EXCEPTION_EXECUTE_HANDLER) {
+            Status = GetExceptionCode();
+        }
+        goto Exit;
+    }
+
+    if (Message.Command == PyasCommandAuthorizeUnload) {
+        DriverState = InterlockedCompareExchange(
+            &GlobalData.DriverState,
+            PyasDriverStateCold,
+            PyasDriverStateCold
+        );
+
+        if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) {
+            Status = STATUS_DEVICE_NOT_READY;
+            goto Exit;
+        }
 
-    if (Message->Command == PyasCommandAuthorizeUnload) {
         InterlockedExchange(&GlobalData.UnloadAuthorized, 1);
+        goto Exit;
     }
-    else if (Message->Command == PyasCommandRevokeUnload) {
+
+    if (Message.Command == PyasCommandRevokeUnload) {
         InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
-    }
-    else if (Message->Command >= PyasCommandAddWhitelist && Message->Command <= PyasCommandClearRules) {
-        UNICODE_STRING Path;
-        RtlInitUnicodeString(&Path, Message->Path);
-
-        if (Message->Command == PyasCommandAddWhitelist && Path.Length > 0) {
-            AddDynamicWhitelist(&Path);
-        }
-        else if (Message->Command == PyasCommandRemoveWhitelist && Path.Length > 0) {
-            RemoveDynamicWhitelist(&Path);
-        }
-        else if (Message->Command == PyasCommandLoadRuleFile && Path.Length > 0) {
-            LoadRuleFile(&Path);
-        }
-        else if (Message->Command == PyasCommandClearRules) {
-            ClearDynamicRules();
-        }
-    }
-    else {
-        return STATUS_INVALID_DEVICE_REQUEST;
+        goto Exit;
     }
 
-    if (ReturnOutputBufferLength) {
-        *ReturnOutputBufferLength = 0;
+    if (Message.Command < PyasCommandAddWhitelist || Message.Command > PyasCommandClearRules) {
+        Status = STATUS_INVALID_DEVICE_REQUEST;
+        goto Exit;
     }
 
-    return STATUS_SUCCESS;
+    DriverState = InterlockedCompareExchange(
+        &GlobalData.DriverState,
+        PyasDriverStateCold,
+        PyasDriverStateCold
+    );
+
+    if (DriverState != PyasDriverStateRunning) {
+        Status = STATUS_DEVICE_NOT_READY;
+        goto Exit;
+    }
+
+    RtlInitUnicodeString(&Path, Message.Path);
+
+    if (Message.Command == PyasCommandAddWhitelist) {
+        if (Path.Length == 0) {
+            Status = STATUS_INVALID_PARAMETER;
+            goto Exit;
+        }
+
+        AddDynamicWhitelist(&Path);
+        goto Exit;
+    }
+
+    if (Message.Command == PyasCommandRemoveWhitelist) {
+        if (Path.Length == 0) {
+            Status = STATUS_INVALID_PARAMETER;
+            goto Exit;
+        }
+
+        RemoveDynamicWhitelist(&Path);
+        goto Exit;
+    }
+
+    if (Message.Command == PyasCommandLoadRuleFile) {
+        if (Path.Length == 0) {
+            Status = STATUS_INVALID_PARAMETER;
+            goto Exit;
+        }
+
+        Status = LoadRuleFile(&Path);
+        goto Exit;
+    }
+
+    ClearDynamicRules();
+
+Exit:
+    ExReleaseRundownProtection(&GlobalData.PortRundown);
+    return Status;
 }
 
 static NTSTATUS PortConnect(
@@ -159,40 +409,88 @@ static NTSTATUS PortConnect(
     UNREFERENCED_PARAMETER(ConnectionContext);
     UNREFERENCED_PARAMETER(SizeOfContext);
 
-    if (!ConnectionPortCookie) {
-        return STATUS_INVALID_PARAMETER;
+    if (!ConnectionPortCookie) return STATUS_INVALID_PARAMETER;
+    *ConnectionPortCookie = NULL;
+
+    if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) {
+        return STATUS_DELETE_PENDING;
+    }
+
+    LONG DriverState = InterlockedCompareExchange(
+        &GlobalData.DriverState,
+        PyasDriverStateCold,
+        PyasDriverStateCold
+    );
+
+    if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) {
+        return STATUS_DEVICE_NOT_READY;
+    }
+
+    if (!ExAcquireRundownProtection(&GlobalData.PortRundown)) {
+        return STATUS_DELETE_PENDING;
     }
 
     PEPROCESS ConnectingProcess = PsGetCurrentProcess();
     ObReferenceObject(ConnectingProcess);
-    *ConnectionPortCookie = ConnectingProcess;
-
-    SafeSetPyasPid((ULONG)(ULONG_PTR)PsGetCurrentProcessId());
-    InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
 
+    NTSTATUS Status = STATUS_SUCCESS;
     KIRQL OldIrql;
     KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
 
-    GlobalData.ClientPort = ClientPort;
-    GlobalData.PyasProcess = ConnectingProcess;
-    ExReInitializeRundownProtection(&GlobalData.PortRundown);
+    DriverState = InterlockedCompareExchange(
+        &GlobalData.DriverState,
+        PyasDriverStateCold,
+        PyasDriverStateCold
+    );
+
+    if (InterlockedCompareExchange(&GlobalData.PortStopping, 0, 0) != 0) {
+        Status = STATUS_DELETE_PENDING;
+    }
+    else if (DriverState != PyasDriverStateRunning && DriverState != PyasDriverStateStopRetry) {
+        Status = STATUS_DEVICE_NOT_READY;
+    }
+    else if (GlobalData.ClientPort || GlobalData.PyasProcess) {
+        Status = STATUS_DEVICE_BUSY;
+    }
+    else {
+        GlobalData.ClientPort = ClientPort;
+        GlobalData.PyasProcess = ConnectingProcess;
+        *ConnectionPortCookie = ConnectingProcess;
+    }
 
     KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
+
+    if (!NT_SUCCESS(Status)) {
+        ObDereferenceObject(ConnectingProcess);
+        ExReleaseRundownProtection(&GlobalData.PortRundown);
+        return Status;
+    }
+
+    SafeSetPyasPid((ULONG)(ULONG_PTR)PsGetCurrentProcessId());
+    InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
+    ExReleaseRundownProtection(&GlobalData.PortRundown);
     return STATUS_SUCCESS;
 }
 
 static VOID PortDisconnect(PVOID ConnectionCookie) {
+    CloseCurrentClientPort();
+
+    BOOLEAN DisconnectedCurrentClient = FALSE;
+
     KIRQL OldIrql;
     KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
 
-    GlobalData.ClientPort = NULL;
-    GlobalData.PyasProcess = NULL;
+    if (ConnectionCookie && GlobalData.PyasProcess == (PEPROCESS)ConnectionCookie) {
+        GlobalData.PyasProcess = NULL;
+        DisconnectedCurrentClient = TRUE;
+    }
 
     KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
 
-    InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
-    SafeSetPyasPid(0);
-    ExWaitForRundownProtectionRelease(&GlobalData.PortRundown);
+    if (DisconnectedCurrentClient) {
+        InterlockedExchange(&GlobalData.UnloadAuthorized, 0);
+        SafeSetPyasPid(0);
+    }
 
     if (ConnectionCookie) {
         ObDereferenceObject((PEPROCESS)ConnectionCookie);
@@ -235,85 +533,84 @@ extern "C" NTSTATUS DriverEntry(
     KeInitializeSpinLock(&GlobalData.TrackerMutex);
     KeInitializeSpinLock(&GlobalData.PidLock);
     ExInitializeRundownProtection(&GlobalData.PortRundown);
+    KeInitializeEvent(&GlobalData.CleanupCompleteEvent, NotificationEvent, FALSE);
+    InterlockedExchange(&GlobalData.PortStopping, 0);
+    InterlockedExchange(&GlobalData.ClientCloseActive, 0);
+    InterlockedExchange(&GlobalData.DriverState, PyasDriverStateStarting);
 
-    InitializeRulesEngine();
+    NTSTATUS Status = STATUS_SUCCESS;
+    PSECURITY_DESCRIPTOR SecurityDescriptor = NULL;
+
+    Status = InitializeRulesEngine();
+    if (!NT_SUCCESS(Status)) {
+        InterlockedExchange(&GlobalData.DriverState, PyasDriverStateStopped);
+        KeSetEvent(&GlobalData.CleanupCompleteEvent, IO_NO_INCREMENT, FALSE);
+        return Status;
+    }
     g_RulesEngineInitialized = TRUE;
 
-    NTSTATUS status = LoadRulesFromDisk(RegistryPath);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
+    Status = LoadRulesFromDisk(RegistryPath);
+    if (!NT_SUCCESS(Status)) goto Failure;
+
+    Status = FltRegisterFilter(DriverObject, &FilterRegistration, &GlobalData.FilterHandle);
+    if (!NT_SUCCESS(Status)) goto Failure;
+
+    Status = FltBuildDefaultSecurityDescriptor(&SecurityDescriptor, FLT_PORT_ALL_ACCESS);
+    if (!NT_SUCCESS(Status)) goto Failure;
+
+    {
+        UNICODE_STRING PortName;
+        RtlInitUnicodeString(&PortName, PYAS_PORT_NAME);
+
+        OBJECT_ATTRIBUTES ObjectAttributes;
+        InitializeObjectAttributes(
+            &ObjectAttributes,
+            &PortName,
+            OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
+            NULL,
+            SecurityDescriptor
+        );
+
+        Status = FltCreateCommunicationPort(
+            GlobalData.FilterHandle,
+            &GlobalData.ServerPort,
+            &ObjectAttributes,
+            NULL,
+            PortConnect,
+            PortDisconnect,
+            PortMessage,
+            1
+        );
     }
 
-    status = FltRegisterFilter(DriverObject, &FilterRegistration, &GlobalData.FilterHandle);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
-
-    PSECURITY_DESCRIPTOR SecurityDescriptor = NULL;
-    status = FltBuildDefaultSecurityDescriptor(&SecurityDescriptor, FLT_PORT_ALL_ACCESS);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
-
-    UNICODE_STRING PortName;
-    RtlInitUnicodeString(&PortName, PYAS_PORT_NAME);
-
-    OBJECT_ATTRIBUTES ObjectAttributes;
-    InitializeObjectAttributes(
-        &ObjectAttributes,
-        &PortName,
-        OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
-        NULL,
-        SecurityDescriptor
-    );
-
-    status = FltCreateCommunicationPort(
-        GlobalData.FilterHandle,
-        &GlobalData.ServerPort,
-        &ObjectAttributes,
-        NULL,
-        PortConnect,
-        PortDisconnect,
-        PortMessage,
-        1
-    );
-
     FltFreeSecurityDescriptor(SecurityDescriptor);
+    if (!NT_SUCCESS(Status)) goto Failure;
 
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
-
-    status = InitializeRegistryProtection(DriverObject);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
+    Status = InitializeRegistryProtection(DriverObject);
+    if (!NT_SUCCESS(Status)) goto Failure;
     g_RegistryProtectionInitialized = TRUE;
 
-    status = PsSetCreateProcessNotifyRoutineEx(ProcessNotifyCallbackEx, FALSE);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
+    Status = InitializeProcessProtection();
+    if (!NT_SUCCESS(Status)) goto Failure;
+    g_ProcessProtectionInitialized = TRUE;
+
+    Status = PsSetCreateProcessNotifyRoutineEx(ProcessNotifyCallbackEx, FALSE);
+    if (!NT_SUCCESS(Status)) goto Failure;
     g_ProcessNotifyRegistered = TRUE;
 
-    status = InitializeMemoryProtection(DriverObject);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
+    Status = InitializeMemoryProtection(DriverObject);
+    if (!NT_SUCCESS(Status)) goto Failure;
 
-    status = FltStartFiltering(GlobalData.FilterHandle);
-    if (!NT_SUCCESS(status)) {
-        CleanupDriverState();
-        return status;
-    }
+    Status = FltStartFiltering(GlobalData.FilterHandle);
+    if (!NT_SUCCESS(Status)) goto Failure;
 
     GlobalData.Initialized = TRUE;
+    InterlockedExchange(&GlobalData.DriverState, PyasDriverStateRunning);
     return STATUS_SUCCESS;
+
+Failure:
+    CleanupDriverState(TRUE);
+    InterlockedExchange(&GlobalData.DriverState, PyasDriverStateStopped);
+    KeSetEvent(&GlobalData.CleanupCompleteEvent, IO_NO_INCREMENT, FALSE);
+    return Status;
 }

Plugins/Filter/ProtectBoot.cpp

@@ -23,10 +23,12 @@ FLT_PREOP_CALLBACK_STATUS ProtectBoot_PreDeviceControl(PFLT_CALLBACK_DATA Data,
 
         HANDLE Pid = PsGetCurrentProcessId();
         ULONG RuleCode = 0;
+        BOOLEAN Kill = FALSE;
 
-        if (EvaluateDeviceRule(Pid, &RuleCode)) {
+        if (EvaluateDeviceRule(Pid, &RuleCode, &Kill)) {
             UNICODE_STRING MsgStr = RTL_CONSTANT_STRING(L"Disk_Wiper_Attempt");
             SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, MsgStr.Buffer, MsgStr.Length);
+            QueueRuleProcessTermination(Pid, Kill);
             Data->IoStatus.Status = STATUS_ACCESS_DENIED;
             return FLT_PREOP_COMPLETE;
         }

Plugins/Filter/ProtectFile.cpp

@@ -31,8 +31,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreCreate(PFLT_CALLBACK_DATA Data, PCFLT_R
 
     if (Operation != 0) {
         ULONG RuleCode = 0;
-        if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode)) {
+        BOOLEAN Kill = FALSE;
+        if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode, &Kill)) {
             SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length);
+            QueueRuleProcessTermination(Pid, Kill);
             Data->IoStatus.Status = STATUS_ACCESS_DENIED;
             callbackStatus = FLT_PREOP_COMPLETE;
         }
@@ -73,8 +75,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreSetInfo(PFLT_CALLBACK_DATA Data, PCFLT_
     }
 
     ULONG RuleCode = 0;
-    if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode)) {
+    BOOLEAN Kill = FALSE;
+    if (EvaluateFileRule(Pid, &nameInfo->Name, Operation, NULL, 0, &RuleCode, &Kill)) {
         SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length);
+        QueueRuleProcessTermination(Pid, Kill);
         Data->IoStatus.Status = STATUS_ACCESS_DENIED;
         callbackStatus = FLT_PREOP_COMPLETE;
     }
@@ -112,8 +116,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreWrite(PFLT_CALLBACK_DATA Data, PCFLT_RE
     if (Mdl) SystemBuffer = MmGetSystemAddressForMdlSafe(Mdl, NormalPagePriority | MdlMappingNoExecute);
 
     ULONG RuleCode = 0;
-    if (EvaluateFileRule(Pid, &nameInfo->Name, OP_WRITE, SystemBuffer, WriteLength, &RuleCode)) {
+    BOOLEAN Kill = FALSE;
+    if (EvaluateFileRule(Pid, &nameInfo->Name, OP_WRITE, SystemBuffer, WriteLength, &RuleCode, &Kill)) {
         SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length);
+        QueueRuleProcessTermination(Pid, Kill);
         Data->IoStatus.Status = STATUS_ACCESS_DENIED;
         callbackStatus = FLT_PREOP_COMPLETE;
     }

Plugins/Filter/ProtectProcess.cpp

@@ -1,21 +1,285 @@
 #include "DriverCommon.h"
 
-VOID ProcessNotifyCallbackEx(PEPROCESS Process, HANDLE ProcessId, PPS_CREATE_NOTIFY_INFO CreateInfo) {
-    UNREFERENCED_PARAMETER(Process);
-    UNREFERENCED_PARAMETER(ProcessId);
+constexpr auto PROCESS_RELATION_BUCKET_COUNT = 1024;
+constexpr auto PROCESS_RELATION_BUCKET_WAYS = 4;
 
-    if (CreateInfo != NULL) {
-        if (KeGetCurrentIrql() > APC_LEVEL) return;
+typedef HANDLE(NTAPI* PPS_GET_PROCESS_INHERITED_FROM_UNIQUE_PROCESS_ID)(PEPROCESS Process);
 
-        HANDLE ParentId = CreateInfo->CreatingThreadId.UniqueProcess;
-        if (!ParentId) return;
+typedef struct _PROCESS_RELATION_ENTRY {
+    HANDLE ProcessId;
+    HANDLE ParentProcessId;
+    HANDLE CreatorProcessId;
+    LARGE_INTEGER CreateTime;
+    LARGE_INTEGER ParentCreateTime;
+    ULONGLONG Sequence;
+} PROCESS_RELATION_ENTRY, * PPROCESS_RELATION_ENTRY;
 
-        ULONG RuleCode = 0;
-        if (EvaluateProcessRule(ParentId, CreateInfo->ImageFileName, CreateInfo->CommandLine, &RuleCode)) {
-            CreateInfo->CreationStatus = STATUS_ACCESS_DENIED;
-            if (CreateInfo->ImageFileName) {
-                SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)ParentId, CreateInfo->ImageFileName->Buffer, CreateInfo->ImageFileName->Length);
-            }
+static PROCESS_RELATION_ENTRY g_ProcessRelations[PROCESS_RELATION_BUCKET_COUNT][PROCESS_RELATION_BUCKET_WAYS];
+static KSPIN_LOCK g_ProcessRelationLock;
+static EX_RUNDOWN_REF g_ProcessRundown;
+static volatile LONG g_ProcessStopping = 1;
+static volatile LONG64 g_ProcessRelationSequence = 0;
+static BOOLEAN g_ProcessProtectionInitialized = FALSE;
+static PPS_GET_PROCESS_INHERITED_FROM_UNIQUE_PROCESS_ID g_GetInheritedProcessId = NULL;
+
+static ULONG GetProcessRelationBucket(HANDLE ProcessId) {
+    return ((ULONG)(ULONG_PTR)ProcessId >> 2) & (PROCESS_RELATION_BUCKET_COUNT - 1);
+}
+
+static BOOLEAN QueryProcessCreateTime(HANDLE ProcessId, PLARGE_INTEGER CreateTime) {
+    if (CreateTime) CreateTime->QuadPart = 0;
+    if (!ProcessId || !CreateTime || KeGetCurrentIrql() > APC_LEVEL) return FALSE;
+
+    PEPROCESS Process = NULL;
+    NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &Process);
+    if (!NT_SUCCESS(Status)) return FALSE;
+
+    CreateTime->QuadPart = PsGetProcessCreateTimeQuadPart(Process);
+    ObDereferenceObject(Process);
+    return TRUE;
+}
+
+static VOID StoreProcessRelation(
+    PEPROCESS Process,
+    HANDLE ProcessId,
+    HANDLE ParentProcessId,
+    HANDLE CreatorProcessId,
+    PLARGE_INTEGER KnownParentCreateTime
+) {
+    if (!Process || !ProcessId) return;
+
+    PROCESS_RELATION_ENTRY NewEntry = { 0 };
+    NewEntry.ProcessId = ProcessId;
+    NewEntry.ParentProcessId = ParentProcessId;
+    NewEntry.CreatorProcessId = CreatorProcessId;
+    NewEntry.CreateTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process);
+
+    if (KnownParentCreateTime) {
+        NewEntry.ParentCreateTime = *KnownParentCreateTime;
+    }
+    else if (ParentProcessId) {
+        QueryProcessCreateTime(ParentProcessId, &NewEntry.ParentCreateTime);
+    }
+
+    NewEntry.Sequence = (ULONGLONG)InterlockedIncrement64(&g_ProcessRelationSequence);
+
+    ULONG Bucket = GetProcessRelationBucket(ProcessId);
+    ULONG Selected = 0;
+
+    KIRQL OldIrql;
+    KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql);
+
+    for (ULONG Index = 0; Index < PROCESS_RELATION_BUCKET_WAYS; Index++) {
+        PPROCESS_RELATION_ENTRY Entry = &g_ProcessRelations[Bucket][Index];
+
+        if (Entry->ProcessId == ProcessId || Entry->ProcessId == NULL) {
+            Selected = Index;
+            break;
+        }
+
+        if (Entry->Sequence < g_ProcessRelations[Bucket][Selected].Sequence) {
+            Selected = Index;
         }
     }
-}
+
+    g_ProcessRelations[Bucket][Selected] = NewEntry;
+    KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql);
+}
+
+static VOID RemoveProcessRelation(HANDLE ProcessId) {
+    if (!ProcessId) return;
+
+    ULONG Bucket = GetProcessRelationBucket(ProcessId);
+
+    KIRQL OldIrql;
+    KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql);
+
+    for (ULONG Index = 0; Index < PROCESS_RELATION_BUCKET_WAYS; Index++) {
+        PPROCESS_RELATION_ENTRY Entry = &g_ProcessRelations[Bucket][Index];
+        if (Entry->ProcessId == ProcessId) {
+            RtlZeroMemory(Entry, sizeof(*Entry));
+            break;
+        }
+    }
+
+    KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql);
+}
+
+BOOLEAN GetProcessRelation(
+    HANDLE ProcessId,
+    PHANDLE ParentProcessId,
+    PHANDLE CreatorProcessId
+) {
+    if (ParentProcessId) *ParentProcessId = NULL;
+    if (CreatorProcessId) *CreatorProcessId = NULL;
+    if (!ProcessId || !g_ProcessProtectionInitialized || KeGetCurrentIrql() > APC_LEVEL) return FALSE;
+
+    PEPROCESS Process = NULL;
+    NTSTATUS Status = PsLookupProcessByProcessId(ProcessId, &Process);
+    if (!NT_SUCCESS(Status)) return FALSE;
+
+    LARGE_INTEGER CurrentCreateTime = { 0 };
+    CurrentCreateTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process);
+
+    PROCESS_RELATION_ENTRY Snapshot = { 0 };
+    ULONG Bucket = GetProcessRelationBucket(ProcessId);
+
+    KIRQL OldIrql;
+    KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql);
+
+    for (ULONG Index = 0; Index < PROCESS_RELATION_BUCKET_WAYS; Index++) {
+        PPROCESS_RELATION_ENTRY Entry = &g_ProcessRelations[Bucket][Index];
+        if (Entry->ProcessId == ProcessId) {
+            Snapshot = *Entry;
+            break;
+        }
+    }
+
+    KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql);
+
+    BOOLEAN CacheValid = Snapshot.ProcessId == ProcessId &&
+        Snapshot.CreateTime.QuadPart == CurrentCreateTime.QuadPart;
+
+    if (CacheValid && Snapshot.ParentProcessId && Snapshot.ParentCreateTime.QuadPart != 0) {
+        LARGE_INTEGER LiveParentCreateTime = { 0 };
+        if (QueryProcessCreateTime(Snapshot.ParentProcessId, &LiveParentCreateTime) &&
+            LiveParentCreateTime.QuadPart != Snapshot.ParentCreateTime.QuadPart) {
+            CacheValid = FALSE;
+        }
+    }
+
+    if (CacheValid) {
+        if (ParentProcessId) *ParentProcessId = Snapshot.ParentProcessId;
+        if (CreatorProcessId) *CreatorProcessId = Snapshot.CreatorProcessId;
+        ObDereferenceObject(Process);
+        return Snapshot.ParentProcessId != NULL || Snapshot.CreatorProcessId != NULL;
+    }
+
+    if (Snapshot.ProcessId == ProcessId) {
+        RemoveProcessRelation(ProcessId);
+    }
+
+    HANDLE LiveParentProcessId = NULL;
+    LARGE_INTEGER LiveParentCreateTime = { 0 };
+
+    if (g_GetInheritedProcessId) {
+        LiveParentProcessId = g_GetInheritedProcessId(Process);
+    }
+
+    if (!LiveParentProcessId || LiveParentProcessId == ProcessId) {
+        ObDereferenceObject(Process);
+        return FALSE;
+    }
+
+    if (!QueryProcessCreateTime(LiveParentProcessId, &LiveParentCreateTime) ||
+        LiveParentCreateTime.QuadPart > CurrentCreateTime.QuadPart) {
+        ObDereferenceObject(Process);
+        return FALSE;
+    }
+
+    StoreProcessRelation(
+        Process,
+        ProcessId,
+        LiveParentProcessId,
+        NULL,
+        &LiveParentCreateTime
+    );
+
+    if (ParentProcessId) *ParentProcessId = LiveParentProcessId;
+    if (CreatorProcessId) *CreatorProcessId = NULL;
+
+    ObDereferenceObject(Process);
+    return TRUE;
+}
+
+NTSTATUS InitializeProcessProtection() {
+    if (g_ProcessProtectionInitialized) return STATUS_SUCCESS;
+
+    KeInitializeSpinLock(&g_ProcessRelationLock);
+    RtlZeroMemory(g_ProcessRelations, sizeof(g_ProcessRelations));
+    InterlockedExchange64(&g_ProcessRelationSequence, 0);
+    ExInitializeRundownProtection(&g_ProcessRundown);
+
+    UNICODE_STRING RoutineName = RTL_CONSTANT_STRING(L"PsGetProcessInheritedFromUniqueProcessId");
+    g_GetInheritedProcessId = reinterpret_cast<PPS_GET_PROCESS_INHERITED_FROM_UNIQUE_PROCESS_ID>(
+        MmGetSystemRoutineAddress(&RoutineName)
+        );
+
+    InterlockedExchange(&g_ProcessStopping, 0);
+    g_ProcessProtectionInitialized = TRUE;
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS UninitializeProcessProtection() {
+    if (!g_ProcessProtectionInitialized) return STATUS_SUCCESS;
+
+    InterlockedExchange(&g_ProcessStopping, 1);
+    ExWaitForRundownProtectionRelease(&g_ProcessRundown);
+
+    KIRQL OldIrql;
+    KeAcquireSpinLock(&g_ProcessRelationLock, &OldIrql);
+    RtlZeroMemory(g_ProcessRelations, sizeof(g_ProcessRelations));
+    KeReleaseSpinLock(&g_ProcessRelationLock, OldIrql);
+
+    g_GetInheritedProcessId = NULL;
+    g_ProcessProtectionInitialized = FALSE;
+    return STATUS_SUCCESS;
+}
+
+VOID ProcessNotifyCallbackEx(
+    PEPROCESS Process,
+    HANDLE ProcessId,
+    PPS_CREATE_NOTIFY_INFO CreateInfo
+) {
+    if (!g_ProcessProtectionInitialized) return;
+    if (InterlockedCompareExchange(&g_ProcessStopping, 0, 0) != 0) return;
+    if (!ExAcquireRundownProtection(&g_ProcessRundown)) return;
+
+    if (!CreateInfo) {
+        RemoveProcessRelation(ProcessId);
+        ExReleaseRundownProtection(&g_ProcessRundown);
+        return;
+    }
+
+    if (!NT_SUCCESS(CreateInfo->CreationStatus)) {
+        ExReleaseRundownProtection(&g_ProcessRundown);
+        return;
+    }
+
+    HANDLE CreatorPid = CreateInfo->CreatingThreadId.UniqueProcess;
+    HANDLE ParentPid = CreateInfo->ParentProcessId;
+    ULONG RuleCode = 0;
+    BOOLEAN Kill = FALSE;
+
+    BOOLEAN Blocked = EvaluateProcessCreateRule(
+        CreatorPid,
+        ParentPid,
+        ProcessId,
+        CreateInfo->ImageFileName,
+        CreateInfo->CommandLine,
+        CreateInfo->FileOpenNameAvailable ? TRUE : FALSE,
+        CreateInfo->IsSubsystemProcess ? TRUE : FALSE,
+        &RuleCode,
+        &Kill
+    );
+
+    if (Blocked) {
+        CreateInfo->CreationStatus = STATUS_ACCESS_DENIED;
+
+        if (CreateInfo->ImageFileName && CreateInfo->ImageFileName->Buffer) {
+            SendMessageToUser(
+                RuleCode,
+                (ULONG)(ULONG_PTR)CreatorPid,
+                CreateInfo->ImageFileName->Buffer,
+                CreateInfo->ImageFileName->Length
+            );
+        }
+
+        QueueRuleProcessTermination(CreatorPid, Kill);
+    }
+    else {
+        StoreProcessRelation(Process, ProcessId, ParentPid, CreatorPid, NULL);
+    }
+
+    ExReleaseRundownProtection(&g_ProcessRundown);
+}

Plugins/Filter/ProtectRegistry.cpp

@@ -100,8 +100,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
 
             HANDLE Pid = PsGetCurrentProcessId();
             ULONG RuleCode = 0;
-            if (EvaluateRegistryRule(Pid, &FullPath, OP_CREATE | OP_WRITE, &RuleCode)) {
+            BOOLEAN Kill = FALSE;
+            if (EvaluateRegistryRule(Pid, &FullPath, OP_CREATE | OP_WRITE, &RuleCode, &Kill)) {
                 SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
+                QueueRuleProcessTermination(Pid, Kill);
                 status = STATUS_ACCESS_DENIED;
             }
             PyasFree(FullPathBuffer);
@@ -117,8 +119,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
             if (KeyName) {
                 HANDLE Pid = PsGetCurrentProcessId();
                 ULONG RuleCode = 0;
-                if (EvaluateRegistryRule(Pid, KeyName, OP_DELETE, &RuleCode)) {
+                BOOLEAN Kill = FALSE;
+                if (EvaluateRegistryRule(Pid, KeyName, OP_DELETE, &RuleCode, &Kill)) {
                     SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, KeyName->Buffer, KeyName->Length);
+                    QueueRuleProcessTermination(Pid, Kill);
                     status = STATUS_ACCESS_DENIED;
                 }
                 CmCallbackReleaseKeyObjectIDEx(KeyName);
@@ -156,8 +160,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
 
                         HANDLE Pid = PsGetCurrentProcessId();
                         ULONG RuleCode = 0;
-                        if (EvaluateRegistryRule(Pid, &FullPath, OP_WRITE, &RuleCode)) {
+                        BOOLEAN Kill = FALSE;
+                        if (EvaluateRegistryRule(Pid, &FullPath, OP_WRITE, &RuleCode, &Kill)) {
                             SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
+                            QueueRuleProcessTermination(Pid, Kill);
                             status = STATUS_ACCESS_DENIED;
                         }
                         PyasFree(Buffer);
@@ -198,8 +204,10 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
 
                         HANDLE Pid = PsGetCurrentProcessId();
                         ULONG RuleCode = 0;
-                        if (EvaluateRegistryRule(Pid, &FullPath, OP_DELETE, &RuleCode)) {
+                        BOOLEAN Kill = FALSE;
+                        if (EvaluateRegistryRule(Pid, &FullPath, OP_DELETE, &RuleCode, &Kill)) {
                             SendMessageToUser(RuleCode, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
+                            QueueRuleProcessTermination(Pid, Kill);
                             status = STATUS_ACCESS_DENIED;
                         }
                         PyasFree(Buffer);
@@ -215,20 +223,20 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
 }
 
 NTSTATUS InitializeRegistryProtection(PDRIVER_OBJECT DriverObject) {
-    static UNICODE_STRING Altitude = RTL_CONSTANT_STRING(L"320000.PYAS.Cm");
-    NTSTATUS status = CmRegisterCallbackEx(RegistryCallback, &Altitude, DriverObject, NULL, &Cookie, NULL);
+    if (Cookie.QuadPart != 0) return STATUS_SUCCESS;
 
-    if (!NT_SUCCESS(status)) {
-        static UNICODE_STRING FallbackAltitude = RTL_CONSTANT_STRING(L"320000.PYAS.Cm.Fallback");
-        status = CmRegisterCallbackEx(RegistryCallback, &FallbackAltitude, DriverObject, NULL, &Cookie, NULL);
-    }
-
-    return status;
+    UNICODE_STRING Altitude;
+    RtlInitUnicodeString(&Altitude, L"320000.4201");
+    return CmRegisterCallbackEx(RegistryCallback, &Altitude, DriverObject, NULL, &Cookie, NULL);
 }
 
-VOID UninitializeRegistryProtection() {
-    if (Cookie.QuadPart != 0) {
-        CmUnRegisterCallback(Cookie);
+NTSTATUS UninitializeRegistryProtection() {
+    if (Cookie.QuadPart == 0) return STATUS_SUCCESS;
+
+    NTSTATUS Status = CmUnRegisterCallback(Cookie);
+    if (NT_SUCCESS(Status)) {
         Cookie.QuadPart = 0;
     }
-}
+
+    return Status;
+}