Skip to content
This repository was archived by the owner on Jun 20, 2019. It is now read-only.

Commit a732b10

Browse files
pan-wangpan-wang
authored
Fix websocket close handshake issue and race condition when websocket client disconnect without close handshake (#151)
1 parent 509ddc6 commit a732b10

File tree

8 files changed

+76
-32
lines changed

8 files changed

+76
-32
lines changed

‎src/AspNetCore/Inc/environmentvariablehash.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,13 +136,15 @@ class ENVIRONMENT_VAR_HASH : public HASH_TABLE<ENVIRONMENT_VAR_ENTRY, PWSTR>
136136
PVOID pvData
137137
)
138138
{
139+
// best effort copy, ignore the failure
139140
ENVIRONMENT_VAR_ENTRY * pNewEntry = new ENVIRONMENT_VAR_ENTRY();
140141
if (pNewEntry != NULL)
141142
{
142143
pNewEntry->Initialize(pEntry->QueryName(), pEntry->QueryValue());
143144
ENVIRONMENT_VAR_HASH *pHash = static_cast<ENVIRONMENT_VAR_HASH *>(pvData);
144145
DBG_ASSERT(pHash);
145146
pHash->InsertRecord(pNewEntry);
147+
// Need to dereference as InsertRecord references it now
146148
pNewEntry->Dereference();
147149
}
148150
}

‎src/AspNetCore/Inc/forwardinghandler.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -322,6 +322,7 @@ class FORWARDING_HANDLER
322322
BOOL m_fWebSocketUpgrade;
323323
BOOL m_fFinishRequest;
324324
BOOL m_fClientDisconnected;
325+
BOOL m_fHasError;
325326
DWORD m_msStartTime;
326327

327328
DWORD m_BytesToReceive;

‎src/AspNetCore/Inc/serverprocess.h

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -313,10 +313,6 @@ class SERVER_PROCESS
313313
HANDLE m_hListeningProcessHandle;
314314
HANDLE m_hProcessWaitHandle;
315315
HANDLE m_hShutdownHandle;
316-
//
317-
// m_hChildProcessHandle is the handle to process created by
318-
// m_hProcessHandle process if it does.
319-
320316
//
321317
// m_hChildProcessHandle is the handle to process created by
322318
// m_hProcessHandle process if it does.

‎src/AspNetCore/Inc/websockethandler.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,9 @@ class WEBSOCKET_HANDLER
206206
volatile
207207
BOOL _fIndicateCompletionToIis;
208208

209+
volatile
210+
BOOL _fReceivedCloseMsg;
211+
209212
static
210213
LIST_ENTRY sm_RequestsListHead;
211214

‎src/AspNetCore/Src/aspnetcoreconfig.cxx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,7 @@ ASPNETCORE_CONFIG::Populate(
376376
strExpandedEnvValue.Reset();
377377
pEnvVar->Release();
378378
pEnvVar = NULL;
379+
pEntry->Dereference();
379380
pEntry = NULL;
380381
}
381382

‎src/AspNetCore/Src/forwardinghandler.cxx

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,8 @@ m_pAppOfflineHtm(NULL),
5353
m_fErrorHandled(FALSE),
5454
m_fWebSocketUpgrade(FALSE),
5555
m_fFinishRequest(FALSE),
56-
m_fClientDisconnected(FALSE)
56+
m_fClientDisconnected(FALSE),
57+
m_fHasError(FALSE)
5758
{
5859
#ifdef DEBUG
5960
DBGPRINTF((DBG_CONTEXT,
@@ -1171,6 +1172,7 @@ FORWARDING_HANDLER::OnExecuteRequestHandler(
11711172
FALSE
11721173
); // no need to check return hresult
11731174

1175+
m_fHasError = TRUE;
11741176
goto Finished;
11751177
}
11761178

@@ -1368,6 +1370,7 @@ FORWARDING_HANDLER::OnExecuteRequestHandler(
13681370
// Reset status for consistency.
13691371
//
13701372
m_RequestStatus = FORWARDER_DONE;
1373+
m_fHasError = TRUE;
13711374

13721375
pResponse->DisableKernelCache();
13731376
pResponse->GetRawHttpResponse()->EntityChunkCount = 0;
@@ -1592,7 +1595,10 @@ REQUEST_NOTIFICATION_STATUS
15921595
{
15931596
if (m_RequestStatus == FORWARDER_DONE && m_fFinishRequest)
15941597
{
1595-
retVal = RQ_NOTIFICATION_FINISH_REQUEST;
1598+
if (m_fHasError)
1599+
{
1600+
retVal = RQ_NOTIFICATION_FINISH_REQUEST;
1601+
}
15961602
goto Finished;
15971603
}
15981604

@@ -1706,7 +1712,7 @@ REQUEST_NOTIFICATION_STATUS
17061712
// Reset status for consistency.
17071713
//
17081714
m_RequestStatus = FORWARDER_DONE;
1709-
1715+
m_fHasError = TRUE;
17101716
//
17111717
// Do the right thing based on where the error originated from.
17121718
//
@@ -1779,8 +1785,17 @@ REQUEST_NOTIFICATION_STATUS
17791785

17801786
//
17811787
// Finish the request on failure.
1788+
// Let IIS pipeline continue only after receiving handle close callback
1789+
// from WinHttp. This ensures no more callback from WinHttp
17821790
//
1783-
retVal = RQ_NOTIFICATION_FINISH_REQUEST;
1791+
if (m_hRequest != NULL)
1792+
{
1793+
if (WinHttpCloseHandle(m_hRequest))
1794+
{
1795+
m_hRequest = NULL;
1796+
}
1797+
}
1798+
retVal = RQ_NOTIFICATION_PENDING;
17841799

17851800
Finished:
17861801

@@ -2249,7 +2264,6 @@ None
22492264
fDerefForwardingHandler = m_fClientDisconnected && !m_fWebSocketUpgrade;
22502265
}
22512266
m_hRequest = NULL;
2252-
m_pWebSocket = NULL;
22532267
fAnotherCompletionExpected = FALSE;
22542268
break;
22552269
case WINHTTP_CALLBACK_STATUS_CONNECTION_CLOSED:
@@ -2391,11 +2405,13 @@ None
23912405
m_hRequest = NULL;
23922406
}
23932407
}
2408+
2409+
//
2410+
// If the request is a websocket request, initiate cleanup.
2411+
//
23942412
if (m_pWebSocket != NULL)
23952413
{
23962414
m_pWebSocket->TerminateRequest();
2397-
m_pWebSocket->Terminate();
2398-
m_pWebSocket = NULL;
23992415
}
24002416

24012417
if(fEndRequest)

‎src/AspNetCore/Src/serverprocess.cxx

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2280,13 +2280,18 @@ SERVER_PROCESS::SendShutDownSignalInternal(
22802280
TerminateBackendProcess();
22812281
}
22822282
FreeConsole();
2283-
}
22842283

2285-
if (fFreeConsole)
2284+
if (fFreeConsole)
2285+
{
2286+
// IISExpress and hostedwebcore w3wp run as background process
2287+
// have to attach console back to ensure post app_offline scenario still works
2288+
AttachConsole(ATTACH_PARENT_PROCESS);
2289+
}
2290+
}
2291+
else
22862292
{
2287-
// IISExpress and hostedwebcore w3wp run as background process
2288-
// have to attach console back to ensure post app_offline scenario still works
2289-
AttachConsole(ATTACH_PARENT_PROCESS);
2293+
// terminate the backend process immediately instead of waiting for timeout
2294+
TerminateBackendProcess();
22902295
}
22912296
}
22922297

‎src/AspNetCore/Src/websockethandler.cxx

Lines changed: 36 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,15 @@ LIST_ENTRY WEBSOCKET_HANDLER::sm_RequestsListHead;
3535

3636
TRACE_LOG * WEBSOCKET_HANDLER::sm_pTraceLog;
3737

38-
WEBSOCKET_HANDLER::WEBSOCKET_HANDLER():
39-
_pHttpContext ( NULL ),
40-
_pWebSocketContext ( NULL ),
41-
_pHandler ( NULL ),
42-
_dwOutstandingIo ( 0 ),
43-
_fCleanupInProgress ( FALSE ),
44-
_fIndicateCompletionToIis ( FALSE )
38+
WEBSOCKET_HANDLER::WEBSOCKET_HANDLER() :
39+
_pHttpContext(NULL),
40+
_pWebSocketContext(NULL),
41+
_hWebSocketRequest(NULL),
42+
_pHandler(NULL),
43+
_dwOutstandingIo(0),
44+
_fCleanupInProgress(FALSE),
45+
_fIndicateCompletionToIis(FALSE),
46+
_fReceivedCloseMsg(FALSE)
4547
{
4648
DebugPrintf (ASPNETCORE_DEBUG_FLAG_INFO, "WEBSOCKET_HANDLER::WEBSOCKET_HANDLER");
4749

@@ -58,16 +60,20 @@ WEBSOCKET_HANDLER::Terminate(
5860
DebugPrintf (ASPNETCORE_DEBUG_FLAG_INFO, "WEBSOCKET_HANDLER::Terminate");
5961

6062
RemoveRequest();
63+
_fCleanupInProgress = TRUE;
6164

62-
_pWebSocketContext = NULL;
63-
_pHttpContext = NULL;
64-
65+
if (_pHttpContext != NULL)
66+
{
67+
_pHttpContext->CancelIo();
68+
_pHttpContext = NULL;
69+
}
6570
if (_hWebSocketRequest)
6671
{
6772
WinHttpCloseHandle(_hWebSocketRequest);
6873
_hWebSocketRequest = NULL;
6974
}
7075

76+
_pWebSocketContext = NULL;
7177
DeleteCriticalSection(&_RequestLock);
7278

7379
delete this;
@@ -216,6 +222,9 @@ WEBSOCKET_HANDLER::IndicateCompletionToIIS(
216222
// wait for handle close callback and then call IndicateCompletion
217223
// otherwise we may release W3Context too early and cause AV
218224
//_pHttpContext->IndicateCompletion(RQ_NOTIFICATION_PENDING);
225+
// close Websocket handle. This will triger a WinHttp callback
226+
// on handle close, then let IIS pipeline continue.
227+
WinHttpCloseHandle(_hWebSocketRequest);
219228
}
220229

221230
HRESULT
@@ -498,6 +507,12 @@ Routine Description:
498507
}
499508

500509
IncrementOutstandingIo();
510+
//
511+
// Backend end may start close hand shake first
512+
// Need to inidcate no more receive should be called on WinHttp conneciton
513+
//
514+
_fReceivedCloseMsg = TRUE;
515+
_fIndicateCompletionToIis = TRUE;
501516

502517
//
503518
// Send close to IIS.
@@ -947,14 +962,19 @@ Routine Description:
947962
}
948963

949964
//
950-
// Write Completed, initiate next read from backend server.
965+
// Only call read if no close hand shake was received from backend
951966
//
952-
953-
hr = DoWinHttpWebSocketReceive();
954-
if (FAILED (hr))
967+
if (!_fReceivedCloseMsg)
955968
{
956-
cleanupReason = ServerDisconnect;
957-
goto Finished;
969+
//
970+
// Write Completed, initiate next read from backend server.
971+
//
972+
hr = DoWinHttpWebSocketReceive();
973+
if (FAILED(hr))
974+
{
975+
cleanupReason = ServerDisconnect;
976+
goto Finished;
977+
}
958978
}
959979

960980
Finished:

0 commit comments

Comments
 (0)