From cd1080d16cbfdbe791aefe40bfd566f4aaba8ee0 Mon Sep 17 00:00:00 2001 From: "Shankar Seal (from Dev Box)" Date: Mon, 3 Aug 2026 01:51:25 -0700 Subject: [PATCH 1/3] fix mock NMR deregister race. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6e257781-cd5a-4ae3-9943-0c217ccad54e --- .github/workflows/build.yml | 4 +- scripts/Test-FaultInjection.ps1 | 2 +- src/nmr_impl.cpp | 13 ++- src/nmr_impl.h | 3 +- tests/nmr_test.cpp | 141 ++++++++++++++++++++++++++++++++ 5 files changed, 155 insertions(+), 8 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c2c7e9e..6b350c5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -81,7 +81,7 @@ jobs: working-directory: ./${{env.BUILD_PLATFORM}}/${{env.BUILD_CONFIGURATION }} shell: cmd run: | - powershell ..\..\scripts\Test-FaultInjection.ps1 ${{env.DUMP_PATH}} ${{env.TEST_TIMEOUT}} ".\usersim_tests.exe" 4 + powershell ..\..\scripts\Test-FaultInjection.ps1 ${{env.DUMP_PATH}} ${{env.TEST_TIMEOUT}} ".\usersim_tests.exe" 4 "~[no_fi]" build-cmake: timeout-minutes: 15 @@ -151,4 +151,4 @@ jobs: working-directory: ./build/bin/${{env.BUILD_CONFIGURATION}} shell: cmd run: | - powershell ..\..\..\scripts\Test-FaultInjection.ps1 ${{env.DUMP_PATH}} ${{env.TEST_TIMEOUT}} ".\usersim_tests.exe" 4 + powershell ..\..\..\scripts\Test-FaultInjection.ps1 ${{env.DUMP_PATH}} ${{env.TEST_TIMEOUT}} ".\usersim_tests.exe" 4 "~[no_fi]" diff --git a/scripts/Test-FaultInjection.ps1 b/scripts/Test-FaultInjection.ps1 index 402ce3a..15e76d2 100644 --- a/scripts/Test-FaultInjection.ps1 +++ b/scripts/Test-FaultInjection.ps1 @@ -12,7 +12,7 @@ param ($OutputFolder, $Timeout, $TestProgram, $StackDepth) # Gather list of all possible tests -$tests = & $TestProgram "--list-tests" "--verbosity=quiet" +$tests = & $TestProgram "--list-tests" "--verbosity=quiet" "~[no_fi]" $env:CXPLAT_FAULT_INJECTION_SIMULATION = $StackDepth diff --git a/src/nmr_impl.cpp b/src/nmr_impl.cpp index 2b6feb3..66caa95 100644 --- a/src/nmr_impl.cpp +++ b/src/nmr_impl.cpp @@ -189,9 +189,15 @@ nmr_t::bind(_Inout_ client_registration& client, _Inout_ provider_registration& if (!NT_SUCCESS(status)) { unbind_complete(*binding_ptr); } else { + bool should_begin_unbind = false; std::unique_lock l(lock); binding_ptr->client_binding_status = binding_status::Ready; binding_ptr->provider_binding_status = binding_status::Ready; + should_begin_unbind = binding_ptr->client.deregistering || binding_ptr->provider.deregistering; + l.unlock(); + if (should_begin_unbind) { + (void)begin_unbind(*binding_ptr); + } } }}; } @@ -221,12 +227,13 @@ nmr_t::unbind_complete(_Inout_ binding& binding) bindings_changed.notify_all(); } -bool // true if pending, false if complete. +bool nmr_t::begin_unbind(_Inout_ binding& binding) { std::unique_lock l(lock); if (binding.client_binding_status != Ready || binding.provider_binding_status != Ready) { - // Unbind already started. + // A Start binding is already published and contributes to binding_count, so deregistration + // must keep waiting even though detach cannot begin until attach finishes and reaches Ready. return true; } binding.client_binding_status = BeginUnbind; @@ -356,7 +363,7 @@ nmr_t::perform_bind( } template -bool // true if pending, false if complete +bool nmr_t::perform_unbind( _Inout_ initiator_collection_t& initiator_collection, _In_ initiator_collection_t::value_type::first_type initiator_handle) diff --git a/src/nmr_impl.h b/src/nmr_impl.h index a8b7ea7..c10bffd 100644 --- a/src/nmr_impl.h +++ b/src/nmr_impl.h @@ -148,7 +148,6 @@ typedef class nmr_t UnbindComplete ///< Client or provider detach returned STATUS_SUCCESS or called NmrBindingDetachClientComplete ///< or NmrBindingDetachProviderComplete. }; - struct binding { provider_registration& provider; @@ -250,7 +249,7 @@ typedef class nmr_t * @brief Start the process of unbinding a client from a provider. * * @param[in] binding_handle Binding handle to unbind. - * @retval true Either the client or provider returned pending. + * @retval true Either unbind cannot start yet or it is pending/in progress. * @retval false Both the client and provider returned successfully. */ bool diff --git a/tests/nmr_test.cpp b/tests/nmr_test.cpp index 16a8d04..274aba4 100644 --- a/tests/nmr_test.cpp +++ b/tests/nmr_test.cpp @@ -8,6 +8,9 @@ #endif #include "../src/framework.h" #include <../km/netioddk.h> +#include +#include +#include NPIID test_npiid = {0}; @@ -154,6 +157,98 @@ NPI_PROVIDER_CHARACTERISTICS _test_provider_characteristics = { #pragma endregion test_nmr_provider +#pragma region smoke_nmr_client_provider + +// Use a distinct NPI ID so smoke registrations never match the test_npiid registrations above, +// preventing cross-contamination when test_npiid registrations are left over from failed tests. +NPIID smoke_npiid = {1}; + +NPI_REGISTRATION_INSTANCE _smoke_client_registration_instance = { + .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &smoke_npiid}; +NPI_REGISTRATION_INSTANCE _smoke_provider_registration_instance = { + .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &smoke_npiid}; + +static NTSTATUS +_smoke_client_attach_provider( + _In_ HANDLE nmr_binding_handle, + _In_opt_ void* client_context, + _In_ NPI_REGISTRATION_INSTANCE* provider_registration_instance) +{ + UNREFERENCED_PARAMETER(client_context); + UNREFERENCED_PARAMETER(provider_registration_instance); + + void* provider_binding_context = nullptr; + const void* provider_dispatch = nullptr; + return NmrClientAttachProvider( + nmr_binding_handle, + reinterpret_cast(nmr_binding_handle), + TEST_CLIENT_DISPATCH, + &provider_binding_context, + &provider_dispatch); +} + +static NTSTATUS +_smoke_client_detach_provider(_In_ void* client_binding_context) +{ + UNREFERENCED_PARAMETER(client_binding_context); + return STATUS_SUCCESS; +} + +static void +_smoke_client_cleanup_binding_context(_In_ void* client_binding_context) +{ + UNREFERENCED_PARAMETER(client_binding_context); +} + +NPI_CLIENT_CHARACTERISTICS _smoke_client_characteristics = { + .Length = sizeof(NPI_CLIENT_CHARACTERISTICS), + .ClientAttachProvider = (PNPI_CLIENT_ATTACH_PROVIDER_FN)_smoke_client_attach_provider, + .ClientDetachProvider = _smoke_client_detach_provider, + .ClientCleanupBindingContext = _smoke_client_cleanup_binding_context, + .ClientRegistrationInstance = _smoke_client_registration_instance}; + +static NTSTATUS +_smoke_provider_attach_client( + _In_ HANDLE nmr_binding_handle, + _In_opt_ void* provider_context, + _In_ NPI_REGISTRATION_INSTANCE* client_registration_instance, + _In_ void* client_binding_context, + _In_ const void* client_dispatch, + _Outptr_ void** provider_binding_context, + _Outptr_ const void** provider_dispatch) +{ + UNREFERENCED_PARAMETER(nmr_binding_handle); + UNREFERENCED_PARAMETER(provider_context); + UNREFERENCED_PARAMETER(client_registration_instance); + UNREFERENCED_PARAMETER(client_dispatch); + + *provider_binding_context = client_binding_context; + *provider_dispatch = TEST_PROVIDER_DISPATCH; + return STATUS_SUCCESS; +} + +static NTSTATUS +_smoke_provider_detach_client(_In_ void* provider_binding_context) +{ + UNREFERENCED_PARAMETER(provider_binding_context); + return STATUS_SUCCESS; +} + +static void +_smoke_provider_cleanup_binding_context(_In_ void* provider_binding_context) +{ + UNREFERENCED_PARAMETER(provider_binding_context); +} + +NPI_PROVIDER_CHARACTERISTICS _smoke_provider_characteristics = { + .Length = sizeof(NPI_PROVIDER_CHARACTERISTICS), + .ProviderAttachClient = (PNPI_PROVIDER_ATTACH_CLIENT_FN)_smoke_provider_attach_client, + .ProviderDetachClient = _smoke_provider_detach_client, + .ProviderCleanupBindingContext = _smoke_provider_cleanup_binding_context, + .ProviderRegistrationInstance = _smoke_provider_registration_instance}; + +#pragma endregion smoke_nmr_client_provider + TEST_CASE("NmrRegisterClient", "[nmr]") { HANDLE nmr_client_handle; @@ -322,4 +417,50 @@ TEST_CASE("NmrRegisterProvider with async deregister", "[nmr]") REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); +} + +TEST_CASE("concurrent register/deregister smoke", "[nmr][no_fi]") +{ + constexpr size_t iteration_count = 1000; + std::atomic provider_iterations{0}; + std::atomic client_iterations{0}; + + std::thread provider_thread([iteration_count, &provider_iterations]() { + for (size_t i = 0; i < iteration_count; i++) { + provider_iterations++; + HANDLE nmr_provider_handle = nullptr; + NTSTATUS register_status = + NmrRegisterProvider(&_smoke_provider_characteristics, nullptr, &nmr_provider_handle); + if (!NT_SUCCESS(register_status)) { + continue; + } + + NTSTATUS deregister_status = NmrDeregisterProvider(nmr_provider_handle); + if (deregister_status == STATUS_PENDING) { + (void)NmrWaitForProviderDeregisterComplete(nmr_provider_handle); + } + } + }); + + std::thread client_thread([iteration_count, &client_iterations]() { + for (size_t i = 0; i < iteration_count; i++) { + client_iterations++; + HANDLE nmr_client_handle = nullptr; + NTSTATUS register_status = NmrRegisterClient(&_smoke_client_characteristics, nullptr, &nmr_client_handle); + if (!NT_SUCCESS(register_status)) { + continue; + } + + NTSTATUS deregister_status = NmrDeregisterClient(nmr_client_handle); + if (deregister_status == STATUS_PENDING) { + (void)NmrWaitForClientDeregisterComplete(nmr_client_handle); + } + } + }); + + provider_thread.join(); + client_thread.join(); + + std::cout << "provider_iterations=" << provider_iterations.load() << ", client_iterations=" + << client_iterations.load() << std::endl; } \ No newline at end of file From fc90bf0b89906d81bb36fc72656bca25524bf883 Mon Sep 17 00:00:00 2001 From: "Shankar Seal (from Dev Box)" Date: Wed, 5 Aug 2026 16:20:17 -0700 Subject: [PATCH 2/3] PR Feedback. --- tests/nmr_test.cpp | 946 +++++++++++++++++++++++---------------------- 1 file changed, 481 insertions(+), 465 deletions(-) diff --git a/tests/nmr_test.cpp b/tests/nmr_test.cpp index 274aba4..26370a8 100644 --- a/tests/nmr_test.cpp +++ b/tests/nmr_test.cpp @@ -1,466 +1,482 @@ -// Copyright (c) Microsoft Corporation -// SPDX-License-Identifier: MIT - -#if !defined(CMAKE_NUGET) -#include -#else -#include -#endif -#include "../src/framework.h" -#include <../km/netioddk.h> -#include -#include -#include - -NPIID test_npiid = {0}; - -#pragma region test_nmr_client - -#define TEST_CLIENT_DISPATCH ((const void*)0x1) - -NPI_REGISTRATION_INSTANCE _test_client_registration_instance = { - .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &test_npiid}; - -typedef struct -{ - bool allocated; - HANDLE nmr_binding_handle; - void* provider_binding_context; - const void* provider_dispatch; -} test_client_binding_context_t; - -static test_client_binding_context_t _test_client_binding_context = {.allocated = false}; -static bool _test_client_async_deregister = false; - -static NTSTATUS -_test_client_attach_provider( - _In_ HANDLE nmr_binding_handle, - _In_opt_ void* client_context, - _In_ NPI_REGISTRATION_INSTANCE* provider_registration_instance) -{ - // Verify not already bound. - REQUIRE(_test_client_binding_context.allocated == false); - _test_client_binding_context.allocated = true; - - NTSTATUS status = NmrClientAttachProvider( - nmr_binding_handle, - &_test_client_binding_context, - TEST_CLIENT_DISPATCH, - &_test_client_binding_context.provider_binding_context, - &_test_client_binding_context.provider_dispatch); - - if (NT_SUCCESS(status)) { - _test_client_binding_context.nmr_binding_handle = nmr_binding_handle; - } - - return status; -} - -static NTSTATUS -_test_client_detach_provider(_In_ void* client_binding_context) -{ - test_client_binding_context_t* context = (test_client_binding_context_t*)client_binding_context; - REQUIRE(context->allocated); - - _test_client_binding_context.provider_binding_context = nullptr; - _test_client_binding_context.provider_dispatch = nullptr; - - return (_test_client_async_deregister) ? STATUS_PENDING : STATUS_SUCCESS; -} - -static void -_test_client_cleanup_binding_context(_In_ void* client_binding_context) -{ - test_client_binding_context_t* context = (test_client_binding_context_t*)client_binding_context; - REQUIRE(context->allocated); - context->allocated = false; - context->nmr_binding_handle = nullptr; -} - -NPI_CLIENT_CHARACTERISTICS _test_client_characteristics = { - .Length = sizeof(NPI_CLIENT_CHARACTERISTICS), - .ClientAttachProvider = (PNPI_CLIENT_ATTACH_PROVIDER_FN)_test_client_attach_provider, - .ClientDetachProvider = _test_client_detach_provider, - .ClientCleanupBindingContext = _test_client_cleanup_binding_context, - .ClientRegistrationInstance = _test_client_registration_instance}; - -#pragma endregion test_nmr_client -#pragma region test_nmr_provider - -#define TEST_PROVIDER_DISPATCH ((const void*)0x2) - -NPI_REGISTRATION_INSTANCE _test_provider_registration_instance = { - .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &test_npiid}; - -typedef struct -{ - bool allocated; - HANDLE nmr_binding_handle; - void* client_binding_context; - const void* client_dispatch; -} test_provider_binding_context_t; - -static test_provider_binding_context_t _test_provider_binding_context = {.allocated = false}; -static bool _test_provider_async_deregister = false; - -static NTSTATUS -_test_provider_attach_client( - _In_ HANDLE nmr_binding_handle, - _In_opt_ void* provider_context, - _In_ NPI_REGISTRATION_INSTANCE* client_registration_instance, - _In_ void* client_binding_context, - _In_ const void* client_dispatch, - _Outptr_ void** provider_binding_context, - _Outptr_ const void** provider_dispatch) -{ - // Verify not already bound. - REQUIRE(_test_provider_binding_context.allocated == false); - _test_provider_binding_context.allocated = true; - - _test_provider_binding_context.nmr_binding_handle = nmr_binding_handle; - _test_provider_binding_context.client_binding_context = client_binding_context; - _test_provider_binding_context.client_dispatch = client_dispatch; - - *provider_dispatch = TEST_PROVIDER_DISPATCH; - *provider_binding_context = &_test_provider_binding_context; - - return STATUS_SUCCESS; -} - -static NTSTATUS -_test_provider_detach_client(_In_ void* provider_binding_context) -{ - test_provider_binding_context_t* context = (test_provider_binding_context_t*)provider_binding_context; - REQUIRE(context->allocated); - - _test_provider_binding_context.client_binding_context = nullptr; - _test_provider_binding_context.client_dispatch = nullptr; - - return (_test_provider_async_deregister) ? STATUS_PENDING : STATUS_SUCCESS; -} - -static void -_test_provider_cleanup_binding_context(_In_ void* provider_binding_context) -{ - test_provider_binding_context_t* context = (test_provider_binding_context_t*)provider_binding_context; - REQUIRE(context->allocated); - context->allocated = false; - context->nmr_binding_handle = nullptr; -} - -NPI_PROVIDER_CHARACTERISTICS _test_provider_characteristics = { - .Length = sizeof(NPI_PROVIDER_CHARACTERISTICS), - .ProviderAttachClient = (PNPI_PROVIDER_ATTACH_CLIENT_FN)_test_provider_attach_client, - .ProviderDetachClient = _test_provider_detach_client, - .ProviderCleanupBindingContext = _test_provider_cleanup_binding_context, - .ProviderRegistrationInstance = _test_provider_registration_instance}; - -#pragma endregion test_nmr_provider - -#pragma region smoke_nmr_client_provider - -// Use a distinct NPI ID so smoke registrations never match the test_npiid registrations above, -// preventing cross-contamination when test_npiid registrations are left over from failed tests. -NPIID smoke_npiid = {1}; - -NPI_REGISTRATION_INSTANCE _smoke_client_registration_instance = { - .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &smoke_npiid}; -NPI_REGISTRATION_INSTANCE _smoke_provider_registration_instance = { - .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &smoke_npiid}; - -static NTSTATUS -_smoke_client_attach_provider( - _In_ HANDLE nmr_binding_handle, - _In_opt_ void* client_context, - _In_ NPI_REGISTRATION_INSTANCE* provider_registration_instance) -{ - UNREFERENCED_PARAMETER(client_context); - UNREFERENCED_PARAMETER(provider_registration_instance); - - void* provider_binding_context = nullptr; - const void* provider_dispatch = nullptr; - return NmrClientAttachProvider( - nmr_binding_handle, - reinterpret_cast(nmr_binding_handle), - TEST_CLIENT_DISPATCH, - &provider_binding_context, - &provider_dispatch); -} - -static NTSTATUS -_smoke_client_detach_provider(_In_ void* client_binding_context) -{ - UNREFERENCED_PARAMETER(client_binding_context); - return STATUS_SUCCESS; -} - -static void -_smoke_client_cleanup_binding_context(_In_ void* client_binding_context) -{ - UNREFERENCED_PARAMETER(client_binding_context); -} - -NPI_CLIENT_CHARACTERISTICS _smoke_client_characteristics = { - .Length = sizeof(NPI_CLIENT_CHARACTERISTICS), - .ClientAttachProvider = (PNPI_CLIENT_ATTACH_PROVIDER_FN)_smoke_client_attach_provider, - .ClientDetachProvider = _smoke_client_detach_provider, - .ClientCleanupBindingContext = _smoke_client_cleanup_binding_context, - .ClientRegistrationInstance = _smoke_client_registration_instance}; - -static NTSTATUS -_smoke_provider_attach_client( - _In_ HANDLE nmr_binding_handle, - _In_opt_ void* provider_context, - _In_ NPI_REGISTRATION_INSTANCE* client_registration_instance, - _In_ void* client_binding_context, - _In_ const void* client_dispatch, - _Outptr_ void** provider_binding_context, - _Outptr_ const void** provider_dispatch) -{ - UNREFERENCED_PARAMETER(nmr_binding_handle); - UNREFERENCED_PARAMETER(provider_context); - UNREFERENCED_PARAMETER(client_registration_instance); - UNREFERENCED_PARAMETER(client_dispatch); - - *provider_binding_context = client_binding_context; - *provider_dispatch = TEST_PROVIDER_DISPATCH; - return STATUS_SUCCESS; -} - -static NTSTATUS -_smoke_provider_detach_client(_In_ void* provider_binding_context) -{ - UNREFERENCED_PARAMETER(provider_binding_context); - return STATUS_SUCCESS; -} - -static void -_smoke_provider_cleanup_binding_context(_In_ void* provider_binding_context) -{ - UNREFERENCED_PARAMETER(provider_binding_context); -} - -NPI_PROVIDER_CHARACTERISTICS _smoke_provider_characteristics = { - .Length = sizeof(NPI_PROVIDER_CHARACTERISTICS), - .ProviderAttachClient = (PNPI_PROVIDER_ATTACH_CLIENT_FN)_smoke_provider_attach_client, - .ProviderDetachClient = _smoke_provider_detach_client, - .ProviderCleanupBindingContext = _smoke_provider_cleanup_binding_context, - .ProviderRegistrationInstance = _smoke_provider_registration_instance}; - -#pragma endregion smoke_nmr_client_provider - -TEST_CASE("NmrRegisterClient", "[nmr]") -{ - HANDLE nmr_client_handle; - - REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); - - // Verify there was no binding callback, since there are no providers. - REQUIRE(_test_client_binding_context.allocated == false); - - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); -} - -TEST_CASE("NmrRegisterProvider", "[nmr]") -{ - HANDLE nmr_provider_handle; - - REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); - - // Verify there was no binding callback, since there are no clients. - REQUIRE(_test_provider_binding_context.allocated == false); - - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); -} - -TEST_CASE("attach during NmrRegisterProvider", "[nmr]") -{ - HANDLE nmr_client_handle; - REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); - - // Verify there was no binding callback, since there are no providers. - REQUIRE(_test_client_binding_context.allocated == false); - - HANDLE nmr_provider_handle; - REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); - - REQUIRE(_test_client_binding_context.allocated == true); - REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_client_binding_context.provider_binding_context != nullptr); - REQUIRE(_test_client_binding_context.provider_dispatch == TEST_PROVIDER_DISPATCH); - - REQUIRE(_test_provider_binding_context.allocated == true); - REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_provider_binding_context.client_binding_context != nullptr); - REQUIRE(_test_provider_binding_context.client_dispatch == TEST_CLIENT_DISPATCH); - - // Deregister the provider first. - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); - - REQUIRE(_test_client_binding_context.allocated == false); - REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); - REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); - REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); - - REQUIRE(_test_provider_binding_context.allocated == false); - REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); - REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); - REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); - - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); -} - -TEST_CASE("attach during NmrRegisterClient", "[nmr]") -{ - HANDLE nmr_provider_handle; - REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); - - // Verify there was no binding callback, since there are no clients. - REQUIRE(_test_provider_binding_context.allocated == false); - - HANDLE nmr_client_handle; - REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); - - REQUIRE(_test_client_binding_context.allocated == true); - REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_client_binding_context.provider_binding_context != nullptr); - REQUIRE(_test_client_binding_context.provider_dispatch == TEST_PROVIDER_DISPATCH); - - REQUIRE(_test_provider_binding_context.allocated == true); - REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_provider_binding_context.client_binding_context != nullptr); - REQUIRE(_test_provider_binding_context.client_dispatch == TEST_CLIENT_DISPATCH); - - // Deregister the client first. - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); - - REQUIRE(_test_client_binding_context.allocated == false); - REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); - REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); - REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); - - REQUIRE(_test_provider_binding_context.allocated == false); - REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); - REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); - REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); - - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); -} - -TEST_CASE("NmrRegisterClient with async deregister", "[nmr]") -{ - HANDLE nmr_client_handle; - REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); - HANDLE nmr_provider_handle; - REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); - - // Start an asynchronous deregister, as if calls were in progress. - _test_client_async_deregister = true; - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_PENDING); - - // Verify that the binding still exists but no further calls will be initiated. - REQUIRE(_test_client_binding_context.allocated == true); - REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); - REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); - - REQUIRE(_test_provider_binding_context.allocated == true); - REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); - REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); - - // Complete the detach. - NmrClientDetachProviderComplete(_test_client_binding_context.nmr_binding_handle); - REQUIRE(NmrWaitForClientDeregisterComplete(nmr_client_handle) == STATUS_SUCCESS); - _test_client_async_deregister = false; - - // The binding should no longer exist. - REQUIRE(_test_client_binding_context.allocated == false); - REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); - REQUIRE(_test_provider_binding_context.allocated == false); - REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); - - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); -} - -TEST_CASE("NmrRegisterProvider with async deregister", "[nmr]") -{ - HANDLE nmr_client_handle; - REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); - HANDLE nmr_provider_handle; - REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); - - // Start an asynchronous deregister, as if calls were in progress. - _test_provider_async_deregister = true; - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_PENDING); - - // Verify that the binding still exists but no further calls will be initiated. - REQUIRE(_test_client_binding_context.allocated == true); - REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); - REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); - - REQUIRE(_test_provider_binding_context.allocated == true); - REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); - REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); - REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); - - // Complete the detach. - NmrProviderDetachClientComplete(_test_provider_binding_context.nmr_binding_handle); - REQUIRE(NmrWaitForProviderDeregisterComplete(nmr_provider_handle) == STATUS_SUCCESS); - _test_provider_async_deregister = false; - - // The binding should no longer exist. - REQUIRE(_test_client_binding_context.allocated == false); - REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); - REQUIRE(_test_provider_binding_context.allocated == false); - REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); - - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); -} - -TEST_CASE("concurrent register/deregister smoke", "[nmr][no_fi]") -{ - constexpr size_t iteration_count = 1000; - std::atomic provider_iterations{0}; - std::atomic client_iterations{0}; - - std::thread provider_thread([iteration_count, &provider_iterations]() { - for (size_t i = 0; i < iteration_count; i++) { - provider_iterations++; - HANDLE nmr_provider_handle = nullptr; - NTSTATUS register_status = - NmrRegisterProvider(&_smoke_provider_characteristics, nullptr, &nmr_provider_handle); - if (!NT_SUCCESS(register_status)) { - continue; - } - - NTSTATUS deregister_status = NmrDeregisterProvider(nmr_provider_handle); - if (deregister_status == STATUS_PENDING) { - (void)NmrWaitForProviderDeregisterComplete(nmr_provider_handle); - } - } - }); - - std::thread client_thread([iteration_count, &client_iterations]() { - for (size_t i = 0; i < iteration_count; i++) { - client_iterations++; - HANDLE nmr_client_handle = nullptr; - NTSTATUS register_status = NmrRegisterClient(&_smoke_client_characteristics, nullptr, &nmr_client_handle); - if (!NT_SUCCESS(register_status)) { - continue; - } - - NTSTATUS deregister_status = NmrDeregisterClient(nmr_client_handle); - if (deregister_status == STATUS_PENDING) { - (void)NmrWaitForClientDeregisterComplete(nmr_client_handle); - } - } - }); - - provider_thread.join(); - client_thread.join(); - - std::cout << "provider_iterations=" << provider_iterations.load() << ", client_iterations=" - << client_iterations.load() << std::endl; +// Copyright (c) Microsoft Corporation +// SPDX-License-Identifier: MIT + +#if !defined(CMAKE_NUGET) +#include +#else +#include +#endif +#include "../src/framework.h" +#include <../km/netioddk.h> +#include +#include +#include +#include + +NPIID test_npiid = {0}; + +#pragma region test_nmr_client + +#define TEST_CLIENT_DISPATCH ((const void*)0x1) + +NPI_REGISTRATION_INSTANCE _test_client_registration_instance = { + .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &test_npiid}; + +typedef struct +{ + bool allocated; + HANDLE nmr_binding_handle; + void* provider_binding_context; + const void* provider_dispatch; +} test_client_binding_context_t; + +static test_client_binding_context_t _test_client_binding_context = {.allocated = false}; +static bool _test_client_async_deregister = false; + +static NTSTATUS +_test_client_attach_provider( + _In_ HANDLE nmr_binding_handle, + _In_opt_ void* client_context, + _In_ NPI_REGISTRATION_INSTANCE* provider_registration_instance) +{ + // Verify not already bound. + REQUIRE(_test_client_binding_context.allocated == false); + _test_client_binding_context.allocated = true; + + NTSTATUS status = NmrClientAttachProvider( + nmr_binding_handle, + &_test_client_binding_context, + TEST_CLIENT_DISPATCH, + &_test_client_binding_context.provider_binding_context, + &_test_client_binding_context.provider_dispatch); + + if (NT_SUCCESS(status)) { + _test_client_binding_context.nmr_binding_handle = nmr_binding_handle; + } + + return status; +} + +static NTSTATUS +_test_client_detach_provider(_In_ void* client_binding_context) +{ + test_client_binding_context_t* context = (test_client_binding_context_t*)client_binding_context; + REQUIRE(context->allocated); + + _test_client_binding_context.provider_binding_context = nullptr; + _test_client_binding_context.provider_dispatch = nullptr; + + return (_test_client_async_deregister) ? STATUS_PENDING : STATUS_SUCCESS; +} + +static void +_test_client_cleanup_binding_context(_In_ void* client_binding_context) +{ + test_client_binding_context_t* context = (test_client_binding_context_t*)client_binding_context; + REQUIRE(context->allocated); + context->allocated = false; + context->nmr_binding_handle = nullptr; +} + +NPI_CLIENT_CHARACTERISTICS _test_client_characteristics = { + .Length = sizeof(NPI_CLIENT_CHARACTERISTICS), + .ClientAttachProvider = (PNPI_CLIENT_ATTACH_PROVIDER_FN)_test_client_attach_provider, + .ClientDetachProvider = _test_client_detach_provider, + .ClientCleanupBindingContext = _test_client_cleanup_binding_context, + .ClientRegistrationInstance = _test_client_registration_instance}; + +#pragma endregion test_nmr_client +#pragma region test_nmr_provider + +#define TEST_PROVIDER_DISPATCH ((const void*)0x2) + +NPI_REGISTRATION_INSTANCE _test_provider_registration_instance = { + .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &test_npiid}; + +typedef struct +{ + bool allocated; + HANDLE nmr_binding_handle; + void* client_binding_context; + const void* client_dispatch; +} test_provider_binding_context_t; + +static test_provider_binding_context_t _test_provider_binding_context = {.allocated = false}; +static bool _test_provider_async_deregister = false; + +static NTSTATUS +_test_provider_attach_client( + _In_ HANDLE nmr_binding_handle, + _In_opt_ void* provider_context, + _In_ NPI_REGISTRATION_INSTANCE* client_registration_instance, + _In_ void* client_binding_context, + _In_ const void* client_dispatch, + _Outptr_ void** provider_binding_context, + _Outptr_ const void** provider_dispatch) +{ + // Verify not already bound. + REQUIRE(_test_provider_binding_context.allocated == false); + _test_provider_binding_context.allocated = true; + + _test_provider_binding_context.nmr_binding_handle = nmr_binding_handle; + _test_provider_binding_context.client_binding_context = client_binding_context; + _test_provider_binding_context.client_dispatch = client_dispatch; + + *provider_dispatch = TEST_PROVIDER_DISPATCH; + *provider_binding_context = &_test_provider_binding_context; + + return STATUS_SUCCESS; +} + +static NTSTATUS +_test_provider_detach_client(_In_ void* provider_binding_context) +{ + test_provider_binding_context_t* context = (test_provider_binding_context_t*)provider_binding_context; + REQUIRE(context->allocated); + + _test_provider_binding_context.client_binding_context = nullptr; + _test_provider_binding_context.client_dispatch = nullptr; + + return (_test_provider_async_deregister) ? STATUS_PENDING : STATUS_SUCCESS; +} + +static void +_test_provider_cleanup_binding_context(_In_ void* provider_binding_context) +{ + test_provider_binding_context_t* context = (test_provider_binding_context_t*)provider_binding_context; + REQUIRE(context->allocated); + context->allocated = false; + context->nmr_binding_handle = nullptr; +} + +NPI_PROVIDER_CHARACTERISTICS _test_provider_characteristics = { + .Length = sizeof(NPI_PROVIDER_CHARACTERISTICS), + .ProviderAttachClient = (PNPI_PROVIDER_ATTACH_CLIENT_FN)_test_provider_attach_client, + .ProviderDetachClient = _test_provider_detach_client, + .ProviderCleanupBindingContext = _test_provider_cleanup_binding_context, + .ProviderRegistrationInstance = _test_provider_registration_instance}; + +#pragma endregion test_nmr_provider + +#pragma region smoke_nmr_client_provider + +// Use a distinct NPI ID so smoke registrations never match the test_npiid registrations above, +// preventing cross-contamination when test_npiid registrations are left over from failed tests. +NPIID smoke_npiid = {1}; + +NPI_REGISTRATION_INSTANCE _smoke_client_registration_instance = { + .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &smoke_npiid}; +NPI_REGISTRATION_INSTANCE _smoke_provider_registration_instance = { + .Size = sizeof(NPI_REGISTRATION_INSTANCE), .NpiId = &smoke_npiid}; + +static NTSTATUS +_smoke_client_attach_provider( + _In_ HANDLE nmr_binding_handle, + _In_opt_ void* client_context, + _In_ NPI_REGISTRATION_INSTANCE* provider_registration_instance) +{ + UNREFERENCED_PARAMETER(client_context); + UNREFERENCED_PARAMETER(provider_registration_instance); + + void* provider_binding_context = nullptr; + const void* provider_dispatch = nullptr; + return NmrClientAttachProvider( + nmr_binding_handle, + reinterpret_cast(nmr_binding_handle), + TEST_CLIENT_DISPATCH, + &provider_binding_context, + &provider_dispatch); +} + +static NTSTATUS +_smoke_client_detach_provider(_In_ void* client_binding_context) +{ + UNREFERENCED_PARAMETER(client_binding_context); + return STATUS_SUCCESS; +} + +static void +_smoke_client_cleanup_binding_context(_In_ void* client_binding_context) +{ + UNREFERENCED_PARAMETER(client_binding_context); +} + +NPI_CLIENT_CHARACTERISTICS _smoke_client_characteristics = { + .Length = sizeof(NPI_CLIENT_CHARACTERISTICS), + .ClientAttachProvider = (PNPI_CLIENT_ATTACH_PROVIDER_FN)_smoke_client_attach_provider, + .ClientDetachProvider = _smoke_client_detach_provider, + .ClientCleanupBindingContext = _smoke_client_cleanup_binding_context, + .ClientRegistrationInstance = _smoke_client_registration_instance}; + +static NTSTATUS +_smoke_provider_attach_client( + _In_ HANDLE nmr_binding_handle, + _In_opt_ void* provider_context, + _In_ NPI_REGISTRATION_INSTANCE* client_registration_instance, + _In_ void* client_binding_context, + _In_ const void* client_dispatch, + _Outptr_ void** provider_binding_context, + _Outptr_ const void** provider_dispatch) +{ + UNREFERENCED_PARAMETER(nmr_binding_handle); + UNREFERENCED_PARAMETER(provider_context); + UNREFERENCED_PARAMETER(client_registration_instance); + UNREFERENCED_PARAMETER(client_dispatch); + + *provider_binding_context = client_binding_context; + *provider_dispatch = TEST_PROVIDER_DISPATCH; + return STATUS_SUCCESS; +} + +static NTSTATUS +_smoke_provider_detach_client(_In_ void* provider_binding_context) +{ + UNREFERENCED_PARAMETER(provider_binding_context); + return STATUS_SUCCESS; +} + +static void +_smoke_provider_cleanup_binding_context(_In_ void* provider_binding_context) +{ + UNREFERENCED_PARAMETER(provider_binding_context); +} + +NPI_PROVIDER_CHARACTERISTICS _smoke_provider_characteristics = { + .Length = sizeof(NPI_PROVIDER_CHARACTERISTICS), + .ProviderAttachClient = (PNPI_PROVIDER_ATTACH_CLIENT_FN)_smoke_provider_attach_client, + .ProviderDetachClient = _smoke_provider_detach_client, + .ProviderCleanupBindingContext = _smoke_provider_cleanup_binding_context, + .ProviderRegistrationInstance = _smoke_provider_registration_instance}; + +#pragma endregion smoke_nmr_client_provider + +TEST_CASE("NmrRegisterClient", "[nmr]") +{ + HANDLE nmr_client_handle; + + REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); + + // Verify there was no binding callback, since there are no providers. + REQUIRE(_test_client_binding_context.allocated == false); + + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); +} + +TEST_CASE("NmrRegisterProvider", "[nmr]") +{ + HANDLE nmr_provider_handle; + + REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); + + // Verify there was no binding callback, since there are no clients. + REQUIRE(_test_provider_binding_context.allocated == false); + + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); +} + +TEST_CASE("attach during NmrRegisterProvider", "[nmr]") +{ + HANDLE nmr_client_handle; + REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); + + // Verify there was no binding callback, since there are no providers. + REQUIRE(_test_client_binding_context.allocated == false); + + HANDLE nmr_provider_handle; + REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); + + REQUIRE(_test_client_binding_context.allocated == true); + REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_client_binding_context.provider_binding_context != nullptr); + REQUIRE(_test_client_binding_context.provider_dispatch == TEST_PROVIDER_DISPATCH); + + REQUIRE(_test_provider_binding_context.allocated == true); + REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_provider_binding_context.client_binding_context != nullptr); + REQUIRE(_test_provider_binding_context.client_dispatch == TEST_CLIENT_DISPATCH); + + // Deregister the provider first. + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); + + REQUIRE(_test_client_binding_context.allocated == false); + REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); + REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); + REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); + + REQUIRE(_test_provider_binding_context.allocated == false); + REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); + REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); + REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); + + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); +} + +TEST_CASE("attach during NmrRegisterClient", "[nmr]") +{ + HANDLE nmr_provider_handle; + REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); + + // Verify there was no binding callback, since there are no clients. + REQUIRE(_test_provider_binding_context.allocated == false); + + HANDLE nmr_client_handle; + REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); + + REQUIRE(_test_client_binding_context.allocated == true); + REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_client_binding_context.provider_binding_context != nullptr); + REQUIRE(_test_client_binding_context.provider_dispatch == TEST_PROVIDER_DISPATCH); + + REQUIRE(_test_provider_binding_context.allocated == true); + REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_provider_binding_context.client_binding_context != nullptr); + REQUIRE(_test_provider_binding_context.client_dispatch == TEST_CLIENT_DISPATCH); + + // Deregister the client first. + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); + + REQUIRE(_test_client_binding_context.allocated == false); + REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); + REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); + REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); + + REQUIRE(_test_provider_binding_context.allocated == false); + REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); + REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); + REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); + + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); +} + +TEST_CASE("NmrRegisterClient with async deregister", "[nmr]") +{ + HANDLE nmr_client_handle; + REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); + HANDLE nmr_provider_handle; + REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); + + // Start an asynchronous deregister, as if calls were in progress. + _test_client_async_deregister = true; + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_PENDING); + + // Verify that the binding still exists but no further calls will be initiated. + REQUIRE(_test_client_binding_context.allocated == true); + REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); + REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); + + REQUIRE(_test_provider_binding_context.allocated == true); + REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); + REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); + + // Complete the detach. + NmrClientDetachProviderComplete(_test_client_binding_context.nmr_binding_handle); + REQUIRE(NmrWaitForClientDeregisterComplete(nmr_client_handle) == STATUS_SUCCESS); + _test_client_async_deregister = false; + + // The binding should no longer exist. + REQUIRE(_test_client_binding_context.allocated == false); + REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); + REQUIRE(_test_provider_binding_context.allocated == false); + REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); + + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); +} + +TEST_CASE("NmrRegisterProvider with async deregister", "[nmr]") +{ + HANDLE nmr_client_handle; + REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); + HANDLE nmr_provider_handle; + REQUIRE(NmrRegisterProvider(&_test_provider_characteristics, nullptr, &nmr_provider_handle) == STATUS_SUCCESS); + + // Start an asynchronous deregister, as if calls were in progress. + _test_provider_async_deregister = true; + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_PENDING); + + // Verify that the binding still exists but no further calls will be initiated. + REQUIRE(_test_client_binding_context.allocated == true); + REQUIRE(_test_client_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_client_binding_context.provider_binding_context == nullptr); + REQUIRE(_test_client_binding_context.provider_dispatch == nullptr); + + REQUIRE(_test_provider_binding_context.allocated == true); + REQUIRE(_test_provider_binding_context.nmr_binding_handle != nullptr); + REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); + REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); + + // Complete the detach. + NmrProviderDetachClientComplete(_test_provider_binding_context.nmr_binding_handle); + REQUIRE(NmrWaitForProviderDeregisterComplete(nmr_provider_handle) == STATUS_SUCCESS); + _test_provider_async_deregister = false; + + // The binding should no longer exist. + REQUIRE(_test_client_binding_context.allocated == false); + REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); + REQUIRE(_test_provider_binding_context.allocated == false); + REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); + + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); +} + +TEST_CASE("concurrent register/deregister smoke", "[nmr][no_fi]") +{ + constexpr auto run_duration = std::chrono::seconds(30); + const auto test_start = std::chrono::steady_clock::now(); + std::atomic ready_threads{0}; + std::atomic start{false}; + + std::thread provider_thread([&ready_threads, &start, run_duration]() { + ready_threads++; + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + const auto deadline = std::chrono::steady_clock::now() + run_duration; + while (std::chrono::steady_clock::now() < deadline) { + HANDLE nmr_provider_handle = nullptr; + NTSTATUS register_status = + NmrRegisterProvider(&_smoke_provider_characteristics, nullptr, &nmr_provider_handle); + if (!NT_SUCCESS(register_status)) { + continue; + } + + NTSTATUS deregister_status = NmrDeregisterProvider(nmr_provider_handle); + if (deregister_status == STATUS_PENDING) { + (void)NmrWaitForProviderDeregisterComplete(nmr_provider_handle); + } + } + }); + + std::thread client_thread([&ready_threads, &start, run_duration]() { + ready_threads++; + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + const auto deadline = std::chrono::steady_clock::now() + run_duration; + while (std::chrono::steady_clock::now() < deadline) { + HANDLE nmr_client_handle = nullptr; + NTSTATUS register_status = NmrRegisterClient(&_smoke_client_characteristics, nullptr, &nmr_client_handle); + if (!NT_SUCCESS(register_status)) { + continue; + } + + NTSTATUS deregister_status = NmrDeregisterClient(nmr_client_handle); + if (deregister_status == STATUS_PENDING) { + (void)NmrWaitForClientDeregisterComplete(nmr_client_handle); + } + } + }); + + while (ready_threads.load(std::memory_order_acquire) != 2) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + + provider_thread.join(); + client_thread.join(); + + const auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - test_start); + std::cout << "smoke test duration: " << elapsed.count() << " ms" << std::endl; } \ No newline at end of file From 858208763bf5bcce3af93323db36ba5ae3b50548 Mon Sep 17 00:00:00 2001 From: "Shankar Seal (from Dev Box)" Date: Mon, 10 Aug 2026 20:20:12 -0700 Subject: [PATCH 3/3] Mock NMR impl changes. --- src/nmr_impl.cpp | 822 +++++++++++++++++++++++---------------------- src/nmr_impl.h | 62 ++-- src/nmr_um.cpp | 14 +- tests/nmr_test.cpp | 31 +- 4 files changed, 474 insertions(+), 455 deletions(-) diff --git a/src/nmr_impl.cpp b/src/nmr_impl.cpp index 66caa95..e09e4d8 100644 --- a/src/nmr_impl.cpp +++ b/src/nmr_impl.cpp @@ -1,402 +1,420 @@ -// Copyright (c) Microsoft Corporation -// SPDX-License-Identifier: MIT - -#include "nmr_impl.h" - -#define NMR_WAIT_TIMEOUT_SECONDS 10 - -nmr_t::nmr_provider_handle -nmr_t::register_provider(_In_ const NPI_PROVIDER_CHARACTERISTICS& characteristics, _In_opt_ const void* context) -{ - // Add the provider to the list of providers. - nmr_provider_handle provider_handle = add(providers, characteristics, context); - // Notify existing clients about the new provider. - perform_bind(providers, provider_handle, clients); - return provider_handle; -} - -bool -nmr_t::deregister_provider(_In_ nmr_provider_handle provider_handle) -{ - // Block new bindings. - deactivate(providers, provider_handle); - - // If the unbind returned pending, then the caller needs to wait for the unbind to complete. - if (perform_unbind(providers, provider_handle)) { - // Pending unbind. - return true; - } - // Unbind is complete. - remove(providers, provider_handle); - return false; -} - -void -nmr_t::wait_for_deregister_provider(_In_ nmr_provider_handle provider_handle) -{ - // Wait for the unbind to complete. - remove(providers, provider_handle); -} - -nmr_t::nmr_client_handle -nmr_t::register_client(_In_ const NPI_CLIENT_CHARACTERISTICS& characteristics, _In_opt_ const void* context) -{ - // Add the client to the list of clients. - nmr_client_handle client_handle = add(clients, characteristics, context); - // Notify existing providers about the new client. - perform_bind(clients, client_handle, providers); - return client_handle; -} - -bool -nmr_t::deregister_client(_In_ nmr_client_handle client_handle) -{ - // Block new bindings. - deactivate(clients, client_handle); - - // If the unbind returned pending, then the caller needs to wait for the unbind to complete. - if (perform_unbind(clients, client_handle)) { - // Pending unbind. - return true; - } - - // Unbind is complete. - remove(clients, client_handle); - return false; -} - -void -nmr_t::wait_for_deregister_client(_In_ nmr_client_handle client_handle) -{ - // Wait for the unbind to complete. - remove(clients, client_handle); -} - -void -nmr_t::binding_detach_client_complete(_In_ nmr_binding_handle binding_handle) -{ - std::unique_lock l(lock); - auto it = bindings.find(binding_handle); - if (it == bindings.end()) { - throw std::runtime_error("invalid handle"); - } - - nmr_t::binding& binding = *it->second; - - ASSERT(binding.client_binding_status == binding_status::UnbindPending); - binding.client_binding_status = UnbindComplete; - bool complete = (binding.provider_binding_status == binding_status::UnbindComplete); - l.unlock(); - if (complete) { - // Signal the detach complete. - unbind_complete(binding); - } -} - -void -nmr_t::binding_detach_provider_complete(_In_ nmr_binding_handle binding_handle) -{ - std::unique_lock l(lock); - auto it = bindings.find(binding_handle); - if (it == bindings.end()) { - throw std::runtime_error("invalid handle"); - } - - nmr_t::binding& binding = *it->second; - - ASSERT(binding.provider_binding_status == binding_status::UnbindPending); - binding.provider_binding_status = UnbindComplete; - bool complete = (binding.client_binding_status == binding_status::UnbindComplete); - l.unlock(); - if (complete) { - // Signal the detach complete. - unbind_complete(binding); - } -} - -NTSTATUS -nmr_t::client_attach_provider( - _In_ nmr_binding_handle binding_handle, - _In_ __drv_aliasesMem const void* client_binding_context, - _In_ const void* client_dispatch, - _Outptr_ const void** provider_binding_context, - _Outptr_ const void** provider_dispatch) -{ - std::unique_lock l(lock); - // Resolve the binding_handle to the binding. - auto it = bindings.find(binding_handle); - if (it == bindings.end()) { - throw std::runtime_error("invalid handle"); - } - auto& binding = *it->second; - - // Save the client's per binding context and dispatch table. - binding.client_binding_context = client_binding_context; - binding.client_dispatch = client_dispatch; - l.unlock(); - - // Call the provider's attach client. - NTSTATUS status = binding.provider.characteristics.ProviderAttachClient( - const_cast(binding_handle), - const_cast(binding.provider.context), - &binding.client.characteristics.ClientRegistrationInstance, - const_cast(client_binding_context), - client_dispatch, - const_cast(&binding.provider_binding_context), - &binding.provider_dispatch); - - // If successful, save the provider's per binding context and dispatch table. - if (NT_SUCCESS(status)) { - *provider_binding_context = binding.provider_binding_context; - *provider_dispatch = binding.provider_dispatch; - } - return status; -} - -// Assumes caller does NOT have the lock held since we call outside NMR. -std::optional -nmr_t::bind(_Inout_ client_registration& client, _Inout_ provider_registration& provider) -{ - PNPIID client_pnpi = client.characteristics.ClientRegistrationInstance.NpiId; - PNPIID provider_pnpi = provider.characteristics.ProviderRegistrationInstance.NpiId; - - // Match on NPI ID. - if (!client_pnpi || !provider_pnpi || *client_pnpi != *provider_pnpi) { - return std::nullopt; - } - - // Skip if client or provider are deregistering. - if (client.deregistering || provider.deregistering) { - return std::nullopt; - } - - // Acquire references on both client and provider to prevent them from unloading. - _InterlockedIncrement64(&client.binding_count); - _InterlockedIncrement64(&provider.binding_count); - - nmr_t::binding binding = {provider, client}; - auto binding_ptr = std::make_shared(std::move(binding)); - - bindings.insert({binding_ptr.get(), binding_ptr}); - - return {[&client, &provider, binding_ptr, this]() { - NTSTATUS status = client.characteristics.ClientAttachProvider( - reinterpret_cast(binding_ptr.get()), - const_cast(client.context), - &provider.characteristics.ProviderRegistrationInstance); - - // Clean up the binding on a failure. - if (!NT_SUCCESS(status)) { - unbind_complete(*binding_ptr); - } else { - bool should_begin_unbind = false; - std::unique_lock l(lock); - binding_ptr->client_binding_status = binding_status::Ready; - binding_ptr->provider_binding_status = binding_status::Ready; - should_begin_unbind = binding_ptr->client.deregistering || binding_ptr->provider.deregistering; - l.unlock(); - if (should_begin_unbind) { - (void)begin_unbind(*binding_ptr); - } - } - }}; -} - -void -nmr_t::unbind_complete(_Inout_ binding& binding) -{ - std::unique_lock l(lock); - if ((binding.client.characteristics.ClientCleanupBindingContext != nullptr) && - (binding.client_binding_context != nullptr)) { - // Notify the client that that the binding context can be freed if needed. - binding.client.characteristics.ClientCleanupBindingContext(const_cast(binding.client_binding_context)); - } - - if ((binding.provider.characteristics.ProviderCleanupBindingContext != nullptr) && - (binding.provider_binding_context != nullptr)) { - // Notify the provider that that the binding context can be freed if needed. - binding.provider.characteristics.ProviderCleanupBindingContext( - const_cast(binding.provider_binding_context)); - } - - _InterlockedDecrement64(&binding.provider.binding_count); - _InterlockedDecrement64(&binding.client.binding_count); - bindings.erase(&binding); - - // Notify the client or provider to check if they have any pending bindings. - bindings_changed.notify_all(); -} - -bool -nmr_t::begin_unbind(_Inout_ binding& binding) -{ - std::unique_lock l(lock); - if (binding.client_binding_status != Ready || binding.provider_binding_status != Ready) { - // A Start binding is already published and contributes to binding_count, so deregistration - // must keep waiting even though detach cannot begin until attach finishes and reaches Ready. - return true; - } - binding.client_binding_status = BeginUnbind; - binding.provider_binding_status = BeginUnbind; - l.unlock(); - - NTSTATUS client_detach_provider_status = - (binding.client_binding_context) - ? binding.client.characteristics.ClientDetachProvider(const_cast(binding.client_binding_context)) - : STATUS_SUCCESS; - NTSTATUS provider_detach_client_status = - (binding.provider_binding_context) - ? binding.provider.characteristics.ProviderDetachClient(const_cast(binding.provider_binding_context)) - : STATUS_SUCCESS; - - // Take a lock to make sure we don't replace UnbindComplete with UnbindPending, - // since if one of the above returned pending, the completion could get called - // sequently which would change the status on the binding any time before we - // grab a lock. - l.lock(); - binding.provider_binding_status = (client_detach_provider_status == STATUS_PENDING && - binding.provider_binding_status != binding_status::UnbindComplete) - ? binding_status::UnbindPending - : binding_status::UnbindComplete; - binding.client_binding_status = (provider_detach_client_status == STATUS_PENDING && - binding.client_binding_status != binding_status::UnbindComplete) - ? binding_status::UnbindPending - : binding_status::UnbindComplete; - bool complete = - ((binding.client_binding_status == binding_status::UnbindComplete) && - (binding.provider_binding_status == binding_status::UnbindComplete)); - l.unlock(); - - if (complete) { - unbind_complete(binding); - return false; - } - return true; -} - -template -collection_t::value_type::first_type -nmr_t::add(_Inout_ collection_t& collection, _In_ const characteristics_t& characteristics, _In_ const void* context) -{ - std::unique_lock l(lock); - auto handle = reinterpret_cast(next_handle++); - collection.insert({handle, {characteristics, context}}); - return handle; -} - -template -void -nmr_t::deactivate(_Inout_ collection_t& collection, _Inout_ collection_t::value_type::first_type handle) -{ - std::unique_lock l(lock); - auto it = collection.find(handle); - if (it == collection.end()) { - throw std::runtime_error("invalid handle"); - } - - // Block new bindings. - it->second.deregistering = true; -} - -template -void -nmr_t::remove(_Inout_ collection_t& collection, _In_ collection_t::value_type::first_type handle) -{ - std::unique_lock l(lock); - auto it = collection.find(handle); - if (it == collection.end()) { - throw std::runtime_error("invalid handle"); - } - - // Wait for bindings to reach zero if requested. - if (it->second.binding_count > 0) { - for (;;) { - // Wait NMR_WAIT_TIMEOUT_SECONDS seconds for bindings to reach zero. - if (bindings_changed.wait_for(l, std::chrono::seconds(NMR_WAIT_TIMEOUT_SECONDS), [&]() { - return it->second.binding_count == 0; - })) { - break; - } - // Assert and continue waiting if bindings are still not zero. - CXPLAT_DEBUG_ASSERT(it->second.binding_count == 0); - } - } - - collection.erase(it); -} - -template -void -nmr_t::perform_bind( - _Inout_ initiator_collection_t& initiator_collection, - _In_ initiator_collection_t::value_type::first_type initiator_handle, - _Inout_ target_collection_t& target_collection) -{ - // Queue up the bind for each target to performed outside the lock. - std::vector pending_actions; - std::unique_lock l(lock); - auto it = initiator_collection.find(initiator_handle); - if (it == initiator_collection.end()) { - throw std::runtime_error("invalid handle"); - } - auto& initiator = it->second; - for (auto& [target_handle, target] : target_collection) { - // If the initiator is a client, then the target must be a provider. - if constexpr (std::is_same::value) { - auto result = bind(initiator, target); - if (result.has_value()) { - pending_actions.push_back(result.value()); - } - } - // If the initiator is a provider, then the target must be a client. - if constexpr (std::is_same::value) { - auto result = bind(target, initiator); - if (result.has_value()) { - pending_actions.push_back(result.value()); - } - } - } - l.unlock(); - for (auto& action : pending_actions) { - action(); - } -} - -template -bool -nmr_t::perform_unbind( - _Inout_ initiator_collection_t& initiator_collection, - _In_ initiator_collection_t::value_type::first_type initiator_handle) -{ - bool pending = false; - std::vector> bindings_to_unbind; - std::unique_lock l(lock); - auto it = initiator_collection.find(initiator_handle); - if (it == initiator_collection.end()) { - throw std::runtime_error("invalid handle"); - } - auto& initiator = it->second; - // Find all the bindings that have the initiator as the client or provider. - for (auto& [binding_handle, binding_reference] : bindings) { - auto& binding = *binding_reference; - - // If the initiator is a client, then the target must be a provider. - if constexpr (std::is_same::value) { - if (&binding.client == &initiator) { - bindings_to_unbind.push_back(binding_reference); - } - } - // If the initiator is a provider, then the target must be a client. - if constexpr (std::is_same::value) { - if (&binding.provider == &initiator) { - bindings_to_unbind.push_back(binding_reference); - } - } - } - l.unlock(); - for (auto& binding_reference : bindings_to_unbind) { - auto& binding = *binding_reference; - pending |= begin_unbind(binding); - } - return pending; -} +// Copyright (c) Microsoft Corporation +// SPDX-License-Identifier: MIT + +#include "nmr_impl.h" + +#define NMR_WAIT_TIMEOUT_SECONDS 10 + +nmr_t::nmr_provider_handle +nmr_t::register_provider(_In_ const NPI_PROVIDER_CHARACTERISTICS& characteristics, _In_opt_ const void* context) +{ + // Add the provider to the list of providers. + nmr_provider_handle provider_handle = add(providers, characteristics, context); + // Notify existing clients about the new provider. + perform_bind(providers, provider_handle, clients); + return provider_handle; +} + +void +nmr_t::deregister_provider(_In_ nmr_provider_handle provider_handle) +{ + // Block new bindings. + deactivate(providers, provider_handle); + + // The caller always waits for deregistration to complete. + perform_unbind(providers, provider_handle); +} + +void +nmr_t::wait_for_deregister_provider(_In_ nmr_provider_handle provider_handle) +{ + // Wait for the unbind to complete. + remove(providers, provider_handle); +} + +nmr_t::nmr_client_handle +nmr_t::register_client(_In_ const NPI_CLIENT_CHARACTERISTICS& characteristics, _In_opt_ const void* context) +{ + // Add the client to the list of clients. + nmr_client_handle client_handle = add(clients, characteristics, context); + // Notify existing providers about the new client. + perform_bind(clients, client_handle, providers); + return client_handle; +} + +void +nmr_t::deregister_client(_In_ nmr_client_handle client_handle) +{ + // Block new bindings. + deactivate(clients, client_handle); + + // The caller always waits for deregistration to complete. + perform_unbind(clients, client_handle); +} + +void +nmr_t::wait_for_deregister_client(_In_ nmr_client_handle client_handle) +{ + // Wait for the unbind to complete. + remove(clients, client_handle); +} + +void +nmr_t::binding_detach_client_complete(_In_ nmr_binding_handle binding_handle) +{ + std::unique_lock l(lock); + auto it = bindings.find(binding_handle); + if (it == bindings.end()) { + throw std::runtime_error("invalid handle"); + } + + auto binding_reference = it->second; + nmr_t::binding& binding = *binding_reference; + + ASSERT(binding.client_binding_status == binding_status::UnbindPending); + binding.client_binding_status = UnbindComplete; + bool complete = (binding.provider_binding_status == binding_status::UnbindComplete); + l.unlock(); + if (complete) { + // Signal the detach complete. + unbind_complete(binding); + } +} + +void +nmr_t::binding_detach_provider_complete(_In_ nmr_binding_handle binding_handle) +{ + std::unique_lock l(lock); + auto it = bindings.find(binding_handle); + if (it == bindings.end()) { + throw std::runtime_error("invalid handle"); + } + + auto binding_reference = it->second; + nmr_t::binding& binding = *binding_reference; + + ASSERT(binding.provider_binding_status == binding_status::UnbindPending); + binding.provider_binding_status = UnbindComplete; + bool complete = (binding.client_binding_status == binding_status::UnbindComplete); + l.unlock(); + if (complete) { + // Signal the detach complete. + unbind_complete(binding); + } +} + +NTSTATUS +nmr_t::client_attach_provider( + _In_ nmr_binding_handle binding_handle, + _In_ __drv_aliasesMem const void* client_binding_context, + _In_ const void* client_dispatch, + _Outptr_ const void** provider_binding_context, + _Outptr_ const void** provider_dispatch) +{ + std::unique_lock l(lock); + // Resolve the binding_handle to the binding. + auto it = bindings.find(binding_handle); + if (it == bindings.end()) { + throw std::runtime_error("invalid handle"); + } + auto& binding = *it->second; + + // Save the client's per binding context and dispatch table. + binding.client_binding_context = client_binding_context; + binding.client_dispatch = client_dispatch; + l.unlock(); + + // Call the provider's attach client. + NTSTATUS status = binding.provider.characteristics.ProviderAttachClient( + const_cast(binding_handle), + const_cast(binding.provider.context), + &binding.client.characteristics.ClientRegistrationInstance, + const_cast(client_binding_context), + client_dispatch, + const_cast(&binding.provider_binding_context), + &binding.provider_dispatch); + + // If successful, save the provider's per binding context and dispatch table. + if (NT_SUCCESS(status)) { + *provider_binding_context = binding.provider_binding_context; + *provider_dispatch = binding.provider_dispatch; + } + return status; +} + +// Assumes caller does NOT have the lock held since we call outside NMR. +std::optional +nmr_t::bind(_Inout_ client_module& client, _Inout_ provider_module& provider) +{ + PNPIID client_pnpi = client.characteristics.ClientRegistrationInstance.NpiId; + PNPIID provider_pnpi = provider.characteristics.ProviderRegistrationInstance.NpiId; + + // Match on NPI ID. + if (!client_pnpi || !provider_pnpi || *client_pnpi != *provider_pnpi) { + return std::nullopt; + } + + // Skip if client or provider are deregistering. + if (client.deregistering || provider.deregistering) { + return std::nullopt; + } + + // Hold both modules while the attach action is pending. + client.pending_bind_ops++; + provider.pending_bind_ops++; + + nmr_t::binding binding = {provider, client}; + auto binding_ptr = std::make_shared(std::move(binding)); + + bindings.insert({binding_ptr.get(), binding_ptr}); + + return {[&client, &provider, binding_ptr, this]() { + NTSTATUS status = client.characteristics.ClientAttachProvider( + reinterpret_cast(binding_ptr.get()), + const_cast(client.context), + &provider.characteristics.ProviderRegistrationInstance); + + // Clean up the binding on a failure. + if (!NT_SUCCESS(status)) { + unbind_complete(*binding_ptr); + } else { + bool should_begin_unbind = false; + std::unique_lock l(lock); + binding_ptr->client_binding_status = binding_status::Ready; + binding_ptr->provider_binding_status = binding_status::Ready; + binding_ptr->attached = true; + client.pending_bind_ops--; + provider.pending_bind_ops--; + client.bindings.push_back(binding_ptr); + provider.bindings.push_back(binding_ptr); + should_begin_unbind = binding_ptr->client.deregistering || binding_ptr->provider.deregistering; + if (should_begin_unbind) { + binding_ptr->client_binding_status = binding_status::LateBind; + binding_ptr->provider_binding_status = binding_status::LateBind; + } + l.unlock(); + if (should_begin_unbind) { + begin_unbind(*binding_ptr); + } + } + }}; +} + +void +nmr_t::unbind_complete(_Inout_ binding& binding) +{ + std::unique_lock l(lock); + if (binding.cleanup_started) { + return; + } + binding.cleanup_started = true; + + if ((binding.client.characteristics.ClientCleanupBindingContext != nullptr) && + (binding.client_binding_context != nullptr)) { + // Notify the client that that the binding context can be freed if needed. + binding.client.characteristics.ClientCleanupBindingContext(const_cast(binding.client_binding_context)); + } + + if ((binding.provider.characteristics.ProviderCleanupBindingContext != nullptr) && + (binding.provider_binding_context != nullptr)) { + // Notify the provider that that the binding context can be freed if needed. + binding.provider.characteristics.ProviderCleanupBindingContext( + const_cast(binding.provider_binding_context)); + } + + if (binding.attached) { + auto remove_binding = [&binding](auto& module) { + auto it = std::find_if(module.bindings.begin(), module.bindings.end(), [&binding](const auto& entry) { + return entry.get() == &binding; + }); + if (it != module.bindings.end()) { + module.bindings.erase(it); + } + }; + remove_binding(binding.provider); + remove_binding(binding.client); + } else { + CXPLAT_DEBUG_ASSERT(binding.client.pending_bind_ops > 0); + CXPLAT_DEBUG_ASSERT(binding.provider.pending_bind_ops > 0); + binding.client.pending_bind_ops--; + binding.provider.pending_bind_ops--; + } + + bindings.erase(&binding); + + // Notify the client or provider to check if they have any pending bindings. + bindings_changed.notify_all(); +} + +void +nmr_t::begin_unbind(_Inout_ binding& binding) +{ + std::unique_lock l(lock); + const bool client_ready = binding.client_binding_status == Ready || binding.client_binding_status == LateBind; + const bool provider_ready = binding.provider_binding_status == Ready || binding.provider_binding_status == LateBind; + if (!client_ready || !provider_ready) { + // A non-ready binding is either still being attached or already being unbound. + // In either case, another caller must not start a second unbind operation. + return; + } + binding.client_binding_status = BeginUnbind; + binding.provider_binding_status = BeginUnbind; + l.unlock(); + + NTSTATUS client_detach_provider_status = + (binding.client_binding_context) + ? binding.client.characteristics.ClientDetachProvider(const_cast(binding.client_binding_context)) + : STATUS_SUCCESS; + NTSTATUS provider_detach_client_status = + (binding.provider_binding_context) + ? binding.provider.characteristics.ProviderDetachClient(const_cast(binding.provider_binding_context)) + : STATUS_SUCCESS; + + // Take a lock to make sure we don't replace UnbindComplete with UnbindPending, + // since if one of the above returned pending, the completion could get called + // sequently which would change the status on the binding any time before we + // grab a lock. + l.lock(); + binding.provider_binding_status = (client_detach_provider_status == STATUS_PENDING && + binding.provider_binding_status != binding_status::UnbindComplete) + ? binding_status::UnbindPending + : binding_status::UnbindComplete; + binding.client_binding_status = (provider_detach_client_status == STATUS_PENDING && + binding.client_binding_status != binding_status::UnbindComplete) + ? binding_status::UnbindPending + : binding_status::UnbindComplete; + bool complete = + ((binding.client_binding_status == binding_status::UnbindComplete) && + (binding.provider_binding_status == binding_status::UnbindComplete)); + l.unlock(); + + if (complete) { + unbind_complete(binding); + } +} + +template +collection_t::value_type::first_type +nmr_t::add(_Inout_ collection_t& collection, _In_ const characteristics_t& characteristics, _In_ const void* context) +{ + std::unique_lock l(lock); + auto handle = reinterpret_cast(next_handle++); + collection.insert({handle, {characteristics, context}}); + return handle; +} + +template +void +nmr_t::deactivate(_Inout_ collection_t& collection, _Inout_ collection_t::value_type::first_type handle) +{ + std::unique_lock l(lock); + auto it = collection.find(handle); + if (it == collection.end()) { + throw std::runtime_error("invalid handle"); + } + + // Block new bindings. + it->second.deregistering = true; +} + +template +void +nmr_t::remove(_Inout_ collection_t& collection, _In_ collection_t::value_type::first_type handle) +{ + std::unique_lock l(lock); + auto it = collection.find(handle); + if (it == collection.end()) { + throw std::runtime_error("invalid handle"); + } + + // Wait until there are no attached bindings and no bind operations still pending + // before cleaning up the client or provider module. + if (it->second.pending_bind_ops > 0 || !it->second.bindings.empty()) { + for (;;) { + // Wait NMR_WAIT_TIMEOUT_SECONDS seconds for all bindings to be released. + if (bindings_changed.wait_for(l, std::chrono::seconds(NMR_WAIT_TIMEOUT_SECONDS), [&]() { + return it->second.pending_bind_ops == 0 && it->second.bindings.empty(); + })) { + break; + } + // Assert and continue waiting if bindings are still not zero. + CXPLAT_DEBUG_ASSERT(it->second.pending_bind_ops == 0 && it->second.bindings.empty()); + } + } + + collection.erase(it); +} + +template +void +nmr_t::perform_bind( + _Inout_ initiator_collection_t& initiator_collection, + _In_ initiator_collection_t::value_type::first_type initiator_handle, + _Inout_ target_collection_t& target_collection) +{ + // Queue up the bind for each target to performed outside the lock. + std::vector pending_actions; + std::unique_lock l(lock); + auto it = initiator_collection.find(initiator_handle); + if (it == initiator_collection.end()) { + throw std::runtime_error("invalid handle"); + } + auto& initiator = it->second; + for (auto& [target_handle, target] : target_collection) { + // If the initiator is a client, then the target must be a provider. + if constexpr (std::is_same::value) { + auto result = bind(initiator, target); + if (result.has_value()) { + pending_actions.push_back(result.value()); + } + } + // If the initiator is a provider, then the target must be a client. + if constexpr (std::is_same::value) { + auto result = bind(target, initiator); + if (result.has_value()) { + pending_actions.push_back(result.value()); + } + } + } + l.unlock(); + for (auto& action : pending_actions) { + action(); + } +} + +template +void +nmr_t::perform_unbind( + _Inout_ initiator_collection_t& initiator_collection, + _In_ initiator_collection_t::value_type::first_type initiator_handle) +{ + std::vector> bindings_to_unbind; + std::unique_lock l(lock); + auto it = initiator_collection.find(initiator_handle); + if (it == initiator_collection.end()) { + throw std::runtime_error("invalid handle"); + } + auto& initiator = it->second; + // Find all the bindings that have the initiator as the client or provider. + for (auto& [binding_handle, binding_reference] : bindings) { + auto& binding = *binding_reference; + + // If the initiator is a client, then the target must be a provider. + if constexpr (std::is_same::value) { + if (&binding.client == &initiator) { + bindings_to_unbind.push_back(binding_reference); + } + } + // If the initiator is a provider, then the target must be a client. + if constexpr (std::is_same::value) { + if (&binding.provider == &initiator) { + bindings_to_unbind.push_back(binding_reference); + } + } + } + l.unlock(); + for (auto& binding_reference : bindings_to_unbind) { + auto& binding = *binding_reference; + begin_unbind(binding); + } +} diff --git a/src/nmr_impl.h b/src/nmr_impl.h index c10bffd..5b19345 100644 --- a/src/nmr_impl.h +++ b/src/nmr_impl.h @@ -5,6 +5,7 @@ #include "platform.h" #include <../km/netioddk.h> +#include #include #include #include @@ -28,19 +29,17 @@ typedef class nmr_t * * @param[in] characteristics Characteristics of the provider. * @param[in] context Context passed to the provider. - * @return Handle to the provider registration. + * @return Handle to the provider module. */ nmr_provider_handle register_provider(_In_ const NPI_PROVIDER_CHARACTERISTICS& characteristics, _In_opt_ const void* context); /** - * @brief Deregister a provider. + * @brief Deregister a provider. The caller must wait for completion. * * @param[in] provider_handle Handle to the provider. - * @retval true Caller needs to wait for the deregistration to complete. - * @retval false Deregistration is complete. */ - bool + void deregister_provider(_In_ nmr_provider_handle provider_handle); /** @@ -56,19 +55,17 @@ typedef class nmr_t * * @param[in] characteristics Characteristics of the client. * @param[in] context Context passed to the client. - * @return Handle to the client registration. + * @return Handle to the client module. */ nmr_client_handle register_client(_In_ const NPI_CLIENT_CHARACTERISTICS& characteristics, _In_opt_ const void* context); /** - * @brief Deregister a client. + * @brief Deregister a client. The caller must wait for completion. * * @param[in] client_handle Handle to the client. - * @retval true Caller needs to wait for the deregistration to complete. - * @retval false Deregistration is complete. */ - bool + void deregister_client(_In_ nmr_client_handle client_handle); /** @@ -122,19 +119,23 @@ typedef class nmr_t } private: - struct client_registration + struct binding; + + struct client_module { const NPI_CLIENT_CHARACTERISTICS characteristics = {}; const void* context = nullptr; - volatile long long binding_count = 0; + size_t pending_bind_ops = 0; + std::vector> bindings; bool deregistering = false; }; - struct provider_registration + struct provider_module { const NPI_PROVIDER_CHARACTERISTICS characteristics = {}; const void* context = nullptr; - volatile long long binding_count = 0; + size_t pending_bind_ops = 0; + std::vector> bindings; bool deregistering = false; }; @@ -142,6 +143,7 @@ typedef class nmr_t { Start = 0, ///< Initial state. Binding has been created but ClientAttachProvider has not been called. Ready, ///< ClientAttachProvider has been called and returned STATUS_SUCCESS. + LateBind, ///< Attach completed after one of the modules started deregistering. BeginUnbind, ///< Client or provider has called NmrDeregisterClient or NmrDeregisterProvider but detach has not ///< yet been called. UnbindPending, ///< Client or provider detach returned STATUS_PENDING. @@ -150,14 +152,16 @@ typedef class nmr_t }; struct binding { - provider_registration& provider; - client_registration& client; + provider_module& provider; + client_module& client; const void* provider_binding_context = nullptr; const void* provider_dispatch = nullptr; binding_status provider_binding_status = Start; const void* client_binding_context = nullptr; const void* client_dispatch = nullptr; binding_status client_binding_status = Start; + bool attached = false; + bool cleanup_started = false; }; typedef std::function pending_action_t; @@ -218,11 +222,9 @@ typedef class nmr_t * * @param[in, out] initiator_collection Collection containing the initiator (can be either provider or client). * @param[in] handle Handle to the initiator (can be either provider or client). - * @retval true One or more bindings returned pending. - * @retval false All bindings where successfully removed. */ template - bool + void perform_unbind( _Inout_ initiator_collection_t& initiator_collection, _In_ initiator_collection_t::value_type::first_type initiator_handle); @@ -232,10 +234,10 @@ typedef class nmr_t * * @param[in, out] client Client to attempt to bind. * @param[in, out] provider Provider to attempt to bind to. - * @return Contains a function to perform the bind if successful. + * @return Contains a function to perform the bind if accepted. */ std::optional - bind(_Inout_ client_registration& client, _Inout_ provider_registration& provider); + bind(_Inout_ client_module& client, _Inout_ provider_module& provider); /** * @brief Finish the process of unbinding a client from a provider. @@ -249,26 +251,22 @@ typedef class nmr_t * @brief Start the process of unbinding a client from a provider. * * @param[in] binding_handle Binding handle to unbind. - * @retval true Either unbind cannot start yet or it is pending/in progress. - * @retval false Both the client and provider returned successfully. */ - bool + void begin_unbind(_Inout_ binding& binding); // Binding handle is a pointer to the binding. std::map> bindings; - // Provider and client handles are incremented for each new provider or client. - std::map providers; - std::map clients; + // Provider and client handles. + std::map providers; + std::map clients; size_t next_handle = 1; std::condition_variable bindings_changed; - std::mutex lock; // Protects all of the instance variables above, - // as well as the client_binding_status and provider_binding_status - // of each binding. client.binding_count and provider.binding_count - // on the other hand are not protected by this lock but instead use - // interlocked operations. + std::mutex lock; // Protects all of the instance variables above, as well as + // the client_binding_status and provider_binding_status of + // each binding. static nmr_t singleton; } nmr_t; diff --git a/src/nmr_um.cpp b/src/nmr_um.cpp index ff8d029..a282ad3 100644 --- a/src/nmr_um.cpp +++ b/src/nmr_um.cpp @@ -28,11 +28,8 @@ NTSTATUS NmrDeregisterProvider(_In_ HANDLE nmr_provider_handle) { try { - if (nmr_t::get().deregister_provider(nmr_provider_handle)) { - return STATUS_PENDING; - } else { - return STATUS_SUCCESS; - } + nmr_t::get().deregister_provider(nmr_provider_handle); + return STATUS_PENDING; } catch (std::bad_alloc) { return STATUS_NO_MEMORY; } @@ -81,11 +78,8 @@ NTSTATUS NmrDeregisterClient(_In_ HANDLE nmr_client_handle) { try { - if (nmr_t::get().deregister_client(nmr_client_handle)) { - return STATUS_PENDING; - } else { - return STATUS_SUCCESS; - } + nmr_t::get().deregister_client(nmr_client_handle); + return STATUS_PENDING; } catch (std::bad_alloc) { return STATUS_NO_MEMORY; } diff --git a/tests/nmr_test.cpp b/tests/nmr_test.cpp index 26370a8..79d7d41 100644 --- a/tests/nmr_test.cpp +++ b/tests/nmr_test.cpp @@ -7,6 +7,7 @@ #include #endif #include "../src/framework.h" + #include <../km/netioddk.h> #include #include @@ -253,13 +254,14 @@ NPI_PROVIDER_CHARACTERISTICS _smoke_provider_characteristics = { TEST_CASE("NmrRegisterClient", "[nmr]") { HANDLE nmr_client_handle; - + REQUIRE(NmrRegisterClient(&_test_client_characteristics, nullptr, &nmr_client_handle) == STATUS_SUCCESS); // Verify there was no binding callback, since there are no providers. REQUIRE(_test_client_binding_context.allocated == false); - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(nmr_client_handle) == STATUS_SUCCESS); } TEST_CASE("NmrRegisterProvider", "[nmr]") @@ -271,7 +273,8 @@ TEST_CASE("NmrRegisterProvider", "[nmr]") // Verify there was no binding callback, since there are no clients. REQUIRE(_test_provider_binding_context.allocated == false); - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(nmr_provider_handle) == STATUS_SUCCESS); } TEST_CASE("attach during NmrRegisterProvider", "[nmr]") @@ -296,7 +299,8 @@ TEST_CASE("attach during NmrRegisterProvider", "[nmr]") REQUIRE(_test_provider_binding_context.client_dispatch == TEST_CLIENT_DISPATCH); // Deregister the provider first. - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(nmr_provider_handle) == STATUS_SUCCESS); REQUIRE(_test_client_binding_context.allocated == false); REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); @@ -308,7 +312,8 @@ TEST_CASE("attach during NmrRegisterProvider", "[nmr]") REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(nmr_client_handle) == STATUS_SUCCESS); } TEST_CASE("attach during NmrRegisterClient", "[nmr]") @@ -333,7 +338,8 @@ TEST_CASE("attach during NmrRegisterClient", "[nmr]") REQUIRE(_test_provider_binding_context.client_dispatch == TEST_CLIENT_DISPATCH); // Deregister the client first. - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(nmr_client_handle) == STATUS_SUCCESS); REQUIRE(_test_client_binding_context.allocated == false); REQUIRE(_test_client_binding_context.nmr_binding_handle == nullptr); @@ -345,7 +351,8 @@ TEST_CASE("attach during NmrRegisterClient", "[nmr]") REQUIRE(_test_provider_binding_context.client_binding_context == nullptr); REQUIRE(_test_provider_binding_context.client_dispatch == nullptr); - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(nmr_provider_handle) == STATUS_SUCCESS); } TEST_CASE("NmrRegisterClient with async deregister", "[nmr]") @@ -381,7 +388,8 @@ TEST_CASE("NmrRegisterClient with async deregister", "[nmr]") REQUIRE(_test_provider_binding_context.allocated == false); REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); - REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterProvider(nmr_provider_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(nmr_provider_handle) == STATUS_SUCCESS); } TEST_CASE("NmrRegisterProvider with async deregister", "[nmr]") @@ -417,7 +425,8 @@ TEST_CASE("NmrRegisterProvider with async deregister", "[nmr]") REQUIRE(_test_provider_binding_context.allocated == false); REQUIRE(_test_provider_binding_context.nmr_binding_handle == nullptr); - REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_SUCCESS); + REQUIRE(NmrDeregisterClient(nmr_client_handle) == STATUS_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(nmr_client_handle) == STATUS_SUCCESS); } TEST_CASE("concurrent register/deregister smoke", "[nmr][no_fi]") @@ -476,7 +485,7 @@ TEST_CASE("concurrent register/deregister smoke", "[nmr][no_fi]") provider_thread.join(); client_thread.join(); - const auto elapsed = std::chrono::duration_cast( - std::chrono::steady_clock::now() - test_start); + const auto elapsed = + std::chrono::duration_cast(std::chrono::steady_clock::now() - test_start); std::cout << "smoke test duration: " << elapsed.count() << " ms" << std::endl; } \ No newline at end of file