User/austinl/ps get process exit status (#191)

inc/usersim/ps.h

@@ -14,6 +14,16 @@ USERSIM_API
 HANDLE
 PsGetCurrentProcessId();
 
+typedef NTSTATUS (*PGETPROCESSEXITSTATUS)(_In_ PEPROCESS process);
+
+USERSIM_API
+NTSTATUS
+PsGetProcessExitStatus(_In_ PEPROCESS Process);
+
+USERSIM_API
+void
+usersime_set_process_exit_status_callback(_In_ PGETPROCESSEXITSTATUS callback);
+
 USERSIM_API
 _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE PsGetCurrentThreadId();
 

src/ps.cpp

@@ -15,6 +15,25 @@ _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE PsGetCurrentThreadId()
     return (HANDLE)(uintptr_t)GetCurrentThreadId();
 }
 
+static PGETPROCESSEXITSTATUS _usersim_get_process_exit_status_callback = NULL;
+
+NTSTATUS
+PsGetProcessExitStatus(_In_ PEPROCESS Process) {
+    if (_usersim_get_process_exit_status_callback != NULL) {
+        return _usersim_get_process_exit_status_callback(Process);
+    }
+
+    // Fall back to a failure code
+    return -1;
+}
+
+USERSIM_API
+void
+usersime_set_process_exit_status_callback(_In_ PGETPROCESSEXITSTATUS callback)
+{
+    _usersim_get_process_exit_status_callback = callback;
+}
+
 static PCREATE_PROCESS_NOTIFY_ROUTINE_EX _usersim_process_creation_notify_routine = NULL;
 
 USERSIM_API

tests/ps_test.cpp

@@ -41,4 +41,24 @@ TEST_CASE("PsSetCreateProcessNotifyRoutineEx", "[ps]")
     // Remove the routine. Should succeed.
     status = PsSetCreateProcessNotifyRoutineEx(notify_routine, TRUE);
     REQUIRE(status == STATUS_SUCCESS);
+}
+
+TEST_CASE("PsGetProcessExitStatus", "[ps]")
+{
+    // If no callback is installed, we default to -1
+    auto status = PsGetProcessExitStatus((PEPROCESS)0);
+    REQUIRE(status == -1);
+
+    usersime_set_process_exit_status_callback([](PEPROCESS proc) -> NTSTATUS { return ((int)proc)  +  1; });
+
+    status = PsGetProcessExitStatus((PEPROCESS)0);
+    REQUIRE(status == 1);
+
+    status = PsGetProcessExitStatus((PEPROCESS)1234);
+    REQUIRE(status == 1235);
+
+    // Setting back to a NULL callback reverts to returning -1
+    usersime_set_process_exit_status_callback(NULL);
+    status = PsGetProcessExitStatus((PEPROCESS)0);
+    REQUIRE(status == -1);
 }

tests/tests.vcxproj

@@ -156,6 +156,7 @@
     <ClCompile Include="mm_test.cpp" />
     <ClCompile Include="nmr_test.cpp" />
     <ClCompile Include="ob_test.cpp" />
+    <ClCompile Include="ps_test.cpp" />
     <ClCompile Include="rtl_test.cpp" />
     <ClCompile Include="se_test.cpp" />
     <ClCompile Include="wdf_test.cpp" />

tests/tests.vcxproj.filters

@@ -48,5 +48,8 @@
     <ClCompile Include="io_test.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="ps_test.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
 </Project>