diff --git a/Source/WebSocket/hcwebsocket.cpp b/Source/WebSocket/hcwebsocket.cpp index f496528c..9674c443 100644 --- a/Source/WebSocket/hcwebsocket.cpp +++ b/Source/WebSocket/hcwebsocket.cpp @@ -205,12 +205,13 @@ HRESULT WebSocket::ConnectAsync( HRESULT CALLBACK WebSocket::ConnectAsyncProvider(XAsyncOp op, XAsyncProviderData const* data) { ConnectContext* context{ static_cast(data->context) }; - auto& ws{ context->observer->websocket }; switch (op) { case XAsyncOp::Begin: { + assert(context->observer); + auto& ws{ context->observer->websocket }; std::unique_lock lock{ ws->m_stateMutex }; RETURN_HR_IF(E_UNEXPECTED, ws->m_state != State::Initial); @@ -255,11 +256,19 @@ void CALLBACK WebSocket::ConnectComplete(XAsyncBlock* async) ConnectContext* context{ static_cast(async->context) }; auto& ws{ context->observer->websocket }; - assert(ws->m_state == State::Connecting); - + // We can be put into the Disconnected state if a spontaneous error occurs between the connection process completing and this callback being invoked. + // We need to be able to handle that scenario here. HRESULT hr = HCGetWebSocketConnectResult(&context->internalAsyncBlock, &context->result); - std::unique_lock lock{ ws->m_stateMutex }; + const bool bIsDisconnected = (ws->m_state == State::Disconnected); + if (bIsDisconnected && !FAILED(hr)) + { + HC_TRACE_WARNING(WEBSOCKET, "WebSocket::ConnectComplete [%p] encountered a spontaneous disconnection. This implies the connection process was successful, but we otherwise had to close the connection (handle=%p)", ws.get(), context->observer.get()); + hr = E_FAIL; + } + + assert(ws->m_state == State::Connecting || bIsDisconnected); + assert(context->observer.get() == context->result.websocket || FAILED(hr) || bIsDisconnected); if (SUCCEEDED(hr) && SUCCEEDED(context->result.errorCode)) { // Connect was sucessful. Allocate ProviderContext to ensure WebSocket lifetime until it is reclaimed in WebSocket::CloseFunc @@ -528,7 +537,10 @@ void CALLBACK WebSocket::CloseFunc( std::unique_lock stateLock{ websocket->m_stateMutex }; if (!websocket->m_providerContext) { - HC_TRACE_ERROR(WEBSOCKET, "Unexpected call to WebSocket::CloseFunc will be ignored!"); + // It's possible for our websocket to get closed before we finish connecting. m_providerContext only gets populated when the connection process is 100% completed. + // If we're still in the process of connecting, mark as disconnected and let ConnectComplete handle the cleanup. + HC_TRACE_WARNING(WEBSOCKET, "Call to WebSocket::CloseFunc without providerContext. This means that we're aborting the connection process unexpectedly."); + websocket->m_state = State::Disconnected; return; } diff --git a/Tests/UnitTests/Tests/WebsocketTests.cpp b/Tests/UnitTests/Tests/WebsocketTests.cpp index 0d58a590..5a789650 100644 --- a/Tests/UnitTests/Tests/WebsocketTests.cpp +++ b/Tests/UnitTests/Tests/WebsocketTests.cpp @@ -7,6 +7,7 @@ #include "DefineTestMacros.h" #include "Utils.h" #include "../global/global.h" +#include "WebSocket/hcwebsocket.h" #include #pragma warning(disable:4389) @@ -135,6 +136,171 @@ HRESULT CALLBACK Test_Internal_HCWebSocketConnectAsync( }); } +bool g_HCWebSocketConnectAndClose_Called = false; +HRESULT CALLBACK Test_Internal_HCWebSocketConnectAsyncAndClose( + _In_z_ PCSTR uri, + _In_z_ PCSTR subProtocol, + _In_ HCWebsocketHandle websocket, + _Inout_ XAsyncBlock* asyncBlock, + _In_opt_ void* context, + _In_ HCPerformEnv env + ) +{ + UNREFERENCED_PARAMETER(uri); + UNREFERENCED_PARAMETER(subProtocol); + UNREFERENCED_PARAMETER(context); + UNREFERENCED_PARAMETER(env); + + return XAsyncBegin(asyncBlock, websocket, nullptr, __FUNCTION__, + [](XAsyncOp op, const XAsyncProviderData* data) + { + auto websocket = static_cast(data->context); + + switch (op) + { + case XAsyncOp::Begin: + { + g_HCWebSocketConnectAndClose_Called = true; + RETURN_IF_FAILED(XTaskQueueSubmitCallback(data->async->queue, XTaskQueuePort::Work, websocket, + [](void* context, bool canceled) + { + if (canceled) + { + return; + } + + auto websocket = static_cast(context); + HCWebSocketCloseEventFunction closeFunc = nullptr; + void* closeContext = nullptr; + HRESULT hr = HCWebSocketGetEventFunctions(websocket, nullptr, nullptr, &closeFunc, &closeContext); + if (SUCCEEDED(hr) && closeFunc != nullptr) + { + closeFunc(websocket, HCWebSocketCloseStatus::AbnormalClose, closeContext); + } + })); + + XAsyncComplete(data->async, S_OK, sizeof(WebSocketCompletionResult)); + return S_OK; + } + case XAsyncOp::GetResult: + { + RETURN_HR_IF(E_NOT_SUFFICIENT_BUFFER, data->bufferSize < sizeof(WebSocketCompletionResult)); + + auto result = static_cast(data->buffer); + ZeroMemory(result, sizeof(WebSocketCompletionResult)); + result->errorCode = S_OK; + result->websocket = websocket; + return S_OK; + } + default: + return S_OK; + } + }); +} + +class TestWebSocketConnectAndCloseProvider : public IWebSocketProvider +{ +public: + HRESULT ConnectAsync( + xbox::httpclient::String const& uri, + xbox::httpclient::String const& subprotocol, + HCWebsocketHandle websocketHandle, + XAsyncBlock* async + ) noexcept override + { + UNREFERENCED_PARAMETER(uri); + UNREFERENCED_PARAMETER(subprotocol); + + m_websocket = websocketHandle; + return XAsyncBegin(async, this, nullptr, __FUNCTION__, + [](XAsyncOp op, const XAsyncProviderData* data) + { + auto provider = static_cast(data->context); + + switch (op) + { + case XAsyncOp::Begin: + { + provider->m_connectCalled = true; + RETURN_IF_FAILED(XTaskQueueSubmitCallback(data->async->queue, XTaskQueuePort::Work, provider->m_websocket, + [](void* context, bool canceled) + { + if (canceled) + { + return; + } + + auto websocket = static_cast(context); + HCWebSocketCloseEventFunction closeFunc = nullptr; + void* closeContext = nullptr; + HRESULT hr = HCWebSocketGetEventFunctions(websocket, nullptr, nullptr, &closeFunc, &closeContext); + if (SUCCEEDED(hr) && closeFunc != nullptr) + { + closeFunc(websocket, HCWebSocketCloseStatus::AbnormalClose, closeContext); + } + })); + + XAsyncComplete(data->async, S_OK, sizeof(WebSocketCompletionResult)); + return S_OK; + } + case XAsyncOp::GetResult: + { + RETURN_HR_IF(E_NOT_SUFFICIENT_BUFFER, data->bufferSize < sizeof(WebSocketCompletionResult)); + + auto result = static_cast(data->buffer); + ZeroMemory(result, sizeof(WebSocketCompletionResult)); + result->errorCode = S_OK; + result->websocket = provider->m_websocket; + return S_OK; + } + default: + return S_OK; + } + }); + } + + HRESULT SendAsync( + HCWebsocketHandle websocketHandle, + const char* message, + XAsyncBlock* async + ) noexcept override + { + UNREFERENCED_PARAMETER(websocketHandle); + UNREFERENCED_PARAMETER(message); + UNREFERENCED_PARAMETER(async); + return E_UNEXPECTED; + } + + HRESULT SendBinaryAsync( + HCWebsocketHandle websocketHandle, + const uint8_t* payloadBytes, + uint32_t payloadSize, + XAsyncBlock* async + ) noexcept override + { + UNREFERENCED_PARAMETER(websocketHandle); + UNREFERENCED_PARAMETER(payloadBytes); + UNREFERENCED_PARAMETER(payloadSize); + UNREFERENCED_PARAMETER(async); + return E_UNEXPECTED; + } + + HRESULT Disconnect( + HCWebsocketHandle websocketHandle, + HCWebSocketCloseStatus closeStatus + ) noexcept override + { + UNREFERENCED_PARAMETER(websocketHandle); + UNREFERENCED_PARAMETER(closeStatus); + return E_UNEXPECTED; + } + + bool m_connectCalled{ false }; + +private: + HCWebsocketHandle m_websocket{ nullptr }; +}; + bool g_HCWebSocketSendMessage_Called = false; HRESULT CALLBACK Test_Internal_HCWebSocketSendMessageAsync( _In_ HCWebsocketHandle websocket, @@ -325,6 +491,62 @@ DEFINE_TEST_CLASS(WebsocketTests) HCCleanup(); } + DEFINE_TEST_CASE(TestConnectFailsWhenDisconnectedDuringCompletion) + { +#ifdef UNITTEST_TE + return; +#else + TestWebSocketConnectAndCloseProvider provider; + auto websocket = std::make_shared(1, provider); + + XAsyncBlock asyncBlock{}; + VERIFY_SUCCEEDED(XTaskQueueCreate(XTaskQueueDispatchMode::Manual, XTaskQueueDispatchMode::Manual, &asyncBlock.queue)); + + VERIFY_ARE_EQUAL(S_OK, websocket->ConnectAsync(http_internal_string{ "test" }, http_internal_string{ "subProtoTest" }, &asyncBlock)); + WebSocketCompletionResult connectResult{}; + XAsyncBlock sendAsyncBlock{}; + HRESULT connectStatus = E_PENDING; + HRESULT getConnectResultHr = E_PENDING; + HRESULT sendHr = S_OK; + HRESULT disconnectHr = S_OK; + + for (uint32_t attempt = 0; attempt < 8 && connectStatus == E_PENDING; ++attempt) + { + auto timeout = attempt == 0 ? 100u : 0u; + while (XTaskQueueDispatch(asyncBlock.queue, XTaskQueuePort::Work, timeout)) {} + while (XTaskQueueDispatch(asyncBlock.queue, XTaskQueuePort::Completion, timeout)) {} + connectStatus = XAsyncGetStatus(&asyncBlock, false); + } + + if (connectStatus == E_FAIL) + { + getConnectResultHr = HCGetWebSocketConnectResult(&asyncBlock, &connectResult); + } + else if (connectStatus == E_PENDING) + { + XAsyncCancel(&asyncBlock); + while (XTaskQueueDispatch(asyncBlock.queue, XTaskQueuePort::Work, 0)) {} + while (XTaskQueueDispatch(asyncBlock.queue, XTaskQueuePort::Completion, 0)) {} + connectStatus = XAsyncGetStatus(&asyncBlock, false); + } + + getConnectResultHr = HCGetWebSocketConnectResult(&asyncBlock, &connectResult); + sendHr = websocket->SendAsync("test", &sendAsyncBlock); + disconnectHr = websocket->Disconnect(); + + while (XTaskQueueDispatch(asyncBlock.queue, XTaskQueuePort::Work, 0)) {} + while (XTaskQueueDispatch(asyncBlock.queue, XTaskQueuePort::Completion, 0)) {} + XTaskQueueCloseHandle(asyncBlock.queue); + websocket.reset(); + + VERIFY_ARE_EQUAL(true, provider.m_connectCalled); + VERIFY_IS_TRUE(FAILED(connectStatus)); + VERIFY_IS_TRUE(FAILED(getConnectResultHr)); + VERIFY_ARE_EQUAL(E_UNEXPECTED, sendHr); + VERIFY_ARE_EQUAL(E_UNEXPECTED, disconnectHr); +#endif + } + DEFINE_TEST_CASE(TestRequestHeaders) {