diff --git a/inc/usersim/ke.h b/inc/usersim/ke.h index 4413cc8..ad16fd2 100644 --- a/inc/usersim/ke.h +++ b/inc/usersim/ke.h @@ -149,6 +149,76 @@ USERSIM_API ULONG KeQueryActiveProcessorCountEx(_In_ USHORT group_number); +typedef enum _KE_PROCESSOR_CHANGE_NOTIFY_STATE +{ + KeProcessorAddStartNotify, + KeProcessorAddCompleteNotify, + KeProcessorAddFailureNotify, +} KE_PROCESSOR_CHANGE_NOTIFY_STATE; + +typedef struct _KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT +{ + KE_PROCESSOR_CHANGE_NOTIFY_STATE State; + ULONG NtNumber; +} KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT; +typedef KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT* PKE_PROCESSOR_CHANGE_NOTIFY_CONTEXT; + +#define KE_PROCESSOR_CHANGE_ADD_EXISTING 0x00000001 + +typedef VOID (*KE_PROCESSOR_CHANGE_CALLBACK)( + _In_ void* callback_context, + _In_ PKE_PROCESSOR_CHANGE_NOTIFY_CONTEXT change_context, + _Inout_ PNTSTATUS operation_status); + +USERSIM_API +void* +KeRegisterProcessorChangeCallback( + _In_ KE_PROCESSOR_CHANGE_CALLBACK callback, _In_opt_ void* callback_context, _In_ ULONG flags); + +USERSIM_API +void +KeDeregisterProcessorChangeCallback(_In_ void* callback_handle); + +USERSIM_API +_Must_inspect_result_ NTSTATUS +usersim_set_active_processor_count(_In_ ULONG active_processor_count); + +USERSIM_API +void +usersim_reset_active_processor_count(); + +/** + * @brief Start a simulated processor hot-add for the next inactive processor. + * + * The test hook models processor additions serially. Callers must start with + * the current active processor count and must finish the same processor with + * usersim_notify_processor_add_complete() or usersim_notify_processor_add_failure() + * before starting another add. + */ +USERSIM_API +_Must_inspect_result_ NTSTATUS +usersim_notify_processor_add_start(_In_ ULONG processor_index); + +/** + * @brief Complete a simulated processor hot-add started by usersim_notify_processor_add_start(). + * + * The processor index must match the currently pending add and becomes visible + * via KeQueryActiveProcessorCount*() once the completion notification succeeds. + */ +USERSIM_API +_Must_inspect_result_ NTSTATUS +usersim_notify_processor_add_complete(_In_ ULONG processor_index); + +/** + * @brief Fail a simulated processor hot-add started by usersim_notify_processor_add_start(). + * + * The processor index must match the currently pending add. Failing an add + * leaves the active processor count unchanged. + */ +USERSIM_API +_Must_inspect_result_ NTSTATUS +usersim_notify_processor_add_failure(_In_ ULONG processor_index); + USERSIM_API KAFFINITY KeSetSystemAffinityThreadEx(KAFFINITY affinity); diff --git a/src/ke.cpp b/src/ke.cpp index cf26c71..e3041cd 100644 --- a/src/ke.cpp +++ b/src/ke.cpp @@ -5,7 +5,10 @@ #include "usersim/ke.h" #include "utilities.h" +#include +#include #include +#include #include #include #include @@ -32,13 +35,67 @@ thread_local int static uint32_t _usersim_original_priority_class; static std::vector _usersim_dispatch_locks; +static std::mutex _usersim_processor_change_callback_lock; +static std::mutex _usersim_processor_change_state_lock; +static std::atomic _usersim_active_processor_count_override = 0; +static std::optional _usersim_pending_processor_add; static TP_POOL* _usersim_threadpool = nullptr; static std::optional _usersim_threadpool_callback_environment = std::nullopt; +typedef struct _usersim_processor_change_callback_registration +{ + KE_PROCESSOR_CHANGE_CALLBACK callback; + void* callback_context; +} usersim_processor_change_callback_registration_t; + +static std::vector> + _usersim_processor_change_callbacks; + static NTSTATUS _wait_for_kevent(_Inout_ KEVENT* event, _In_opt_ PLARGE_INTEGER timeout); +static _Must_inspect_result_ NTSTATUS +_usersim_invoke_processor_change_callback( + _In_ const std::shared_ptr& registration, + _In_ ULONG processor_index, + _In_ KE_PROCESSOR_CHANGE_NOTIFY_STATE state, + _Inout_ NTSTATUS& operation_status) +{ + KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT change_context = { + .State = state, + .NtNumber = processor_index, + }; + + registration->callback(registration->callback_context, &change_context, &operation_status); + return operation_status; +} + +static _Must_inspect_result_ NTSTATUS +_usersim_notify_processor_change(_In_ ULONG processor_index, _In_ KE_PROCESSOR_CHANGE_NOTIFY_STATE state) +{ + if (processor_index >= KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS)) { + return STATUS_INVALID_PARAMETER; + } + + std::vector> callbacks; + { + std::lock_guard lock(_usersim_processor_change_callback_lock); + callbacks = _usersim_processor_change_callbacks; + } + + NTSTATUS operation_status = STATUS_SUCCESS; + + for (const auto& registration : callbacks) { + _usersim_invoke_processor_change_callback(registration, processor_index, state, operation_status); + if (!NT_SUCCESS(operation_status)) { + break; + } + } + + return operation_status; +} + usersim_result_t usersim_initialize_irql() { @@ -356,10 +413,169 @@ ULONG KeQueryMaximumProcessorCountEx(_In_ USHORT group_number) { return GetMaximumProcessorCount(group_number); } ULONG -KeQueryActiveProcessorCount() { return KeQueryMaximumProcessorCount(); } +KeQueryActiveProcessorCount() { return KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); } ULONG -KeQueryActiveProcessorCountEx(_In_ USHORT group_number) { return KeQueryMaximumProcessorCountEx(group_number); } +KeQueryActiveProcessorCountEx(_In_ USHORT group_number) +{ + ULONG active_processor_count_override = _usersim_active_processor_count_override.load(std::memory_order_acquire); + if (active_processor_count_override == 0) { + return KeQueryMaximumProcessorCountEx(group_number); + } + + return min(active_processor_count_override, KeQueryMaximumProcessorCountEx(group_number)); +} + +void* +KeRegisterProcessorChangeCallback( + _In_ KE_PROCESSOR_CHANGE_CALLBACK callback, _In_opt_ void* callback_context, _In_ ULONG flags) +{ + std::shared_ptr registration( + new (std::nothrow) usersim_processor_change_callback_registration_t{ + .callback = callback, .callback_context = callback_context}); + if (registration == nullptr) { + return nullptr; + } + + { + std::lock_guard lock(_usersim_processor_change_callback_lock); + _usersim_processor_change_callbacks.push_back(registration); + } + + if ((flags & KE_PROCESSOR_CHANGE_ADD_EXISTING) != 0) { + ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); + NTSTATUS status = STATUS_SUCCESS; + for (ULONG processor_index = 0; processor_index < active_processor_count; processor_index++) { + status = _usersim_invoke_processor_change_callback( + registration, processor_index, KeProcessorAddStartNotify, status); + if (!NT_SUCCESS(status)) { + KeDeregisterProcessorChangeCallback(registration.get()); + return nullptr; + } + } + + status = STATUS_SUCCESS; + for (ULONG processor_index = 0; processor_index < active_processor_count; processor_index++) { + status = _usersim_invoke_processor_change_callback( + registration, processor_index, KeProcessorAddCompleteNotify, status); + if (!NT_SUCCESS(status)) { + KeDeregisterProcessorChangeCallback(registration.get()); + return nullptr; + } + } + } + + return registration.get(); +} + +void +KeDeregisterProcessorChangeCallback(_In_ void* callback_handle) +{ + std::lock_guard lock(_usersim_processor_change_callback_lock); + auto iterator = std::find_if( + _usersim_processor_change_callbacks.begin(), + _usersim_processor_change_callbacks.end(), + [callback_handle](const std::shared_ptr& registration) { + return registration.get() == callback_handle; + }); + if (iterator != _usersim_processor_change_callbacks.end()) { + _usersim_processor_change_callbacks.erase(iterator); + } +} + +_Must_inspect_result_ NTSTATUS +usersim_set_active_processor_count(_In_ ULONG active_processor_count) +{ + ULONG maximum_processor_count = KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS); + if (active_processor_count == 0 || active_processor_count > maximum_processor_count) { + return STATUS_INVALID_PARAMETER; + } + + std::lock_guard lock(_usersim_processor_change_state_lock); + _usersim_pending_processor_add.reset(); + _usersim_active_processor_count_override.store(active_processor_count, std::memory_order_release); + return STATUS_SUCCESS; +} + +void +usersim_reset_active_processor_count() +{ + std::lock_guard lock(_usersim_processor_change_state_lock); + _usersim_pending_processor_add.reset(); + _usersim_active_processor_count_override.store(0, std::memory_order_release); +} + +_Must_inspect_result_ NTSTATUS +usersim_notify_processor_add_start(_In_ ULONG processor_index) +{ + { + std::lock_guard lock(_usersim_processor_change_state_lock); + ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); + if (_usersim_pending_processor_add.has_value() || processor_index != active_processor_count) { + return STATUS_INVALID_PARAMETER; + } + + _usersim_pending_processor_add = processor_index; + } + + NTSTATUS status = _usersim_notify_processor_change(processor_index, KeProcessorAddStartNotify); + if (!NT_SUCCESS(status)) { + std::lock_guard lock(_usersim_processor_change_state_lock); + if (_usersim_pending_processor_add == processor_index) { + _usersim_pending_processor_add.reset(); + } + } + + return status; +} + +_Must_inspect_result_ NTSTATUS +usersim_notify_processor_add_complete(_In_ ULONG processor_index) +{ + { + std::lock_guard lock(_usersim_processor_change_state_lock); + ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); + if (!_usersim_pending_processor_add.has_value() || *_usersim_pending_processor_add != processor_index || + processor_index != active_processor_count) { + return STATUS_INVALID_PARAMETER; + } + } + + NTSTATUS status = _usersim_notify_processor_change(processor_index, KeProcessorAddCompleteNotify); + + std::lock_guard lock(_usersim_processor_change_state_lock); + if (_usersim_pending_processor_add == processor_index) { + _usersim_pending_processor_add.reset(); + } + if (NT_SUCCESS(status)) { + ULONG active_processor_count_override = + _usersim_active_processor_count_override.load(std::memory_order_relaxed); + _usersim_active_processor_count_override.store( + max(active_processor_count_override, processor_index + 1), std::memory_order_release); + } + + return status; +} + +_Must_inspect_result_ NTSTATUS +usersim_notify_processor_add_failure(_In_ ULONG processor_index) +{ + { + std::lock_guard lock(_usersim_processor_change_state_lock); + if (!_usersim_pending_processor_add.has_value() || *_usersim_pending_processor_add != processor_index) { + return STATUS_INVALID_PARAMETER; + } + } + + NTSTATUS status = _usersim_notify_processor_change(processor_index, KeProcessorAddFailureNotify); + + std::lock_guard lock(_usersim_processor_change_state_lock); + if (_usersim_pending_processor_add == processor_index) { + _usersim_pending_processor_add.reset(); + } + + return status; +} KAFFINITY KeSetSystemAffinityThreadEx(KAFFINITY affinity) @@ -1093,4 +1309,4 @@ KeExpandKernelStackAndCalloutEx( return STATUS_SUCCESS; } -#pragma endregion events \ No newline at end of file +#pragma endregion events diff --git a/tests/ke_test.cpp b/tests/ke_test.cpp index 1fedaee..08816ce 100644 --- a/tests/ke_test.cpp +++ b/tests/ke_test.cpp @@ -10,6 +10,7 @@ #include "usersim/mm.h" #include +#include TEST_CASE("irql", "[ke]") { @@ -128,6 +129,133 @@ TEST_CASE("processor count", "[ke]") } } +typedef struct _processor_change_notification_record +{ + KE_PROCESSOR_CHANGE_NOTIFY_STATE state; + ULONG processor_index; +} processor_change_notification_record_t; + +typedef struct _processor_change_callback_context +{ + std::vector notifications; + ULONG invocation_count = 0; + void* handle_to_deregister = nullptr; +} processor_change_callback_context_t; + +static VOID +_record_processor_change_callback( + _In_ void* callback_context, + _In_ PKE_PROCESSOR_CHANGE_NOTIFY_CONTEXT change_context, + _Inout_ PNTSTATUS operation_status) +{ + UNREFERENCED_PARAMETER(operation_status); + + auto* context = reinterpret_cast(callback_context); + context->notifications.push_back({change_context->State, change_context->NtNumber}); + context->invocation_count++; +} + +static VOID +_deregister_processor_change_callback( + _In_ void* callback_context, + _In_ PKE_PROCESSOR_CHANGE_NOTIFY_CONTEXT change_context, + _Inout_ PNTSTATUS operation_status) +{ + UNREFERENCED_PARAMETER(change_context); + UNREFERENCED_PARAMETER(operation_status); + + auto* context = reinterpret_cast(callback_context); + context->invocation_count++; + if (context->handle_to_deregister != nullptr) { + KeDeregisterProcessorChangeCallback(context->handle_to_deregister); + } +} + +TEST_CASE("processor change callback add existing", "[ke]") +{ + processor_change_callback_context_t existing_callback_context = {}; + processor_change_callback_context_t add_existing_callback_context = {}; + void* existing_handle = + KeRegisterProcessorChangeCallback(_record_processor_change_callback, &existing_callback_context, 0); + REQUIRE(existing_handle != nullptr); + + void* add_existing_handle = KeRegisterProcessorChangeCallback( + _record_processor_change_callback, &add_existing_callback_context, KE_PROCESSOR_CHANGE_ADD_EXISTING); + REQUIRE(add_existing_handle != nullptr); + + ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); + REQUIRE(existing_callback_context.invocation_count == 0); + REQUIRE(add_existing_callback_context.notifications.size() == active_processor_count * 2); + for (ULONG processor_index = 0; processor_index < active_processor_count; processor_index++) { + const auto& start_notification = add_existing_callback_context.notifications[processor_index]; + REQUIRE(start_notification.state == KeProcessorAddStartNotify); + REQUIRE(start_notification.processor_index == processor_index); + + const auto& complete_notification = + add_existing_callback_context.notifications[processor_index + active_processor_count]; + REQUIRE(complete_notification.state == KeProcessorAddCompleteNotify); + REQUIRE(complete_notification.processor_index == processor_index); + } + + KeDeregisterProcessorChangeCallback(existing_handle); + KeDeregisterProcessorChangeCallback(add_existing_handle); +} + +TEST_CASE("processor change callback deregistration is deferred until notifications finish", "[ke]") +{ + if (KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS) < 2) { + return; + } + + processor_change_callback_context_t deregistering_callback_context = {}; + processor_change_callback_context_t removed_callback_context = {}; + + void* deregistering_handle = + KeRegisterProcessorChangeCallback(_deregister_processor_change_callback, &deregistering_callback_context, 0); + REQUIRE(deregistering_handle != nullptr); + + void* removed_handle = + KeRegisterProcessorChangeCallback(_record_processor_change_callback, &removed_callback_context, 0); + REQUIRE(removed_handle != nullptr); + deregistering_callback_context.handle_to_deregister = removed_handle; + + REQUIRE(usersim_set_active_processor_count(1) == STATUS_SUCCESS); + REQUIRE(usersim_notify_processor_add_start(1) == STATUS_SUCCESS); + REQUIRE(deregistering_callback_context.invocation_count == 1); + REQUIRE(removed_callback_context.invocation_count == 1); + + REQUIRE(usersim_notify_processor_add_failure(1) == STATUS_SUCCESS); + REQUIRE(deregistering_callback_context.invocation_count == 2); + REQUIRE(removed_callback_context.invocation_count == 1); + + KeDeregisterProcessorChangeCallback(removed_handle); + KeDeregisterProcessorChangeCallback(deregistering_handle); + usersim_reset_active_processor_count(); +} + +TEST_CASE("processor add notifications are serialized", "[ke]") +{ + if (KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS) < 2) { + return; + } + + REQUIRE(usersim_set_active_processor_count(1) == STATUS_SUCCESS); + + REQUIRE(usersim_notify_processor_add_start(2) == STATUS_INVALID_PARAMETER); + REQUIRE(usersim_notify_processor_add_start(1) == STATUS_SUCCESS); + REQUIRE(usersim_notify_processor_add_start(1) == STATUS_INVALID_PARAMETER); + REQUIRE(usersim_notify_processor_add_complete(2) == STATUS_INVALID_PARAMETER); + REQUIRE(usersim_notify_processor_add_complete(1) == STATUS_SUCCESS); + REQUIRE(KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS) == 2); + + REQUIRE(usersim_notify_processor_add_complete(1) == STATUS_INVALID_PARAMETER); + REQUIRE(usersim_notify_processor_add_start(2) == STATUS_SUCCESS); + REQUIRE(usersim_notify_processor_add_failure(2) == STATUS_SUCCESS); + REQUIRE(KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS) == 2); + + usersim_reset_active_processor_count(); +} + TEST_CASE("semaphore", "[ke]") { KSEMAPHORE semaphore;