Merge pull request #449 from justcoding121/master

Beta

Author: justcoding121

Examples/Titanium.Web.Proxy.Examples.Basic/ProxyTestController.cs

@@ -149,8 +149,9 @@ private async Task OnBeforeTunnelConnectRequest(object sender, TunnelConnectSess
             }
         }
 
-        private async Task OnBeforeTunnelConnectResponse(object sender, TunnelConnectSessionEventArgs e)
+        private Task OnBeforeTunnelConnectResponse(object sender, TunnelConnectSessionEventArgs e)
         {
+            return Task.FromResult(false);
         }
 
         // intecept & cancel redirect or update requests

Titanium.Web.Proxy/ExplicitClientHandler.cs

@@ -140,8 +140,9 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
                         {
                             // test server HTTP/2 support
                             // todo: this is a hack, because Titanium does not support HTTP protocol changing currently
-                            var connection = await getServerConnection(connectArgs, true,
-                                SslExtensions.Http2ProtocolAsList, true, cancellationToken);
+                            var connection = await tcpConnectionFactory.GetServerConnection(this, connectArgs,
+                                                    isConnect: true, applicationProtocols: SslExtensions.Http2ProtocolAsList,
+                                                    noCache: true, cancellationToken: cancellationToken);
 
                             http2Supported = connection.NegotiatedApplicationProtocol == SslApplicationProtocol.Http2;
 
@@ -153,8 +154,9 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
 
                         //don't pass cancellation token here
                         //it could cause floating server connections when client exits
-                        prefetchConnectionTask = getServerConnection(connectArgs, true,
-                                null, false, CancellationToken.None);
+                        prefetchConnectionTask = tcpConnectionFactory.GetServerConnection(this, connectArgs,
+                                                isConnect: true, applicationProtocols: null, noCache: false,
+                                                cancellationToken: CancellationToken.None);
 
                         try
                         {
@@ -201,6 +203,12 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
                         {
                             decryptSsl = false;
                         }
+
+                        if(!decryptSsl)
+                        {
+                            await tcpConnectionFactory.Release(prefetchConnectionTask, true);
+                            prefetchConnectionTask = null; 
+                        }
                     }
 
                     if (cancellationTokenSource.IsCancellationRequested)
@@ -214,8 +222,9 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.WebSession.Response,
                         // create new connection to server.
                         // If we detected that client tunnel CONNECTs without SSL by checking for empty client hello then 
                         // this connection should not be HTTPS.
-                        var connection = await getServerConnection(connectArgs,
-                            true, SslExtensions.Http2ProtocolAsList, true, cancellationToken);
+                        var connection = await tcpConnectionFactory.GetServerConnection(this, connectArgs,
+                                                        isConnect: true, applicationProtocols: SslExtensions.Http2ProtocolAsList,
+                                                        noCache: true, cancellationToken: cancellationToken);
 
                         try
                         {
@@ -247,8 +256,6 @@ await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
                                 (buffer, offset, count) => { connectArgs.OnDataSent(buffer, offset, count); },
                                 (buffer, offset, count) => { connectArgs.OnDataReceived(buffer, offset, count); },
                                 connectArgs.CancellationTokenSource, ExceptionFunc);
-
-
                         }
                         finally
                         {
@@ -284,9 +291,9 @@ await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
                             throw new Exception($"HTTP/2 Protocol violation. Empty string expected, '{line}' received");
                         }
 
-                        var connection = await getServerConnection(connectArgs, true,
-                                                        SslExtensions.Http2ProtocolAsList, true,
-                                                        cancellationToken);
+                        var connection = await tcpConnectionFactory.GetServerConnection(this, connectArgs,
+                                                        isConnect: true, applicationProtocols: SslExtensions.Http2ProtocolAsList,
+                                                        noCache: true, cancellationToken: cancellationToken);
                         try
                         {
                             await connection.StreamWriter.WriteLineAsync("PRI * HTTP/2.0", cancellationToken);
@@ -335,19 +342,9 @@ await handleHttpSessionRequest(endPoint, clientConnection, clientStream, clientS
             }
             finally
             {
-                if (!calledRequestHandler
-                        && prefetchConnectionTask != null)
+                if (!calledRequestHandler)
                 {
-                    TcpServerConnection prefetchedConnection = null;
-                    try
-                    {
-                        prefetchedConnection = await prefetchConnectionTask;
-
-                    }
-                    finally
-                    {
-                        await tcpConnectionFactory.Release(prefetchedConnection, closeServerConnection);
-                    }
+                    await tcpConnectionFactory.Release(prefetchConnectionTask, closeServerConnection);
                 }
 
                 clientStream.Dispose();

Titanium.Web.Proxy/Models/ProxyAuthenticationContext.cs

@@ -0,0 +1,53 @@
+namespace Titanium.Web.Proxy.Models
+{
+    public enum ProxyAuthenticationResult
+    {
+        /// <summary>
+        /// Indicates the authentication request was successful
+        /// </summary>
+        Success,
+        /// <summary>
+        /// Indicates the authentication request failed
+        /// </summary>
+        Failure,
+        /// <summary>
+        /// Indicates that this stage of the authentication request succeeded
+        /// And a second pass of the handshake needs to occur
+        /// </summary>
+        ContinuationNeeded
+    }
+
+    /// <summary>
+    /// A context container for authentication flows
+    /// </summary>
+    public class ProxyAuthenticationContext
+    {
+        /// <summary>
+        /// The result of the current authentication request
+        /// </summary>
+        public ProxyAuthenticationResult Result { get; set; }
+
+        /// <summary>
+        /// An optional continuation token to return to the caller if set
+        /// </summary>
+        public string Continuation { get; set; }
+
+        public static ProxyAuthenticationContext Failed()
+        {
+            return new ProxyAuthenticationContext
+            {
+                Result = ProxyAuthenticationResult.Failure,
+                Continuation = null
+            };
+        }
+
+        public static ProxyAuthenticationContext Succeeded()
+        {
+            return new ProxyAuthenticationContext
+            {
+                Result = ProxyAuthenticationResult.Success,
+                Continuation = null
+            };
+        }
+    }
+}

Titanium.Web.Proxy/Network/RetryPolicy.cs

@@ -0,0 +1,91 @@
+using Polly;
+using System;
+using System.Threading.Tasks;
+using Titanium.Web.Proxy.Network.Tcp;
+
+namespace Titanium.Web.Proxy.Network
+{
+    internal class RetryPolicy<T> where T : Exception
+    {
+        private readonly int retries;
+        private readonly TcpConnectionFactory tcpConnectionFactory;
+
+        private TcpServerConnection currentConnection;
+        private Policy policy;
+
+        internal RetryPolicy(int retries, TcpConnectionFactory tcpConnectionFactory)
+        {
+            this.retries = retries;
+            this.tcpConnectionFactory = tcpConnectionFactory;
+
+            policy = getRetryPolicy();
+        }
+
+        /// <summary>
+        ///     Execute and retry the given action until retry number of times.
+        /// </summary>
+        /// <param name="action">The action to retry with return value specifying whether caller should continue execution.</param>
+        /// <param name="generator">The Tcp connection generator to be invoked to get new connection for retry.</param>
+        /// <param name="initialConnection">Initial Tcp connection to use.</param>
+        /// <returns>Returns the latest connection used and the latest exception if any.</returns>
+        internal async Task<RetryResult> ExecuteAsync(Func<TcpServerConnection, Task<bool>> action,
+            Func<Task<TcpServerConnection>> generator, TcpServerConnection initialConnection)
+        {
+            currentConnection = initialConnection;
+            Exception exception = null;
+            bool @continue = true;
+
+            try
+            {
+                //retry on error with polly policy
+                //do not use polly context to store connection; it does not save states b/w attempts
+                await policy.ExecuteAsync(async () =>
+                {
+                    //setup connection
+                    currentConnection = currentConnection as TcpServerConnection ??
+                                      await generator();
+                    //try
+                    @continue = await action(currentConnection);
+
+                });
+            }
+            catch (Exception e) { exception = e; }
+
+            return new RetryResult(currentConnection, exception, @continue);
+        }
+
+        //get the policy
+        private Policy getRetryPolicy()
+        {
+            return Policy.Handle<T>()
+                    .RetryAsync(retries,
+                        onRetryAsync: onRetry);
+        }
+
+        //before retry clear connection
+        private async Task onRetry(Exception ex, int attempt)
+        {
+            if (currentConnection != null)
+            {
+                //close connection on error
+                await tcpConnectionFactory.Release(currentConnection, true);
+                currentConnection = null;
+            }
+        }
+    }
+
+    internal class RetryResult
+    {
+        internal bool IsSuccess => Exception == null;
+        internal TcpServerConnection LatestConnection { get; }
+        internal Exception Exception { get; }
+        internal bool Continue { get; }
+
+        internal RetryResult(TcpServerConnection lastConnection, Exception exception, bool @continue)
+        {
+            LatestConnection = lastConnection;
+            Exception = exception;
+            Continue = @continue;
+        }
+    }
+}