Merge pull request #665 from justcoding121/beta

Stable

Author: honfika

src/Titanium.Web.Proxy/EventArguments/LimitedStream.cs

@@ -3,6 +3,7 @@
 using System.IO;
 using System.Threading.Tasks;
 using Titanium.Web.Proxy.Exceptions;
+using Titanium.Web.Proxy.Helpers;
 using Titanium.Web.Proxy.StreamExtended.BufferPool;
 using Titanium.Web.Proxy.StreamExtended.Network;
 
@@ -11,16 +12,16 @@ namespace Titanium.Web.Proxy.EventArguments
     internal class LimitedStream : Stream
     {
         private readonly IBufferPool bufferPool;
-        private readonly CustomBufferedStream baseStream;
+        private readonly IHttpStreamReader baseReader;
         private readonly bool isChunked;
         private long bytesRemaining;
 
         private bool readChunkTrail;
 
-        internal LimitedStream(CustomBufferedStream baseStream, IBufferPool bufferPool, bool isChunked,
+        internal LimitedStream(IHttpStreamReader baseStream, IBufferPool bufferPool, bool isChunked,
             long contentLength)
         {  
-            this.baseStream = baseStream;
+            this.baseReader = baseStream;
             this.bufferPool = bufferPool;
             this.isChunked = isChunked;
             bytesRemaining = isChunked
@@ -49,12 +50,12 @@ private void getNextChunk()
             if (readChunkTrail)
             {
                 // read the chunk trail of the previous chunk
-                string? s = baseStream.ReadLineAsync().Result;
+                string? s = baseReader.ReadLineAsync().Result;
             }
 
             readChunkTrail = true;
 
-            string? chunkHead = baseStream.ReadLineAsync().Result!;
+            string? chunkHead = baseReader.ReadLineAsync().Result!;
             int idx = chunkHead.IndexOf(";", StringComparison.Ordinal);
             if (idx >= 0)
             {
@@ -73,7 +74,7 @@ private void getNextChunk()
                 bytesRemaining = -1;
 
                 // chunk trail
-                var task = baseStream.ReadLineAsync();
+                var task = baseReader.ReadLineAsync();
                 if (!task.IsCompleted)
                     task.AsTask().Wait();
             }
@@ -119,7 +120,7 @@ public override int Read(byte[] buffer, int offset, int count)
             }
 
             int toRead = (int)Math.Min(count, bytesRemaining);
-            int res = baseStream.Read(buffer, offset, toRead);
+            int res = baseReader.Read(buffer, offset, toRead);
             bytesRemaining -= res;
 
             if (res == 0)

src/Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs

@@ -10,6 +10,7 @@
 using Titanium.Web.Proxy.Http.Responses;
 using Titanium.Web.Proxy.Models;
 using Titanium.Web.Proxy.Network;
+using Titanium.Web.Proxy.Network.Tcp;
 using Titanium.Web.Proxy.StreamExtended.Network;
 
 namespace Titanium.Web.Proxy.EventArguments
@@ -22,8 +23,6 @@ namespace Titanium.Web.Proxy.EventArguments
     /// </summary>
     public class SessionEventArgs : SessionEventArgsBase
     {
-        private static readonly byte[] emptyData = new byte[0];
-
         /// <summary>
         /// Backing field for corresponding public property
         /// </summary>
@@ -37,8 +36,8 @@ public class SessionEventArgs : SessionEventArgsBase
         /// <summary>
         /// Constructor to initialize the proxy
         /// </summary>
-        internal SessionEventArgs(ProxyServer server, ProxyEndPoint endPoint, ProxyClient proxyClient, ConnectRequest? connectRequest, CancellationTokenSource cancellationTokenSource)
-            : base(server, endPoint, proxyClient, connectRequest, new Request(), cancellationTokenSource)
+        internal SessionEventArgs(ProxyServer server, ProxyEndPoint endPoint, TcpClientConnection clientConnection, HttpClientStream clientStream, ConnectRequest? connectRequest, CancellationTokenSource cancellationTokenSource)
+            : base(server, endPoint, clientConnection, clientStream, connectRequest, new Request(), cancellationTokenSource)
         {
         }
 
@@ -66,14 +65,9 @@ public bool ReRequest
         /// </summary>
         public event EventHandler<MultipartRequestPartSentEventArgs>? MultipartRequestPartSent;
 
-        private CustomBufferedStream getStreamReader(bool isRequest)
-        {
-            return isRequest ? ProxyClient.ClientStream : HttpClient.Connection.Stream;
-        }
-
-        private HttpWriter getStreamWriter(bool isRequest)
+        private HttpStream getStream(bool isRequest)
         {
-            return isRequest ? (HttpWriter)ProxyClient.ClientStreamWriter : HttpClient.Connection.StreamWriter;
+            return isRequest ? (HttpStream)ClientStream : HttpClient.Connection.Stream;
         }
 
         /// <summary>
@@ -110,7 +104,10 @@ private async Task readRequestBodyAsync(CancellationToken cancellationToken)
                 else
                 {
                     var body = await readBodyAsync(true, cancellationToken);
-                    request.Body = body;
+                    if (!request.BodyAvailable)
+                    {
+                        request.Body = body;
+                    }
 
                     // Now set the flag to true
                     // So that next time we can deliver body from cache
@@ -182,7 +179,10 @@ private async Task readResponseBodyAsync(CancellationToken cancellationToken)
                 else
                 {
                     var body = await readBodyAsync(false, cancellationToken);
-                    response.Body = body;
+                    if (!response.BodyAvailable)
+                    {
+                        response.Body = body;
+                    }
 
                     // Now set the flag to true
                     // So that next time we can deliver body from cache
@@ -193,21 +193,19 @@ private async Task readResponseBodyAsync(CancellationToken cancellationToken)
 
         private async Task<byte[]> readBodyAsync(bool isRequest, CancellationToken cancellationToken)
         {
-            using (var bodyStream = new MemoryStream())
-            {
-                var writer = new HttpWriter(bodyStream, BufferPool);
-
-                if (isRequest)
-                {
-                    await CopyRequestBodyAsync(writer, TransformationMode.Uncompress, cancellationToken);
-                }
-                else
-                {
-                    await CopyResponseBodyAsync(writer, TransformationMode.Uncompress, cancellationToken);
-                }
+            using var bodyStream = new MemoryStream();
+            using var http = new HttpStream(bodyStream, BufferPool);
 
-                return bodyStream.ToArray();
+            if (isRequest)
+            {
+                await CopyRequestBodyAsync(http, TransformationMode.Uncompress, cancellationToken);
+            }
+            else
+            {
+                await CopyResponseBodyAsync(http, TransformationMode.Uncompress, cancellationToken);
             }
+
+            return bodyStream.ToArray();
         }
 
         /// <summary>
@@ -225,18 +223,16 @@ internal async Task SyphonOutBodyAsync(bool isRequest, CancellationToken cancell
                 return;
             }
 
-            using (var bodyStream = new MemoryStream())
-            {
-                var writer = new HttpWriter(bodyStream, BufferPool);
-                await copyBodyAsync(isRequest, true, writer, TransformationMode.None, null, cancellationToken);
-            }
+            using var bodyStream = new MemoryStream();
+            using var http = new HttpStream(bodyStream, BufferPool);
+            await copyBodyAsync(isRequest, true, http, TransformationMode.None, null, cancellationToken);
         }
 
         /// <summary>
         ///  This is called when the request is PUT/POST/PATCH to read the body
         /// </summary>
         /// <returns></returns>
-        internal async Task CopyRequestBodyAsync(HttpWriter writer, TransformationMode transformation, CancellationToken cancellationToken)
+        internal async Task CopyRequestBodyAsync(IHttpStreamWriter writer, TransformationMode transformation, CancellationToken cancellationToken)
         {
             var request = HttpClient.Request;
 
@@ -245,7 +241,7 @@ internal async Task CopyRequestBodyAsync(HttpWriter writer, TransformationMode t
             // send the request body bytes to server
             if (contentLength > 0 && hasMulipartEventSubscribers && request.IsMultipartFormData)
             {
-                var reader = getStreamReader(true);
+                var reader = getStream(true);
                 var boundary = HttpHelper.GetBoundaryFromContentType(request.ContentType);
 
                 using (var copyStream = new CopyStream(reader, writer, BufferPool))
@@ -275,14 +271,14 @@ internal async Task CopyRequestBodyAsync(HttpWriter writer, TransformationMode t
             }
         }
 
-        internal async Task CopyResponseBodyAsync(HttpWriter writer, TransformationMode transformation, CancellationToken cancellationToken)
+        internal async Task CopyResponseBodyAsync(IHttpStreamWriter writer, TransformationMode transformation, CancellationToken cancellationToken)
         {
             await copyBodyAsync(false, false, writer, transformation, OnDataReceived, cancellationToken);
         }
 
-        private async Task copyBodyAsync(bool isRequest, bool useOriginalHeaderValues, HttpWriter writer, TransformationMode transformation, Action<byte[], int, int>? onCopy, CancellationToken cancellationToken)
+        private async Task copyBodyAsync(bool isRequest, bool useOriginalHeaderValues, IHttpStreamWriter writer, TransformationMode transformation, Action<byte[], int, int>? onCopy, CancellationToken cancellationToken)
         {
-            var stream = getStreamReader(isRequest);
+            var stream = getStream(isRequest);
 
             var requestResponse = isRequest ? (RequestResponseBase)HttpClient.Request : HttpClient.Response;
 
@@ -309,10 +305,8 @@ private async Task copyBodyAsync(bool isRequest, bool useOriginalHeaderValues, H
 
             try
             {
-                using (var bufStream = new CustomBufferedStream(s, BufferPool, true))
-                {
-                    await writer.CopyBodyAsync(bufStream, false, -1, onCopy, cancellationToken);
-                }
+                var http = new HttpStream(s, BufferPool, true);
+                await writer.CopyBodyAsync(http, false, -1, onCopy, cancellationToken);
             }
             finally
             {
@@ -595,7 +589,7 @@ public void Redirect(string url, bool closeServerConnection = false)
             var response = new RedirectResponse();
             response.HttpVersion = HttpClient.Request.HttpVersion;
             response.Headers.AddHeader(KnownHeaders.Location, url);
-            response.Body = emptyData;
+            response.Body = Array.Empty<byte>();
 
             Respond(response, closeServerConnection);
         }

src/Titanium.Web.Proxy/EventArguments/SessionEventArgsBase.cs

@@ -26,7 +26,12 @@ public abstract class SessionEventArgsBase : EventArgs, IDisposable
 
         internal TcpServerConnection ServerConnection => HttpClient.Connection;
 
-        internal TcpClientConnection ClientConnection => ProxyClient.Connection;
+        /// <summary>
+        ///     Holds a reference to client
+        /// </summary>
+        internal TcpClientConnection ClientConnection { get; }
+
+        internal HttpClientStream ClientStream { get; }
 
         protected readonly IBufferPool BufferPool;
         protected readonly ExceptionHandler ExceptionFunc;
@@ -41,25 +46,21 @@ public abstract class SessionEventArgsBase : EventArgs, IDisposable
         ///     Initializes a new instance of the <see cref="SessionEventArgsBase" /> class.
         /// </summary>
         private protected SessionEventArgsBase(ProxyServer server, ProxyEndPoint endPoint,
-            ProxyClient proxyClient, ConnectRequest? connectRequest, Request request, CancellationTokenSource cancellationTokenSource)
+            TcpClientConnection clientConnection, HttpClientStream clientStream, ConnectRequest? connectRequest, Request request, CancellationTokenSource cancellationTokenSource)
         {
             BufferPool = server.BufferPool;
             ExceptionFunc = server.ExceptionFunc;
             TimeLine["Session Created"] = DateTime.Now;
 
             CancellationTokenSource = cancellationTokenSource;
 
-            ProxyClient = proxyClient;
-            HttpClient = new HttpWebClient(connectRequest, request, new Lazy<int>(() => ProxyClient.Connection.GetProcessId(endPoint)));
+            ClientConnection = clientConnection;
+            ClientStream = clientStream;
+            HttpClient = new HttpWebClient(connectRequest, request, new Lazy<int>(() => clientConnection.GetProcessId(endPoint)));
             LocalEndPoint = endPoint;
             EnableWinAuth = server.EnableWinAuth && isWindowsAuthenticationSupported;
         }
 
-        /// <summary>
-        ///     Holds a reference to client
-        /// </summary>
-        internal ProxyClient ProxyClient { get; }
-
         /// <summary>
         ///     Returns a user data for this request/response session which is
         ///     same as the user data of HttpClient.
@@ -93,7 +94,7 @@ public bool EnableWinAuth
         /// <summary>
         ///     Client End Point.
         /// </summary>
-        public IPEndPoint ClientEndPoint => (IPEndPoint)ProxyClient.Connection.RemoteEndPoint;
+        public IPEndPoint ClientEndPoint => (IPEndPoint)ClientConnection.RemoteEndPoint;
 
         /// <summary>
         ///    The web client used to communicate with server for this session.
@@ -106,7 +107,7 @@ public bool EnableWinAuth
         /// <summary>
         ///     Are we using a custom upstream HTTP(S) proxy?
         /// </summary>
-        public ExternalProxy? CustomUpStreamProxyUsed { get; internal set; }
+        public IExternalProxy? CustomUpStreamProxyUsed { get; internal set; }
 
         /// <summary>
         ///     Local endpoint via which we make the request.

src/Titanium.Web.Proxy/EventArguments/TunnelConnectEventArgs.cs

@@ -1,8 +1,10 @@
 using System;
 using System.Threading;
+using Titanium.Web.Proxy.Helpers;
 using Titanium.Web.Proxy.Http;
 using Titanium.Web.Proxy.Models;
 using Titanium.Web.Proxy.Network;
+using Titanium.Web.Proxy.Network.Tcp;
 using Titanium.Web.Proxy.StreamExtended.Network;
 
 namespace Titanium.Web.Proxy.EventArguments
@@ -15,8 +17,8 @@ public class TunnelConnectSessionEventArgs : SessionEventArgsBase
         private bool? isHttpsConnect;
 
         internal TunnelConnectSessionEventArgs(ProxyServer server, ProxyEndPoint endPoint, ConnectRequest connectRequest,
-            ProxyClient proxyClient, CancellationTokenSource cancellationTokenSource)
-            : base(server, endPoint, proxyClient, connectRequest, connectRequest, cancellationTokenSource)
+            TcpClientConnection clientConnection, HttpClientStream clientStream, CancellationTokenSource cancellationTokenSource)
+            : base(server, endPoint, clientConnection, clientStream, connectRequest, connectRequest, cancellationTokenSource)
         {
         }