fix websocket connection issue and some memory leak, and add debug print (#129)

src/AspNetCore/Inc/environmentvariablehash.h

@@ -122,6 +122,8 @@ public:
     {
         STRU     strTemp;
         MULTISZ   *pMultiSz = static_cast<MULTISZ *>(pvData);
+        DBG_ASSERT(pMultiSz);
+        DBG_ASSERT(pEntry);
         strTemp.Copy(pEntry->QueryName());
         strTemp.Append(pEntry->QueryValue());
         pMultiSz->Append(strTemp.QueryStr());
@@ -135,9 +137,14 @@ public:
     )
     {
         ENVIRONMENT_VAR_ENTRY *   pNewEntry = new ENVIRONMENT_VAR_ENTRY();
-        pNewEntry->Initialize(pEntry->QueryName(), pEntry->QueryValue());
-        ENVIRONMENT_VAR_HASH *pHash = static_cast<ENVIRONMENT_VAR_HASH *>(pvData);
-        pHash->InsertRecord(pNewEntry);
+        if (pNewEntry != NULL) 
+        {
+            pNewEntry->Initialize(pEntry->QueryName(), pEntry->QueryValue());
+            ENVIRONMENT_VAR_HASH *pHash = static_cast<ENVIRONMENT_VAR_HASH *>(pvData);
+            DBG_ASSERT(pHash);
+            pHash->InsertRecord(pNewEntry);
+            pNewEntry->Dereference();
+        }
     }
 
 private:

src/AspNetCore/Src/dllmain.cpp

@@ -4,7 +4,12 @@
 #include "precomp.hxx"
 #include <IPHlpApi.h>
 
-//DECLARE_DEBUG_PRINT_OBJECT("Asp.Net Core Module");
+#ifdef DEBUG
+    DECLARE_DEBUG_PRINTS_OBJECT();
+    DECLARE_DEBUG_VARIABLE();
+    DECLARE_PLATFORM_TYPE();
+#endif // DEBUG
+
 
 HTTP_MODULE_ID      g_pModuleId = NULL;
 IHttpServer *       g_pHttpServer = NULL;
@@ -142,6 +147,11 @@ HRESULT
     HRESULT                 hr = S_OK;
     CProxyModuleFactory *   pFactory = NULL;
 
+#ifdef DEBUG
+    CREATE_DEBUG_PRINT_OBJECT("Asp.Net Core Module");
+    g_dwDebugFlags = DEBUG_FLAGS_ANY;
+#endif // DEBUG
+
     CREATE_DEBUG_PRINT_OBJECT;
 
     LoadGlobalConfiguration();

src/AspNetCore/Src/forwarderconnection.cxx

@@ -33,6 +33,19 @@ FORWARDER_CONNECTION::Initialize(
         goto Finished;
     }
 
+    //
+    // Since WinHttp will not emit WINHTTP_CALLBACK_STATUS_HANDLE_CLOSING
+    // when closing WebSocket handle on Win8. Register callback at Connect level as a workaround
+    //
+    if (WinHttpSetStatusCallback(m_hConnection,
+                                 FORWARDING_HANDLER::OnWinHttpCompletion,
+                                 WINHTTP_CALLBACK_FLAG_HANDLES,
+                                 NULL) == WINHTTP_INVALID_STATUS_CALLBACK)
+    {
+        hr = HRESULT_FROM_WIN32(GetLastError());
+        goto Finished;
+    }
+
 Finished:
 
     return hr;

src/AspNetCore/Src/forwardinghandler.cxx

@@ -2,6 +2,8 @@
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 #include "precomp.hxx"
+#include <dbgutil.h>
+
 
 // Just to be aware of the FORWARDING_HANDLER object size.
 C_ASSERT(sizeof(FORWARDING_HANDLER) <= 632);
@@ -53,6 +55,11 @@ m_fWebSocketUpgrade(FALSE),
 m_fFinishRequest(FALSE),
 m_fClientDisconnected(FALSE)
 {
+#ifdef DEBUG
+    DBGPRINTF((DBG_CONTEXT,
+        "FORWARDING_HANDLER::FORWARDING_HANDLER \n"));
+#endif // DEBUG
+
     InitializeSRWLock(&m_RequestLock);
 }
 
@@ -63,6 +70,11 @@ FORWARDING_HANDLER::~FORWARDING_HANDLER(
     //
     // Destructor has started.
     //
+#ifdef DEBUG
+    DBGPRINTF((DBG_CONTEXT,
+        "FORWARDING_HANDLER::~FORWARDING_HANDLER \n"));
+#endif // DEBUG
+
     m_Signature = FORWARDING_HANDLER_SIGNATURE_FREE;
 
     //
@@ -2101,6 +2113,11 @@ None
         DBG_ASSERT(TlsGetValue(g_dwTlsIndex) == this);
     }
 
+#ifdef DEBUG
+    DBGPRINTF((DBG_CONTEXT,
+        "FORWARDING_HANDLER::OnWinHttpCompletionInternal %x -- %d --%p\n", dwInternetStatus, m_fWebSocketUpgrade, m_pW3Context));
+#endif // DEBUG
+
     fEndRequest = (dwInternetStatus == WINHTTP_CALLBACK_STATUS_HANDLE_CLOSING);
     if (!fEndRequest) 
     {
@@ -2377,6 +2394,7 @@ Finished:
         if (m_pWebSocket != NULL)
         {
             m_pWebSocket->TerminateRequest();
+            m_pWebSocket->Terminate();
             m_pWebSocket = NULL;
         }
 

src/AspNetCore/Src/serverprocess.cxx

@@ -443,7 +443,10 @@ SERVER_PROCESS::InitEnvironmentVariablesTable(
     if (dwResult == 0)
     {
         dwError = GetLastError();
-        if (dwError != ERROR_ENVVAR_NOT_FOUND)
+        // Windows API (e.g., CreateProcess) allows variable with empty string value
+        // in such case dwResult will be 0 and dwError will also be 0
+        // As UI and CMD does not allow empty value, ignore this environment var
+        if (dwError != ERROR_ENVVAR_NOT_FOUND && dwError != ERROR_SUCCESS)
         {
             hr = HRESULT_FROM_WIN32(dwError);
             goto Finished;
@@ -1421,23 +1424,21 @@ SERVER_PROCESS::CheckIfServerIsUp(
             hr = HRESULT_FROM_WIN32(WSAGetLastError());
             goto Finished;
         }
-
-        //
-        // Connected successfully, close socket.
-        //
-        iResult = closesocket(socketCheck);
-        if (iResult == SOCKET_ERROR)
-        {
-            hr = HRESULT_FROM_WIN32(WSAGetLastError());
-            goto Finished;
-        }
-
-        socketCheck = INVALID_SOCKET;
         *pfReady = TRUE;
     }
 
 Finished:
 
+    if (socketCheck != INVALID_SOCKET)
+    {
+        iResult = closesocket(socketCheck);
+        if (iResult == SOCKET_ERROR)
+        {
+            hr = HRESULT_FROM_WIN32(WSAGetLastError());
+        }
+        socketCheck = INVALID_SOCKET;
+    }
+
     if (pTCPInfo != NULL)
     {
         HeapFree(GetProcessHeap(), 0, pTCPInfo);
@@ -1650,6 +1651,12 @@ SERVER_PROCESS::IsDebuggerIsAttached(
                     dwPid);
 
             BOOL returnValue = CheckRemoteDebuggerPresent(hProcess, &fDebuggerPresent);
+            if (hProcess != NULL)
+            {
+                CloseHandle(hProcess);
+                hProcess = NULL;
+            }
+
             if (!returnValue)
             {
                 goto Finished;