Fix websocket close handshake issue and race condition when websocket client disconnect without close handshake (#151)

src/AspNetCore/Inc/environmentvariablehash.h

@@ -136,6 +136,7 @@ public:
         PVOID                     pvData
     )
     {
+        // best effort copy, ignore the failure
         ENVIRONMENT_VAR_ENTRY *   pNewEntry = new ENVIRONMENT_VAR_ENTRY();
         if (pNewEntry != NULL) 
         {
@@ -143,6 +144,7 @@ public:
             ENVIRONMENT_VAR_HASH *pHash = static_cast<ENVIRONMENT_VAR_HASH *>(pvData);
             DBG_ASSERT(pHash);
             pHash->InsertRecord(pNewEntry);
+            // Need to dereference as InsertRecord references it now
             pNewEntry->Dereference();
         }
     }

src/AspNetCore/Inc/forwardinghandler.h

@@ -322,6 +322,7 @@ private:
     BOOL                                m_fWebSocketUpgrade;
     BOOL                                m_fFinishRequest;
     BOOL                                m_fClientDisconnected;
+    BOOL                                m_fHasError;
     DWORD                               m_msStartTime;
         
     DWORD                               m_BytesToReceive;

src/AspNetCore/Inc/serverprocess.h

@@ -313,10 +313,6 @@ private:
     HANDLE                  m_hListeningProcessHandle;
     HANDLE                  m_hProcessWaitHandle;
     HANDLE                  m_hShutdownHandle;
-    //
-    // m_hChildProcessHandle is the handle to process created by 
-    // m_hProcessHandle process if it does.
-
     //
     // m_hChildProcessHandle is the handle to process created by 
     // m_hProcessHandle process if it does.

src/AspNetCore/Inc/websockethandler.h

@@ -206,6 +206,9 @@ private:
     volatile
     BOOL                _fIndicateCompletionToIis;
 
+    volatile
+    BOOL                _fReceivedCloseMsg;
+
     static
     LIST_ENTRY          sm_RequestsListHead;
 

src/AspNetCore/Src/aspnetcoreconfig.cxx

@@ -376,6 +376,7 @@ ASPNETCORE_CONFIG::Populate(
         strExpandedEnvValue.Reset();
         pEnvVar->Release();
         pEnvVar = NULL;
+        pEntry->Dereference();
         pEntry = NULL;
     }
 

src/AspNetCore/Src/forwardinghandler.cxx

@@ -53,7 +53,8 @@ m_pAppOfflineHtm(NULL),
 m_fErrorHandled(FALSE),
 m_fWebSocketUpgrade(FALSE),
 m_fFinishRequest(FALSE),
-m_fClientDisconnected(FALSE)
+m_fClientDisconnected(FALSE),
+m_fHasError(FALSE)
 {
 #ifdef DEBUG
     DBGPRINTF((DBG_CONTEXT,
@@ -1171,6 +1172,7 @@ FORWARDING_HANDLER::OnExecuteRequestHandler(
             FALSE
         ); // no need to check return hresult
 
+        m_fHasError = TRUE;
         goto Finished;
     }
 
@@ -1368,6 +1370,7 @@ Failure:
     // Reset status for consistency.
     //
     m_RequestStatus = FORWARDER_DONE;
+    m_fHasError = TRUE;
 
     pResponse->DisableKernelCache();
     pResponse->GetRawHttpResponse()->EntityChunkCount = 0;
@@ -1592,7 +1595,10 @@ REQUEST_NOTIFICATION_STATUS
     {
         if (m_RequestStatus == FORWARDER_DONE && m_fFinishRequest)
         {
-            retVal = RQ_NOTIFICATION_FINISH_REQUEST;
+            if (m_fHasError)
+            {
+                retVal = RQ_NOTIFICATION_FINISH_REQUEST;
+            }
             goto Finished;
         }
 
@@ -1706,7 +1712,7 @@ Failure:
     // Reset status for consistency.
     //
     m_RequestStatus = FORWARDER_DONE;
-
+    m_fHasError = TRUE;
     //
     // Do the right thing based on where the error originated from.
     //
@@ -1779,8 +1785,17 @@ Failure:
 
     //
     // Finish the request on failure.
+    // Let IIS pipeline continue only after receiving handle close callback
+    // from WinHttp. This ensures no more callback from WinHttp
     //
-    retVal = RQ_NOTIFICATION_FINISH_REQUEST;
+    if (m_hRequest != NULL)
+    {
+        if (WinHttpCloseHandle(m_hRequest))
+        {
+            m_hRequest = NULL;
+        }
+    }
+    retVal = RQ_NOTIFICATION_PENDING;
 
 Finished:
 
@@ -2249,7 +2264,6 @@ None
             fDerefForwardingHandler = m_fClientDisconnected && !m_fWebSocketUpgrade;
         }
         m_hRequest = NULL;
-        m_pWebSocket = NULL;
         fAnotherCompletionExpected = FALSE;
         break;
     case WINHTTP_CALLBACK_STATUS_CONNECTION_CLOSED:
@@ -2391,11 +2405,13 @@ Finished:
                 m_hRequest = NULL;
             }
         }
+
+        //
+        // If the request is a websocket request, initiate cleanup.
+        //
         if (m_pWebSocket != NULL)
         {
             m_pWebSocket->TerminateRequest();
-            m_pWebSocket->Terminate();
-            m_pWebSocket = NULL;
         }
 
         if(fEndRequest)

src/AspNetCore/Src/serverprocess.cxx

@@ -2280,13 +2280,18 @@ SERVER_PROCESS::SendShutDownSignalInternal(
                 TerminateBackendProcess();
             }
             FreeConsole();
-        }
 
-        if (fFreeConsole)
+            if (fFreeConsole)
+            {
+                // IISExpress and hostedwebcore w3wp run as background process
+                // have to attach console back to ensure post app_offline scenario still works
+                AttachConsole(ATTACH_PARENT_PROCESS);
+            }
+        }
+        else
         {
-            // IISExpress and hostedwebcore w3wp run as background process
+            // terminate the backend process immediately instead of waiting for timeout
-            // have to attach console back to ensure post app_offline scenario still works
+            TerminateBackendProcess();
-            AttachConsole(ATTACH_PARENT_PROCESS);
         }
     }
 

src/AspNetCore/Src/websockethandler.cxx

@@ -35,13 +35,15 @@ LIST_ENTRY WEBSOCKET_HANDLER::sm_RequestsListHead;
 
 TRACE_LOG * WEBSOCKET_HANDLER::sm_pTraceLog;
 
-WEBSOCKET_HANDLER::WEBSOCKET_HANDLER():
+WEBSOCKET_HANDLER::WEBSOCKET_HANDLER() :
-    _pHttpContext ( NULL ),
+    _pHttpContext(NULL),
-    _pWebSocketContext ( NULL ),
+    _pWebSocketContext(NULL),
-    _pHandler ( NULL ),
+    _hWebSocketRequest(NULL),
-    _dwOutstandingIo ( 0 ),
+    _pHandler(NULL),
-    _fCleanupInProgress ( FALSE ),
+    _dwOutstandingIo(0),
-    _fIndicateCompletionToIis ( FALSE )
+    _fCleanupInProgress(FALSE),
+    _fIndicateCompletionToIis(FALSE),
+    _fReceivedCloseMsg(FALSE)
 {
     DebugPrintf (ASPNETCORE_DEBUG_FLAG_INFO, "WEBSOCKET_HANDLER::WEBSOCKET_HANDLER");
 
@@ -58,16 +60,20 @@ WEBSOCKET_HANDLER::Terminate(
     DebugPrintf (ASPNETCORE_DEBUG_FLAG_INFO, "WEBSOCKET_HANDLER::Terminate");
 
     RemoveRequest();
+    _fCleanupInProgress = TRUE;
 
-    _pWebSocketContext = NULL;
+    if (_pHttpContext != NULL)
-    _pHttpContext = NULL;
+    {
-
+        _pHttpContext->CancelIo();
+        _pHttpContext = NULL;
+    }
     if (_hWebSocketRequest)
     {
         WinHttpCloseHandle(_hWebSocketRequest);
         _hWebSocketRequest = NULL;
     }
 
+    _pWebSocketContext = NULL;
     DeleteCriticalSection(&_RequestLock);
 
     delete this;
@@ -216,6 +222,9 @@ WEBSOCKET_HANDLER::IndicateCompletionToIIS(
     // wait for handle close callback and then call IndicateCompletion
     // otherwise we may release W3Context too early and cause AV
     //_pHttpContext->IndicateCompletion(RQ_NOTIFICATION_PENDING);
+    // close Websocket handle. This will triger a WinHttp callback
+    // on handle close, then let IIS pipeline continue.
+    WinHttpCloseHandle(_hWebSocketRequest);
 }
 
 HRESULT
@@ -498,6 +507,12 @@ Routine Description:
         }
 
         IncrementOutstandingIo();
+        //
+        // Backend end may start close hand shake first
+        // Need to inidcate no more receive should be called on WinHttp conneciton
+        //
+        _fReceivedCloseMsg = TRUE;
+        _fIndicateCompletionToIis = TRUE;
 
         //
         // Send close to IIS.
@@ -947,14 +962,19 @@ Routine Description:
     }
 
     //
-    // Write Completed, initiate next read from backend server.
+    // Only call read if no close hand shake was received from backend
     //
-
+    if (!_fReceivedCloseMsg)
-    hr = DoWinHttpWebSocketReceive();
-    if (FAILED (hr))
     {
-        cleanupReason = ServerDisconnect;
+        //
-        goto Finished;
+        // Write Completed, initiate next read from backend server.
+        //
+        hr = DoWinHttpWebSocketReceive();
+        if (FAILED(hr))
+        {
+            cleanupReason = ServerDisconnect;
+            goto Finished;
+        }
     }
 
 Finished: