Fix NetworkState submit/cleanup TOCTOU races

Source/Global/NetworkState.cpp

@@ -151,6 +151,13 @@ bool NetworkState::CanCleanupCancelHttpRequest(XAsyncBlock* async) noexcept
     }
     return false;
 }
+
+void NetworkState::TestSetCleanupStarted(bool started, XAsyncBlock* cleanupAsyncBlock) noexcept
+{
+    std::unique_lock<std::mutex> lock{ m_mutex };
+    m_cleanupStarted = started;
+    m_cleanupAsyncBlock = cleanupAsyncBlock;
+}
 #endif
 
 HRESULT CALLBACK NetworkState::HttpCallPerformAsyncProvider(XAsyncOp op, const XAsyncProviderData* data)
@@ -168,6 +175,15 @@ HRESULT CALLBACK NetworkState::HttpCallPerformAsyncProvider(XAsyncOp op, const X
         RETURN_IF_FAILED(XTaskQueueCreateComposite(workPort, workPort, &performContext->internalAsyncBlock.queue));
 
         std::unique_lock<std::mutex> lock{ state.m_mutex };
+        if (state.m_cleanupStarted)
+        {
+            // Cleanup has already begun and taken (or is taking) its snapshot of
+            // m_activeHttpRequests under this same mutex. Refuse the request instead of inserting
+            // it after the snapshot, which would orphan it (never canceled or awaited) and could
+            // run it against a torn-down provider.
+            lock.unlock();
+            return E_HC_NOT_INITIALISED;
+        }
         state.m_activeHttpRequests.insert(performContext);
         lock.unlock();
 
@@ -181,8 +197,11 @@ HRESULT CALLBACK NetworkState::HttpCallPerformAsyncProvider(XAsyncOp op, const X
     case XAsyncOp::Cleanup:
     {
         std::unique_lock<std::mutex> lock{ state.m_mutex };
-        state.m_activeHttpRequests.erase(performContext);
-        bool scheduleCleanup = state.ScheduleCleanup();
+        // Only a perform that was actually admitted (present in m_activeHttpRequests) may drive the
+        // cleanup wakeup. A perform rejected by the m_cleanupStarted guard was never inserted, so
+        // erase() returns 0 and we must not call ScheduleCleanup()/XAsyncSchedule for it -- doing so
+        // would spuriously (re)schedule cleanup's async block for a request that was never tracked.
+        bool scheduleCleanup = state.m_activeHttpRequests.erase(performContext) != 0 && state.ScheduleCleanup();
         lock.unlock();
 
         // Free performContext before scheduling cleanup to ensure it happens before returing to client
@@ -370,6 +389,13 @@ HRESULT CALLBACK NetworkState::WebSocketConnectAsyncProvider(XAsyncOp op, const
         RETURN_IF_FAILED(XTaskQueueCreateComposite(workPort, workPort, &context->internalAsyncBlock.queue));
 
         std::unique_lock<std::mutex> lock{ state.m_mutex };
+        if (state.m_cleanupStarted)
+        {
+            // See the equivalent guard in HttpCallPerformAsyncProvider: reject connects that arrive
+            // after cleanup has begun rather than orphaning them past the cleanup snapshot.
+            lock.unlock();
+            return E_HC_NOT_INITIALISED;
+        }
         state.m_connectingWebSockets.insert(context->clientAsyncBlock);
         lock.unlock();
 
@@ -471,11 +497,14 @@ void CALLBACK NetworkState::WebSocketClosed(HCWebsocketHandle /*websocket*/, HCW
 }
 #endif // !HC_NOWEBSOCKETS
 
-HRESULT NetworkState::CleanupAsync(UniquePtr<NetworkState> state, XAsyncBlock* async) noexcept
+HRESULT NetworkState::CleanupAsync(NetworkState* state, XAsyncBlock* async) noexcept
 {
-    RETURN_IF_FAILED(XAsyncBegin(async, state.get(), __FUNCTION__, __FUNCTION__, CleanupAsyncProvider));
-    state.release();
-    return S_OK;
+    // NetworkState is not taken by owning pointer here: it remains owned by the http_singleton for
+    // its whole lifetime. Cleanup runs against the still-owned instance and never destroys it, so an
+    // in-flight API caller holding a singleton reference can never observe a moved-from or destroyed
+    // NetworkState (Race B). The instance is destroyed together with the singleton, once the
+    // singleton's use_count gate confirms no other references remain.
+    return XAsyncBegin(async, state, __FUNCTION__, __FUNCTION__, CleanupAsyncProvider);
 }
 
 HRESULT CALLBACK NetworkState::CleanupAsyncProvider(XAsyncOp op, const XAsyncProviderData* data)
@@ -495,17 +524,19 @@ HRESULT CALLBACK NetworkState::CleanupAsyncProvider(XAsyncOp op, const XAsyncPro
         {
             std::unique_lock<std::mutex> lock{ state->m_mutex };
             state->m_cleanupAsyncBlock = data->async;
+            state->m_cleanupStarted = true;
             scheduleCleanup = state->ScheduleCleanup();
 
 #ifndef HC_NOWEBSOCKETS 
             HC_TRACE_VERBOSE(HTTPCLIENT, "NetworkState::CleanupAsyncProvider::Begin: HTTP active=%llu, WebSocket Connecting=%llu, WebSocket Connected=%llu", state->m_activeHttpRequests.size(), state->m_connectingWebSockets.size(), state->m_connectedWebSockets.size());
 #endif
-            // No new HTTP performs can enter m_activeHttpRequests after cleanup begins because
-            // http_singleton::singleton_access(cleanup) detaches the singleton before
-            // NetworkState::CleanupAsync runs. Snapshot requests here, then cancel them after
-            // releasing m_mutex. This prevents a race between holding the global cleanup mutex
-            // across XAsyncCancel and allowing completion to advance a request that cleanup has
-            // already decided to cancel.
+            // Setting m_cleanupStarted above (under m_mutex) closes the admission window: any HTTP
+            // perform or WebSocket connect whose Begin op acquires m_mutex after this point is
+            // refused, so nothing can be inserted into the tracking sets after the snapshot below.
+            // Requests that acquired m_mutex before us are already in m_activeHttpRequests and are
+            // captured by the snapshot here. Snapshot them under the lock and cancel them after
+            // releasing m_mutex; this prevents holding the lock across XAsyncCancel while still
+            // ensuring completion cannot advance a request cleanup has already decided to cancel.
             for (auto activeRequest : state->m_activeHttpRequests)
             {
                 auto expectedState = HttpPerformClientBlockState::CleanupMayCancel;
@@ -566,8 +597,7 @@ HRESULT CALLBACK NetworkState::CleanupAsyncProvider(XAsyncOp op, const XAsyncPro
         {
             XAsyncBlock* cleanupAsyncBlock{ state->m_cleanupAsyncBlock };
 
-            UniquePtr<NetworkState> reclaim{ state };
-            reclaim.reset();
+            // NetworkState stays owned by the http_singleton; do not destroy it here.
             providerCleanupAsyncBlock.reset();
 
             XAsyncComplete(cleanupAsyncBlock, hr, 0);
@@ -589,14 +619,15 @@ HRESULT CALLBACK NetworkState::CleanupAsyncProvider(XAsyncOp op, const XAsyncPro
 void CALLBACK NetworkState::HttpProviderCleanupComplete(XAsyncBlock* async)
 {
     UniquePtr<XAsyncBlock> providerCleanupAsyncBlock{ async };
-    UniquePtr<NetworkState> state{ static_cast<NetworkState*>(providerCleanupAsyncBlock->context) };
+    NetworkState* state{ static_cast<NetworkState*>(providerCleanupAsyncBlock->context) };
     XAsyncBlock* stateCleanupAsyncBlock = state->m_cleanupAsyncBlock;
 
     HRESULT cleanupResult = XAsyncGetStatus(providerCleanupAsyncBlock.get(), false);
     providerCleanupAsyncBlock.reset();
-    state.reset();
 
-    // NetworkState fully cleaned up at this point
+    // NetworkState's providers are cleaned up at this point. The NetworkState instance itself stays
+    // owned by the http_singleton and is destroyed when the singleton is (after its use_count gate),
+    // so an in-flight caller holding a singleton reference never observes it freed.
     XAsyncComplete(stateCleanupAsyncBlock, cleanupResult, 0);
 }
 

Source/Global/NetworkState.h

@@ -34,7 +34,7 @@ class NetworkState
 #endif
 
     static HRESULT CleanupAsync(
-        UniquePtr<NetworkState> networkManager,
+        NetworkState* networkState,
         XAsyncBlock* async
     ) noexcept;
 
@@ -50,6 +50,9 @@ class NetworkState
 
 #ifdef HC_UNITTEST_API
     bool CanCleanupCancelHttpRequest(XAsyncBlock* async) noexcept;
+    // Test seam: simulate that cleanup has begun on this NetworkState without running the
+    // real cleanup flow, so admission-control behavior can be exercised deterministically.
+    void TestSetCleanupStarted(bool started, XAsyncBlock* cleanupAsyncBlock = nullptr) noexcept;
 #endif
 
 #ifndef HC_NOWEBSOCKETS
@@ -94,6 +97,14 @@ class NetworkState
     Set<HttpPerformContext*> m_activeHttpRequests;
     XAsyncBlock* m_cleanupAsyncBlock{ nullptr }; // non-owning
 
+    // Admission-control gate for new network operations. Set (under m_mutex) once cleanup has
+    // begun; while set, new HTTP performs and WebSocket connects are refused so they cannot land
+    // in the tracking sets after the cleanup snapshot has been taken.
+    // All reads and writes must occur under m_mutex; that is what makes a plain bool sufficient
+    // (no atomic needed) and what makes the guard atomic with the tracking-set insert/snapshot. Do
+    // not read this outside m_mutex.
+    bool m_cleanupStarted{ false };
+
 #ifndef HC_NOWEBSOCKETS
     static HRESULT CALLBACK WebSocketConnectAsyncProvider(XAsyncOp op, const XAsyncProviderData* data);
     static void CALLBACK WebSocketConnectComplete(XAsyncBlock* async);

Source/Global/global.cpp

@@ -156,7 +156,7 @@ HRESULT CALLBACK http_singleton::CleanupAsyncProvider(XAsyncOp op, const XAsyncP
             XAsyncSchedule(singletonCleanupAsyncBlock, 0);
         };
 
-        RETURN_IF_FAILED(NetworkState::CleanupAsync(std::move(singleton->m_networkState), performEnvCleanupAsyncBlock.get()));
+        RETURN_IF_FAILED(NetworkState::CleanupAsync(singleton->m_networkState.get(), performEnvCleanupAsyncBlock.get()));
         performEnvCleanupAsyncBlock.release();
 
         return S_OK;
@@ -169,6 +169,19 @@ HRESULT CALLBACK http_singleton::CleanupAsyncProvider(XAsyncOp op, const XAsyncP
         // Note that the use count check here is only valid because we never create
         // a weak_ptr to the singleton. If we did that could cause the use count
         // to increase even though we are the only strong reference
+        //
+        // This gate is also what keeps NetworkState safe to own from the singleton for its whole
+        // lifetime (rather than moving it out during cleanup Begin). An in-flight API caller that
+        // obtained a strong reference via get_http_singleton() before cleanup detached the singleton
+        // keeps use_count above 1 here, so neither the singleton nor the NetworkState it owns is
+        // destroyed until that caller releases its reference. That is what prevents an in-flight
+        // HCHttpCallPerformAsync / HCWebSocketConnectAsync from observing a moved-from or destroyed
+        // m_networkState.
+        //
+        // INVARIANT: no code may hold a get_http_singleton() reference across an async wait or other
+        // blocking operation. Doing so would keep use_count above 1 indefinitely and stall cleanup
+        // here. Every caller today takes the reference as a short-lived local scoped to a single
+        // synchronous operation.
         if (self.use_count() > 1)
         {
             RETURN_IF_FAILED(XAsyncSchedule(data->async, 10));

Tests/UnitTests/Tests/GlobalTests.cpp

@@ -196,6 +196,106 @@ DEFINE_TEST_CLASS(GlobalTests)
         VERIFY_SUCCEEDED(HCHttpCallCloseHandle(call));
         HCCleanup();
     }
+
+    // Race B reproduction: an in-flight API caller (e.g. HCHttpCallPerformAsync) takes a strong
+    // singleton reference via get_http_singleton() and has not yet dereferenced m_networkState.
+    // Cleanup running concurrently must not detach/destroy NetworkState out from under that
+    // reference, otherwise the caller null-derefs the moved-from m_networkState.
+    DEFINE_TEST_CASE(TestCleanupKeepsNetworkStateForInFlightSingletonRef)
+    {
+        VERIFY_SUCCEEDED(HCInitialize(nullptr));
+        PumpedTaskQueue pumpedQueue;
+
+        auto inFlightSingletonRef = get_http_singleton();
+        VERIFY_IS_NOT_NULL(inFlightSingletonRef.get());
+
+        XAsyncBlock cleanupAsyncBlock{ pumpedQueue.queue };
+        VERIFY_SUCCEEDED(HCCleanupAsync(&cleanupAsyncBlock));
+        // XAsyncBegin dispatches the cleanup Begin op synchronously on this thread, so by the time
+        // HCCleanupAsync returns the singleton has been detached. A pre-fix build has already
+        // std::move'd m_networkState out from under our still-live reference at this point.
+        bool networkStateStillValid = (inFlightSingletonRef->m_networkState.get() != nullptr);
+
+        // Release our reference and drain cleanup to completion BEFORE asserting, so a failing
+        // assertion never unwinds the test with an in-flight cleanup still referencing the stack
+        // async block.
+        inFlightSingletonRef.reset();
+        VERIFY_SUCCEEDED(XAsyncGetStatus(&cleanupAsyncBlock, true));
+
+        VERIFY_IS_TRUE(networkStateStillValid);
+    }
+
+    // Race A reproduction: once cleanup has begun on NetworkState, a perform whose Begin op runs
+    // afterwards must be rejected rather than inserted into m_activeHttpRequests. Otherwise it is
+    // orphaned past the cleanup snapshot (never canceled/awaited) and can run against a torn-down
+    // provider.
+    DEFINE_TEST_CASE(TestHttpPerformRejectedAfterCleanupStarted)
+    {
+        VERIFY_SUCCEEDED(HCInitialize(nullptr));
+
+        XTaskQueueHandle queue{ nullptr };
+        VERIFY_SUCCEEDED(XTaskQueueCreate(XTaskQueueDispatchMode::Manual, XTaskQueueDispatchMode::Manual, &queue));
+
+        auto cleanupProvider = [](XAsyncOp op, const XAsyncProviderData* data)
+        {
+            switch (op)
+            {
+            case XAsyncOp::Begin:
+            {
+                return S_OK;
+            }
+            case XAsyncOp::DoWork:
+            {
+                XAsyncComplete(data->async, S_OK, 0);
+                return E_PENDING;
+            }
+            default:
+            {
+                return S_OK;
+            }
+            }
+        };
+
+        XAsyncBlock cleanupAsyncBlock{ queue };
+        VERIFY_SUCCEEDED(XAsyncBegin(&cleanupAsyncBlock, nullptr, nullptr, nullptr, cleanupProvider));
+
+        constexpr char mockUrl[]{ "www.bing.com" };
+        HCMockCallHandle mock{ nullptr };
+        VERIFY_SUCCEEDED(HCMockCallCreate(&mock));
+        VERIFY_SUCCEEDED(HCMockResponseSetStatusCode(mock, 200));
+        VERIFY_SUCCEEDED(HCMockAddMock(mock, "GET", mockUrl, nullptr, 0));
+
+        HCCallHandle call{ nullptr };
+        VERIFY_SUCCEEDED(HCHttpCallCreate(&call));
+        VERIFY_SUCCEEDED(HCHttpCallRequestSetUrl(call, "GET", mockUrl));
+
+        auto httpSingleton = get_http_singleton();
+        VERIFY_IS_NOT_NULL(httpSingleton.get());
+
+        // Simulate cleanup having begun on NetworkState after its tracking-set snapshot. The
+        // cleanup async block is intentionally real and unscheduled so a rejected perform must not
+        // consume cleanup's one allowed schedule.
+        httpSingleton->m_networkState->TestSetCleanupStarted(true, &cleanupAsyncBlock);
+
+        XAsyncBlock performAsyncBlock{ queue };
+        VERIFY_SUCCEEDED(httpSingleton->m_networkState->HttpCallPerformAsync(call, &performAsyncBlock));
+
+        HRESULT performStatus = XAsyncGetStatus(&performAsyncBlock, true);
+        VERIFY_ARE_EQUAL(E_HC_NOT_INITIALISED, performStatus);
+        VERIFY_IS_FALSE(httpSingleton->m_networkState->CanCleanupCancelHttpRequest(&performAsyncBlock));
+
+        HRESULT cleanupScheduleHr = XAsyncSchedule(&cleanupAsyncBlock, 0);
+        VERIFY_IS_TRUE(XTaskQueueDispatch(queue, XTaskQueuePort::Work, 0));
+        VERIFY_SUCCEEDED(XAsyncGetStatus(&cleanupAsyncBlock, true));
+        VERIFY_SUCCEEDED(cleanupScheduleHr);
+
+        httpSingleton->m_networkState->TestSetCleanupStarted(false);
+        httpSingleton.reset();
+
+        VERIFY_SUCCEEDED(HCHttpCallCloseHandle(call));
+        HCCleanup();
+        XTaskQueueCloseHandle(queue);
+    }
 };
 
 NAMESPACE_XBOX_HTTP_CLIENT_TEST_END