Add APIs to handle PROCESSOR_NUMBER

inc/usersim/ke.h

@@ -45,6 +45,14 @@ USERSIM_API
 ULONG
 KeGetCurrentProcessorNumberEx(_Out_opt_ PPROCESSOR_NUMBER ProcNumber);
 
+USERSIM_API
+NTSTATUS
+KeGetProcessorNumberFromIndex(ULONG ProcessorIndex, _Out_ PPROCESSOR_NUMBER ProcNumber);
+
+USERSIM_API
+ULONG
+KeGetProcessorIndexFromNumber(_In_ PPROCESSOR_NUMBER ProcNumber);
+
 #pragma region irqls
 typedef uint8_t KIRQL;
 typedef KIRQL* PKIRQL;
@@ -245,7 +253,7 @@ struct _KDPC
 {
     usersim_list_entry_t entry;
     void* context;
-    CCHAR cpu_id;
+    uint32_t cpu_id;
     void* parameter_1;
     void* parameter_2;
     PKDEFERRED_ROUTINE work_item_routine;
@@ -274,6 +282,10 @@ USERSIM_API
 void
 KeSetTargetProcessorDpc(_Inout_ PRKDPC dpc, CCHAR number);
 
+USERSIM_API
+void
+KeSetTargetProcessorDpcEx(_Inout_ PRKDPC dpc, PPROCESSOR_NUMBER proc_number);
+
 void
 usersim_initialize_dpcs();
 

src/ke.cpp

@@ -638,6 +638,12 @@ KeSetTargetProcessorDpc(_Inout_ PRKDPC dpc, CCHAR number)
     dpc->cpu_id = number;
 }
 
+void
+KeSetTargetProcessorDpcEx(_Inout_ PRKDPC dpc, PPROCESSOR_NUMBER proc_number)
+{
+    dpc->cpu_id = KeGetProcessorIndexFromNumber(proc_number);
+}
+
 BOOLEAN
 KeInsertQueueDpc(_Inout_ PRKDPC dpc, _In_opt_ PVOID system_argument1, _In_opt_ __drv_aliasesMem PVOID system_argument2)
 {

src/platform_user.cpp

@@ -381,7 +381,38 @@ KeGetCurrentProcessorNumberEx(_Out_opt_ PPROCESSOR_NUMBER ProcNumber)
     }
 
     // Compute the CPU index from the group and number.
-    return _usersim_platform_group_to_index_map[processor_number.Group] + processor_number.Number;
+    return KeGetProcessorIndexFromNumber(&processor_number);
+}
+
+ULONG
+KeGetProcessorIndexFromNumber(_In_ PPROCESSOR_NUMBER ProcNumber)
+{
+    // Compute the CPU index from the group and number.
+    return _usersim_platform_group_to_index_map[ProcNumber->Group] + ProcNumber->Number;
+}
+
+NTSTATUS
+KeGetProcessorNumberFromIndex(ULONG ProcessorIndex, _Out_ PPROCESSOR_NUMBER ProcNumber)
+{
+    if (ProcessorIndex >= _usersim_platform_maximum_processor_count) {
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    // Find the group that contains the processor index.
+    for (size_t i = 0; i < _usersim_platform_group_to_index_map.size(); i++) {
+        if (ProcessorIndex < _usersim_platform_group_to_index_map[i]) {
+            // This is the first group with a processor index greater than the input.
+            // The previous group contains the processor index and the number is the difference.
+            ProcNumber->Group = (uint16_t)(i - 1);
+            ProcNumber->Number = (uint8_t)(ProcessorIndex - _usersim_platform_group_to_index_map[i - 1]);
+            return STATUS_SUCCESS;
+        }
+    }
+
+    // The processor index is in the last group.
+    ProcNumber->Group = (uint16_t)(_usersim_platform_group_to_index_map.size() - 1);
+    ProcNumber->Number = (uint8_t)(ProcessorIndex - _usersim_platform_group_to_index_map.back());
+    return STATUS_SUCCESS;
 }
 
 uint64_t

tests/ke_test.cpp

@@ -93,6 +93,12 @@ TEST_CASE("processor count", "[ke]")
     REQUIRE(KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS) == count);
     REQUIRE(KeQueryActiveProcessorCount() == count);
     REQUIRE(KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS) == count);
+
+    for (uint32_t i = 0; i < count; i++) {
+        PROCESSOR_NUMBER processor_number;
+        REQUIRE(NT_SUCCESS(KeGetProcessorNumberFromIndex(i, &processor_number)));
+        REQUIRE(KeGetProcessorIndexFromNumber(&processor_number) == i);
+    }
 }
 
 TEST_CASE("semaphore", "[ke]")
@@ -437,4 +443,4 @@ TEST_CASE("event", "[ke]")
 
     REQUIRE(wait_status == STATUS_TIMEOUT);
     REQUIRE(end_time - start_time >= 1000);
-}
+}