PR feedback

Signed-off-by: Alan Jowett <[email protected]>
microsoft · Feb 17, 2024 · 2465cbc · 2465cbc
1 parent 41d4eae
commit 2465cbc
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 33 deletions.
diff --git a/src/ke.cpp b/src/ke.cpp
@@ -620,25 +620,6 @@ class _usersim_emulated_dpc
     bool terminate;
 };
 
-NTSTATUS
-KeGetProcessorNumberFromIndex(ULONG ProcessorIndex, _Out_ PPROCESSOR_NUMBER ProcNumber)
-{
-    if (ProcessorIndex >= GetMaximumProcessorCount(ALL_PROCESSOR_GROUPS)) {
-        return STATUS_INVALID_PARAMETER;
-    }
-
-    ProcNumber->Group = (WORD)(ProcessorIndex / 64);
-    ProcNumber->Number = (UCHAR)(ProcessorIndex % 64);
-    ProcNumber->Reserved = 0;
-    return STATUS_SUCCESS;
-}
-
-ULONG
-KeGetProcessorIndexFromNumber(_In_ PPROCESSOR_NUMBER ProcNumber)
-{
-    return (ULONG)ProcNumber->Group * 64 + (ULONG)ProcNumber->Number;
-}
-
 void
 KeInitializeDpc(
     _Out_ __drv_aliasesMem PRKDPC dpc,

diff --git a/src/platform_user.cpp b/src/platform_user.cpp
@@ -381,7 +381,39 @@ 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

diff --git a/tests/ke_test.cpp b/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]")
@@ -438,16 +444,3 @@ TEST_CASE("event", "[ke]")
     REQUIRE(wait_status == STATUS_TIMEOUT);
     REQUIRE(end_time - start_time >= 1000);
 }
-
-TEST_CASE("processor count", "[ke]")
-{
-    uint32_t count = KeQueryMaximumProcessorCount();
-
-    REQUIRE(count > 0);
-
-    for (uint32_t i = 0; i < count; i++) {
-        PROCESSOR_NUMBER processor_number;
-        REQUIRE(NT_SUCCESS(KeGetProcessorNumberFromIndex(i)));
-        REQUIRE(KeGetProcessorIndexFromNumber(&processor_number) == i);
-    }
-}