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..e09e4d8 100644 --- a/src/nmr_impl.cpp +++ b/src/nmr_impl.cpp @@ -1,395 +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 { - std::unique_lock l(lock); - binding_ptr->client_binding_status = binding_status::Ready; - binding_ptr->provider_binding_status = binding_status::Ready; - } - }}; -} - -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 // true if pending, false if complete. -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. - 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 // true if pending, false if complete -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 a8b7ea7..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,23 +143,25 @@ 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. UnbindComplete ///< Client or provider detach returned STATUS_SUCCESS or called NmrBindingDetachClientComplete ///< or NmrBindingDetachProviderComplete. }; - 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; @@ -219,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); @@ -233,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. @@ -250,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 the client or provider returned pending. - * @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 16a8d04..79d7d41 100644 --- a/tests/nmr_test.cpp +++ b/tests/nmr_test.cpp @@ -1,325 +1,491 @@ -// Copyright (c) Microsoft Corporation -// SPDX-License-Identifier: MIT - -#if !defined(CMAKE_NUGET) -#include -#else -#include -#endif -#include "../src/framework.h" -#include <../km/netioddk.h> - -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 - -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); +// 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_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForProviderDeregisterComplete(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_PENDING); + REQUIRE(NmrWaitForClientDeregisterComplete(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