Merge pull request #446 from justcoding121/master

Fix prefetch connection leak

Author: justcoding121

Titanium.Web.Proxy/ExplicitClientHandler.cs

@@ -150,8 +150,12 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
                         }
 
                         SslStream sslStream = null;
+
+                        //don't pass cancellation token here
+                        //it could cause floating server connections when client exits
                         prefetchConnectionTask = getServerConnection(connectArgs, true,
-                                null, false, cancellationToken);
+                                null, false, CancellationToken.None);
+
                         try
                         {
                             sslStream = new SslStream(clientStream);
@@ -250,7 +254,7 @@ await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
                         {
                             await tcpConnectionFactory.Release(connection, true);
                         }
-                        calledRequestHandler = true;
+
                         return;
                     }
                 }
@@ -331,19 +335,27 @@ await handleHttpSessionRequest(endPoint, clientConnection, clientStream, clientS
             }
             finally
             {
+                if (!calledRequestHandler
+                        && prefetchConnectionTask != null)
+                {
+                    TcpServerConnection prefetchedConnection = null;
+                    try
+                    {
+                        prefetchedConnection = await prefetchConnectionTask;
+
+                    }
+                    finally
+                    {
+                        await tcpConnectionFactory.Release(prefetchedConnection, closeServerConnection);
+                    }
+                }
+
                 clientStream.Dispose();
 
                 if (!cancellationTokenSource.IsCancellationRequested)
                 {
                     cancellationTokenSource.Cancel();
                 }
-
-                if (!calledRequestHandler
-                        && prefetchConnectionTask != null)
-                {
-                    var connection = await prefetchConnectionTask;
-                    await tcpConnectionFactory.Release(connection, closeServerConnection);
-                }
             }
         }
     }

Titanium.Web.Proxy/ProxyServer.cs

@@ -815,12 +815,12 @@ private Policy retryPolicy<T>() where T : Exception
                 .RetryAsync(retries,
                     onRetryAsync: async (ex, i, context) =>
                     {
-                        if (context.ContainsKey("connection"))
+                        if (context["connection"] != null)
                         {
                             //close connection on error
                             var connection = (TcpServerConnection)context["connection"];
                             await tcpConnectionFactory.Release(connection, true);
-                            context.Remove("connection");
+                            context["connection"] = null;
                         }
 
                     });

Titanium.Web.Proxy/RequestHandler.cs

@@ -58,14 +58,14 @@ private async Task handleHttpSessionRequest(ProxyEndPoint endPoint, TcpClientCon
             CancellationTokenSource cancellationTokenSource, string httpsConnectHostname, ConnectRequest connectRequest,
             Task<TcpServerConnection> prefetchConnectionTask = null)
         {
-            var cancellationToken = cancellationTokenSource.Token;
-
             var prefetchTask = prefetchConnectionTask;
             TcpServerConnection connection = null;
             bool closeServerConnection = false;
 
             try
             {
+                var cancellationToken = cancellationTokenSource.Token;
+
                 // Loop through each subsequest request on this particular client connection
                 // (assuming HTTP connection is kept alive by client)
                 while (true)
@@ -204,21 +204,13 @@ await clientStreamWriter.WriteResponseAsync(args.WebSession.Response,
                                 connection = null;
                             }
 
-                            var contextData = new Dictionary<string, object>();
-                            if (connection != null)
-                            {
-                                contextData.Add("connection", connection);
-                            }
-
                             //for connection pool retry fails until cache is exhausted   
                             await retryPolicy<ServerConnectionException>().ExecuteAsync(async (context) =>
                             {
-
-                               connection = context.ContainsKey("connection")? 
-                                    (TcpServerConnection)context["connection"]
-                                    : await getServerConnection(args, false,
-                                                     clientConnection.NegotiatedApplicationProtocol, 
-                                                     false, cancellationToken);
+                                connection = context["connection"] as TcpServerConnection ??
+                                            await getServerConnection(args, false,
+                                                clientConnection.NegotiatedApplicationProtocol,
+                                                false, cancellationToken);
 
                                 context["connection"] = connection;
 
@@ -235,8 +227,7 @@ await handleWebSocketUpgrade(httpCmd, args, request,
                                 // construct the web request that we are going to issue on behalf of the client.
                                 await handleHttpSessionRequestInternal(connection, args);
 
-                            }, contextData);
-
+                            }, new Dictionary<string, object> { { "connection", connection } });
 
                             //user requested
                             if (args.WebSession.CloseServerConnection)
@@ -288,12 +279,21 @@ await handleWebSocketUpgrade(httpCmd, args, request,
             finally
             {
                 await tcpConnectionFactory.Release(connection,
-                    closeServerConnection);
+                        closeServerConnection);
+
 
                 if (prefetchTask != null)
                 {
-                    await tcpConnectionFactory.Release(await prefetchTask,
-                            closeServerConnection);
+                    TcpServerConnection prefetchedConnection = null;
+                    try
+                    {
+                        prefetchedConnection = await prefetchTask;
+
+                    }
+                    finally
+                    {
+                        await tcpConnectionFactory.Release(prefetchedConnection, closeServerConnection);
+                    }
                 }
             }
         }

Titanium.Web.Proxy/TransparentClientHandler.cs

@@ -63,9 +63,11 @@ private async Task handleClient(TransparentProxyEndPoint endPoint, TcpClientConn
 
                     if (endPoint.DecryptSsl && args.DecryptSsl)
                     {
+                        //don't pass cancellation token here
+                        //it could cause floating server connections when client exits
                         prefetchConnectionTask = tcpConnectionFactory.GetClient(httpsHostName, endPoint.Port,
                                 null, true, null, false, this,
-                                UpStreamEndPoint, UpStreamHttpsProxy, false, cancellationToken);
+                                UpStreamEndPoint, UpStreamHttpsProxy, false, CancellationToken.None);
 
                         SslStream sslStream = null;
 
@@ -163,8 +165,16 @@ await handleHttpSessionRequest(endPoint, clientConnection, clientStream, clientS
                 if (!calledRequestHandler
                    && prefetchConnectionTask != null)
                 {
-                    var connection = await prefetchConnectionTask;
-                    await tcpConnectionFactory.Release(connection, closeServerConnection);
+                    TcpServerConnection prefetchedConnection = null;
+                    try
+                    {
+                        prefetchedConnection = await prefetchConnectionTask;
+                       
+                    }
+                    finally
+                    {
+                        await tcpConnectionFactory.Release(prefetchedConnection, closeServerConnection);
+                    }
                 }
 
                 clientStream.Dispose();