Add support for simulating PsSetCreateProcessNotifyRoutineEx

Signed-off-by: Alan Jowett (from Dev Box) <[email protected]>

inc/usersim/ps.h

@@ -4,6 +4,7 @@
 #pragma once
 #include "..\src\platform.h"
 #include "usersim\io.h"
+#include "usersim\rtl.h"
 
 CXPLAT_EXTERN_C_BEGIN
 
@@ -16,4 +17,32 @@ PsGetCurrentProcessId();
 USERSIM_API
 _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE PsGetCurrentThreadId();
 
+typedef struct _PS_CREATE_NOTIFY_INFO
+{
+    _In_ SIZE_T Size;
+    union
+    {
+        _In_ ULONG Flags;
+        struct
+        {
+            _In_ ULONG FileOpenNameAvailable : 1;
+            _In_ ULONG IsSubsystemProcess : 1;
+            _In_ ULONG Reserved : 30;
+        };
+    };
+    _In_ HANDLE ParentProcessId;
+    _In_ HANDLE CreatingThreadId;
+    _Inout_ struct _FILE_OBJECT* FileObject;
+    _In_ PCUNICODE_STRING ImageFileName;
+    _In_opt_ PCUNICODE_STRING CommandLine;
+    _Inout_ NTSTATUS CreationStatus;
+} PS_CREATE_NOTIFY_INFO, *PPS_CREATE_NOTIFY_INFO;
+
+typedef VOID (*PCREATE_PROCESS_NOTIFY_ROUTINE_EX)(
+    _Inout_ PEPROCESS process, _In_ HANDLE process_id, _Inout_opt_ PPS_CREATE_NOTIFY_INFO create_info);
+
+USERSIM_API
+NTSTATUS
+PsSetCreateProcessNotifyRoutineEx(_In_ PCREATE_PROCESS_NOTIFY_ROUTINE_EX notify_routine, _In_ BOOLEAN remove);
+
 CXPLAT_EXTERN_C_END

src/ps.cpp

@@ -1,14 +1,53 @@
 // Copyright (c) Microsoft Corporation
 // SPDX-License-Identifier: MIT
 
-#include "platform.h"
 #include "kernel_um.h"
+#include "platform.h"
 #include "usersim/ps.h"
 
 // Ps* functions.
 
 HANDLE
 PsGetCurrentProcessId() { return (HANDLE)(uintptr_t)GetCurrentProcessId(); }
 
-_IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE
-PsGetCurrentThreadId() { return (HANDLE)(uintptr_t)GetCurrentThreadId(); }
+_IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE PsGetCurrentThreadId()
+{
+    return (HANDLE)(uintptr_t)GetCurrentThreadId();
+}
+
+static PCREATE_PROCESS_NOTIFY_ROUTINE_EX _usersim_process_creation_notify_routine = NULL;
+
+USERSIM_API
+NTSTATUS
+PsSetCreateProcessNotifyRoutineEx(_In_ PCREATE_PROCESS_NOTIFY_ROUTINE_EX notify_routine, _In_ BOOLEAN remove)
+{
+    if (notify_routine == NULL) {
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    // Remove the routine.
+    if (remove) {
+        if (_usersim_process_creation_notify_routine == NULL) {
+            return STATUS_INVALID_PARAMETER;
+        }
+        _usersim_process_creation_notify_routine = NULL;
+        return STATUS_SUCCESS;
+    }
+    // Only one routine is supported.
+    if (_usersim_process_creation_notify_routine != NULL) {
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    // Set the routine.
+    _usersim_process_creation_notify_routine = notify_routine;
+    return STATUS_SUCCESS;
+}
+
+void
+usersime_invoke_process_creation_notify_routine(
+    _Inout_ PEPROCESS process, _In_ HANDLE process_id, _Inout_opt_ PPS_CREATE_NOTIFY_INFO create_info)
+{
+    if (_usersim_process_creation_notify_routine != NULL) {
+        _usersim_process_creation_notify_routine(process, process_id, create_info);
+    }
+}

tests/CMakeLists.txt

@@ -18,6 +18,7 @@ add_executable(usersim_tests
   mm_test.cpp
   nmr_test.cpp
   ob_test.cpp
+  ps_test.cpp
   rtl_test.cpp
   se_test.cpp
   wdf_test.cpp

tests/ps_test.cpp

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: MIT
+
+#if !defined(CMAKE_NUGET)
+#include <catch2/catch_all.hpp>
+#else
+#include <catch2/catch.hpp>
+#endif
+#include "usersim/ps.h"
+
+TEST_CASE("PsSetCreateProcessNotifyRoutineEx", "[ps]")
+{
+    auto notify_routine = [](PEPROCESS, HANDLE, PPS_CREATE_NOTIFY_INFO) {};
+    // First call should succeed.
+    auto status = PsSetCreateProcessNotifyRoutineEx(notify_routine, FALSE);
+    REQUIRE(status == STATUS_SUCCESS);
+    // Try to set the routine again, should fail.
+    status = PsSetCreateProcessNotifyRoutineEx(notify_routine, FALSE);
+    REQUIRE(status == STATUS_INVALID_PARAMETER);
+
+    // Remove the routine. Should succeed.
+    status = PsSetCreateProcessNotifyRoutineEx(notify_routine, TRUE);
+    REQUIRE(status == STATUS_SUCCESS);
+
+    // Try to remove the routine again, should fail.
+    status = PsSetCreateProcessNotifyRoutineEx(notify_routine, TRUE);
+    REQUIRE(status == STATUS_INVALID_PARAMETER);
+
+    // Try to remove with a NULL routine, should fail.
+    status = PsSetCreateProcessNotifyRoutineEx(NULL, TRUE);
+    REQUIRE(status == STATUS_INVALID_PARAMETER);
+
+    // Try set with a NULL routine, should fail.
+    status = PsSetCreateProcessNotifyRoutineEx(NULL, FALSE);
+    REQUIRE(status == STATUS_INVALID_PARAMETER);
+
+    // Set the routine again, should succeed.
+    status = PsSetCreateProcessNotifyRoutineEx(notify_routine, FALSE);
+    REQUIRE(status == STATUS_SUCCESS);
+
+    // Remove the routine. Should succeed.
+    status = PsSetCreateProcessNotifyRoutineEx(notify_routine, TRUE);
+    REQUIRE(status == STATUS_SUCCESS);
+}