fix performance issues

Author: justcoding121

Titanium.Web.Proxy/Models/ConnectRequest.cs

@@ -14,10 +14,6 @@ namespace Titanium.Web.Proxy.Network
     /// </summary>
     internal class CustomSslStream : SslStream
     {
-        /// <summary>
-        /// Holds the current session
-        /// </summary>
-        internal object Param { get; set; }
 
         internal CustomSslStream(Stream innerStream, bool leaveInnerStreamOpen, RemoteCertificateValidationCallback userCertificateValidationCallback, LocalCertificateSelectionCallback clientCertificateSelectionCallback) 
             :base(innerStream, leaveInnerStreamOpen, userCertificateValidationCallback, clientCertificateSelectionCallback)

Titanium.Web.Proxy/Network/CustomSslStream.cs

@@ -40,11 +40,8 @@ internal class TcpConnectionManager
         static Dictionary<string, List<TcpConnection>> connectionCache = new Dictionary<string, List<TcpConnection>>();
         static SemaphoreSlim connectionAccessLock = new SemaphoreSlim(1);
 
+
         internal static async Task<TcpConnection> GetClient(string hostname, int port, bool isHttps, Version version)
-        {
-            return await GetClient(null, hostname, port, isHttps, version);
-        }
-        internal static async Task<TcpConnection> GetClient(ConnectRequest connectRequest, string hostname, int port, bool isHttps, Version version)
         {
             List<TcpConnection> cachedConnections = null;
             TcpConnection cached = null;
@@ -71,14 +68,24 @@ internal static async Task<TcpConnection> GetClient(ConnectRequest connectReques
                 finally { connectionAccessLock.Release(); }
 
                 if (cached != null && !cached.TcpClient.Client.IsConnected())
+                {
+                    cached.TcpClient.Client.Dispose();
+                    cached.TcpClient.Close();
                     continue;
+                }
 
                 if (cached == null)
                     break;
             }
 
             if (cached == null)
-                cached = await CreateClient(connectRequest, hostname, port, isHttps, version).ConfigureAwait(false);
+                cached = await CreateClient(hostname, port, isHttps, version).ConfigureAwait(false);
+
+            if (cachedConnections == null || cachedConnections.Count() <= 2)
+            {
+                await CreateClient(hostname, port, isHttps, version)
+                              .ContinueWith(async (x) => { if (x.Status == TaskStatus.RanToCompletion) await ReleaseClient(x.Result); });
+            }
 
             return cached;
         }
@@ -88,7 +95,7 @@ internal static string GetConnectionKey(string hostname, int port, bool isHttps,
             return string.Format("{0}:{1}:{2}:{3}:{4}", hostname.ToLower(), port, isHttps, version.Major, version.Minor);
         }
 
-        private static async Task<TcpConnection> CreateClient(ConnectRequest connectRequest, string hostname, int port, bool isHttps, Version version)
+        private static async Task<TcpConnection> CreateClient(string hostname, int port, bool isHttps, Version version)
         {
             TcpClient client;
             Stream stream;
@@ -132,7 +139,6 @@ private static async Task<TcpConnection> CreateClient(ConnectRequest connectRequ
                 {
                     sslStream = new CustomSslStream(stream, true, new RemoteCertificateValidationCallback(ProxyServer.ValidateServerCertificate),
                         new LocalCertificateSelectionCallback(ProxyServer.SelectClientCertificate));
-                    sslStream.Param  = connectRequest;
                     await sslStream.AuthenticateAsClientAsync(hostname, null, Constants.SupportedProtocols, false).ConfigureAwait(false);
                     stream = (Stream)sslStream;
                 }
@@ -172,6 +178,7 @@ private static async Task<TcpConnection> CreateClient(ConnectRequest connectRequ
 
         internal static async Task ReleaseClient(TcpConnection connection)
         {
+
             connection.LastAccess = DateTime.Now;
             var key = GetConnectionKey(connection.HostName, connection.port, connection.IsHttps, connection.Version);
             await connectionAccessLock.WaitAsync();

Titanium.Web.Proxy/Network/TcpConnectionManager.cs

@@ -72,29 +72,23 @@ private static async void HandleClient(ExplicitProxyEndPoint endPoint, TcpClient
 
                     await WriteConnectResponse(clientStreamWriter, version).ConfigureAwait(false);
 
-                    var certificate = await CertManager.CreateCertificate(httpRemoteUri.Host, false);
+
 
                     SslStream sslStream = null;
 
                     try
                     {
-                        var connectRequest = new ConnectRequest() { Stream = clientStream, Uri = httpRemoteUri };
 
-                        await TcpConnectionManager.GetClient(connectRequest, httpRemoteUri.Host, httpRemoteUri.Port, true, version).ConfigureAwait(false);
+                        await TcpConnectionManager.GetClient(httpRemoteUri.Host, httpRemoteUri.Port, true, version).ConfigureAwait(false);
+
+                        sslStream = new SslStream(clientStream, true);
+                        var certificate = await CertManager.CreateCertificate(httpRemoteUri.Host, false);
+                        //Successfully managed to authenticate the client using the fake certificate
+                        await sslStream.AuthenticateAsServerAsync(certificate, false,
+                            Constants.SupportedProtocols, false).ConfigureAwait(false);
+                        //HTTPS server created - we can now decrypt the client's traffic
+                        clientStream = sslStream;
 
-                        if (clientStream is SslStream)
-                        {
-                            sslStream = clientStream as SslStream;
-                        }
-                        else
-                        {
-                            sslStream = new SslStream(clientStream, true);
-                            //Successfully managed to authenticate the client using the fake certificate
-                            await sslStream.AuthenticateAsServerAsync(certificate, false,
-                                Constants.SupportedProtocols, false).ConfigureAwait(false);
-                            //HTTPS server created - we can now decrypt the client's traffic
-                            clientStream = sslStream;
-                        }
 
                         clientStreamReader = new CustomBinaryReader(sslStream);
                         clientStreamWriter = new StreamWriter(sslStream);
@@ -192,8 +186,6 @@ await HandleHttpSessionRequest(tcpClient, httpCmd, clientStream, clientStreamRea
         private static async Task HandleHttpSessionRequest(TcpClient client, string httpCmd, Stream clientStream,
             CustomBinaryReader clientStreamReader, StreamWriter clientStreamWriter, bool isHttps)
         {
-            TcpConnection connection = null;
-            string lastRequest = null;
 
             while (true)
             {
@@ -276,12 +268,8 @@ await TcpHelper.SendRaw(clientStream, httpCmd, args.WebSession.Request.RequestHe
                     }
 
                     //construct the web request that we are going to issue on behalf of the client.
-                    connection = connection == null ?
-                       await TcpConnectionManager.GetClient(args.WebSession.Request.RequestUri.Host, args.WebSession.Request.RequestUri.Port, args.IsHttps, version).ConfigureAwait(false)
-                        : lastRequest != args.WebSession.Request.RequestUri.Host ? await TcpConnectionManager.GetClient(args.WebSession.Request.RequestUri.Host, args.WebSession.Request.RequestUri.Port, args.IsHttps, version).ConfigureAwait(false)
-                            : connection;
+                    var connection = await TcpConnectionManager.GetClient(args.WebSession.Request.RequestUri.Host, args.WebSession.Request.RequestUri.Port, args.IsHttps, version).ConfigureAwait(false);
 
-                    lastRequest = TcpConnectionManager.GetConnectionKey(args.WebSession.Request.RequestUri.Host, args.WebSession.Request.RequestUri.Port, args.IsHttps, version);
 
                     args.WebSession.Request.RequestLocked = true;
 
@@ -350,11 +338,12 @@ await WriteResponseStatus(args.WebSession.Response.HttpVersion, "417",
                     //if connection is closing exit
                     if (args.WebSession.Response.ResponseKeepAlive == false)
                     {
-                        connection.TcpClient.Close();
                         Dispose(client, clientStream, clientStreamReader, clientStreamWriter, args);
                         return;
                     }
 
+                   await TcpConnectionManager.ReleaseClient(connection);
+
                     // read the next request 
                     httpCmd = await clientStreamReader.ReadLineAsync().ConfigureAwait(false);
 
@@ -367,8 +356,6 @@ await WriteResponseStatus(args.WebSession.Response.HttpVersion, "417",
 
             }
 
-            if (connection != null)
-                await TcpConnectionManager.ReleaseClient(connection);
         }
 
         private static async Task WriteConnectResponse(StreamWriter clientStreamWriter, Version httpVersion)
@@ -460,8 +447,6 @@ internal static bool ValidateServerCertificate(
           X509Chain chain,
           SslPolicyErrors sslPolicyErrors)
         {
-            var customSslStream = sender as CustomSslStream;
-
             if (ServerCertificateValidationCallback != null)
             {
                 var args = new CertificateValidationEventArgs();
@@ -510,43 +495,25 @@ internal static X509Certificate SelectClientCertificate(
             X509Certificate clientCertificate = null;
             var customSslStream = sender as CustomSslStream;
 
-            if (customSslStream.Param is ConnectRequest && remoteCertificate != null)
-            {
-                var connectRequest = customSslStream.Param as ConnectRequest;
-
-                var sslStream = new SslStream(connectRequest.Stream, true);
-
-                var certificate = CertManager.CreateCertificate(connectRequest.Uri.Host, false).Result;
-                //Successfully managed to authenticate the client using the fake certificate
-                sslStream.AuthenticateAsServerAsync(certificate, true,
-                   Constants.SupportedProtocols, false).Wait();
-
-                connectRequest.Stream = sslStream;
-
-                clientCertificate = sslStream.RemoteCertificate;
-               
-            }
-            else if (customSslStream.Param is ConnectRequest)
+            if (acceptableIssuers != null &&
+                acceptableIssuers.Length > 0 &&
+                localCertificates != null &&
+                localCertificates.Count > 0)
             {
-                if (acceptableIssuers != null &&
-                    acceptableIssuers.Length > 0 &&
-                    localCertificates != null &&
-                    localCertificates.Count > 0)
+                // Use the first certificate that is from an acceptable issuer.
+                foreach (X509Certificate certificate in localCertificates)
                 {
-                    // Use the first certificate that is from an acceptable issuer.
-                    foreach (X509Certificate certificate in localCertificates)
-                    {
-                        string issuer = certificate.Issuer;
-                        if (Array.IndexOf(acceptableIssuers, issuer) != -1)
-                            clientCertificate = certificate;
-                    }
+                    string issuer = certificate.Issuer;
+                    if (Array.IndexOf(acceptableIssuers, issuer) != -1)
+                        clientCertificate = certificate;
                 }
-
-                if (localCertificates != null &&
-                    localCertificates.Count > 0)
-                    clientCertificate = localCertificates[0];
             }
 
+            if (localCertificates != null &&
+                localCertificates.Count > 0)
+                clientCertificate = localCertificates[0];
+
+
             if (ClientCertificateSelectionCallback != null)
             {
                 var args = new CertificateSelectionEventArgs();

Titanium.Web.Proxy/RequestHandler.cs

@@ -62,7 +62,6 @@
     <Compile Include="Decompression\ZlibDecompression.cs" />
     <Compile Include="EventArguments\CertificateSelectionEventArgs.cs" />
     <Compile Include="EventArguments\CertificateValidationEventArgs.cs" />
-    <Compile Include="Models\ConnectRequest.cs" />
     <Compile Include="Network\ProxyClient.cs" />
     <Compile Include="Exceptions\BodyNotFoundException.cs" />
     <Compile Include="Extensions\HttpWebResponseExtensions.cs" />