Add process create and exit time APIs (#193)

Signed-off-by: Austin Lamb <[email protected]>
microsoft · May 3, 2024 · ad2276a · ad2276a
1 parent 5f70910
commit ad2276a
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 2 deletions.
diff --git a/inc/usersim/ps.h b/inc/usersim/ps.h
@@ -23,6 +23,24 @@ USERSIM_API
 void
 usersime_set_process_exit_status_callback(_In_ PGETPROCESSEXITSTATUS callback);
 
+typedef LONGLONG (*PGETPROCESSCREATETIMEQUADPART)(_In_ PEPROCESS process);
+
+USERSIM_API
+_IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI LONGLONG PsGetProcessCreateTimeQuadPart(_In_ PEPROCESS Process);
+
+USERSIM_API
+void
+usersime_set_process_create_time_quadpart_callback(_In_ PGETPROCESSCREATETIMEQUADPART callback);
+
+typedef LARGE_INTEGER (*PGETPROCESSEXITTIME)(VOID);
+
+USERSIM_API
+_IRQL_requires_max_(APC_LEVEL) NTKERNELAPI LARGE_INTEGER PsGetProcessExitTime(VOID);
+
+USERSIM_API
+void
+usersime_set_process_exit_time_callback(_In_ PGETPROCESSEXITTIME callback);
+
 USERSIM_API
 _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE PsGetCurrentThreadId();
 

diff --git a/src/ps.cpp b/src/ps.cpp
@@ -17,8 +17,10 @@ _IRQL_requires_max_(DISPATCH_LEVEL) NTKERNELAPI HANDLE PsGetCurrentThreadId()
 
 static PGETPROCESSEXITSTATUS _usersim_get_process_exit_status_callback = NULL;
 
+USERSIM_API
 NTSTATUS
-PsGetProcessExitStatus(_In_ PEPROCESS Process) {
+PsGetProcessExitStatus(_In_ PEPROCESS Process)
+{
     if (_usersim_get_process_exit_status_callback != NULL) {
         return _usersim_get_process_exit_status_callback(Process);
     }
@@ -34,6 +36,47 @@ usersime_set_process_exit_status_callback(_In_ PGETPROCESSEXITSTATUS callback)
     _usersim_get_process_exit_status_callback = callback;
 }
 
+static PGETPROCESSCREATETIMEQUADPART _usersim_get_process_create_time_quadpart_callback = NULL;
+
+USERSIM_API
+LONGLONG
+PsGetProcessCreateTimeQuadPart(_In_ PEPROCESS Process)
+{
+    if (_usersim_get_process_create_time_quadpart_callback != NULL) {
+        return _usersim_get_process_create_time_quadpart_callback(Process);
+    }
+
+    // Fall back to the beginning of time
+    return 0;
+}
+
+USERSIM_API
+void
+usersime_set_process_create_time_quadpart_callback(_In_ PGETPROCESSCREATETIMEQUADPART callback)
+{
+    _usersim_get_process_create_time_quadpart_callback = callback;
+}
+
+static PGETPROCESSEXITTIME _usersim_get_process_exit_time_callback = NULL;
+
+USERSIM_API
+LARGE_INTEGER PsGetProcessExitTime(VOID)
+{
+    if (_usersim_get_process_exit_time_callback != NULL) {
+        return _usersim_get_process_exit_time_callback();
+    }
+
+    // Fall back to the beginning of time
+    LARGE_INTEGER li = {0};
+    return li;
+}
+
+void
+usersime_set_process_exit_time_callback(_In_ PGETPROCESSEXITTIME callback)
+{
+    _usersim_get_process_exit_time_callback = callback;
+}
+
 static PCREATE_PROCESS_NOTIFY_ROUTINE_EX _usersim_process_creation_notify_routine = NULL;
 
 USERSIM_API

diff --git a/tests/ps_test.cpp b/tests/ps_test.cpp
@@ -49,7 +49,7 @@ TEST_CASE("PsGetProcessExitStatus", "[ps]")
     auto status = PsGetProcessExitStatus((PEPROCESS)0);
     REQUIRE(status == -1);
 
-    usersime_set_process_exit_status_callback([](PEPROCESS proc) -> NTSTATUS { return ((int)proc)  +  1; });
+    usersime_set_process_exit_status_callback([](PEPROCESS proc) -> NTSTATUS { return ((int)proc) + 1; });
 
     status = PsGetProcessExitStatus((PEPROCESS)0);
     REQUIRE(status == 1);
@@ -61,4 +61,41 @@ TEST_CASE("PsGetProcessExitStatus", "[ps]")
     usersime_set_process_exit_status_callback(NULL);
     status = PsGetProcessExitStatus((PEPROCESS)0);
     REQUIRE(status == -1);
+}
+
+TEST_CASE("PsGetProcessCreateTimeQuadPart", "[ps]")
+{
+    // If no callback is installed, we default to 0
+    auto time = PsGetProcessCreateTimeQuadPart((PEPROCESS)0);
+    REQUIRE(time == 0);
+
+    usersime_set_process_create_time_quadpart_callback([](PEPROCESS proc) -> LONGLONG { return ((int)proc) + 1; });
+
+    time = PsGetProcessCreateTimeQuadPart((PEPROCESS)0);
+    REQUIRE(time == 1);
+
+    time = PsGetProcessCreateTimeQuadPart((PEPROCESS)1234);
+    REQUIRE(time == 1235);
+
+    // Setting back to a NULL callback reverts to returning 0
+    usersime_set_process_create_time_quadpart_callback(NULL);
+    time = PsGetProcessCreateTimeQuadPart((PEPROCESS)0);
+    REQUIRE(time == 0);
+}
+
+TEST_CASE("PsGetProcessExitTime", "[ps]")
+{
+    // If no callback is installed, we default to 0
+    auto time = PsGetProcessExitTime();
+    REQUIRE(time.QuadPart == 0);
+
+    usersime_set_process_exit_time_callback([]() -> LARGE_INTEGER { return {123}; });
+
+    time = PsGetProcessExitTime();
+    REQUIRE(time.QuadPart == 123);
+
+    // Setting back to a NULL callback reverts to returning 0
+    usersime_set_process_exit_time_callback(NULL);
+    time = PsGetProcessExitTime();
+    REQUIRE(time.QuadPart == 0);
 }