Merge pull request #395 from justcoding121/develop

merge to beta

Author: honfika

Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs

@@ -107,15 +107,18 @@ public bool ReRequest
         /// <summary>
         /// Constructor to initialize the proxy
         /// </summary>
-        internal SessionEventArgs(int bufferSize,
-            ProxyEndPoint endPoint,
-            ExceptionHandler exceptionFunc)
+        internal SessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc)
+            : this(bufferSize, endPoint, exceptionFunc, null)
+        {
+        }
+
+        protected SessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ExceptionHandler exceptionFunc, Request request)
         {
             this.bufferSize = bufferSize;
             this.exceptionFunc = exceptionFunc;
 
             ProxyClient = new ProxyClient();
-            WebSession = new HttpWebClient(bufferSize);
+            WebSession = new HttpWebClient(bufferSize, request);
             LocalEndPoint = endPoint;
 
             WebSession.ProcessId = new Lazy<int>(() =>
@@ -138,6 +141,21 @@ internal SessionEventArgs(int bufferSize,
             });
         }
 
+        private CustomBufferedStream GetStream(bool isRequest)
+        {
+            return isRequest ? ProxyClient.ClientStream : WebSession.ServerConnection.Stream;
+        }
+
+        private CustomBinaryReader GetStreamReader(bool isRequest)
+        {
+            return isRequest ? ProxyClient.ClientStreamReader : WebSession.ServerConnection.StreamReader;
+        }
+
+        private HttpWriter GetStreamWriter(bool isRequest)
+        {
+            return isRequest ? (HttpWriter)ProxyClient.ClientStreamWriter : WebSession.ServerConnection.StreamWriter;
+        }
+
         /// <summary>
         /// Read request body content as bytes[] for current session
         /// </summary>
@@ -150,7 +168,7 @@ private async Task ReadRequestBodyAsync()
             //If not already read (not cached yet)
             if (!request.IsBodyRead)
             {
-                var body = await ReadBodyAsync(ProxyClient.ClientStreamReader, true);
+                var body = await ReadBodyAsync(true);
                 request.Body = body;
 
                 //Now set the flag to true
@@ -165,8 +183,8 @@ private async Task ReadRequestBodyAsync()
         /// </summary>
         internal async Task ClearResponse()
         {
-            //siphon out the body
-            await ReadResponseBodyAsync();
+            //syphon out the response body from server
+            await SyphonOutBodyAsync(false);
             WebSession.Response = new Response();
         }
 
@@ -225,7 +243,7 @@ private async Task ReadResponseBodyAsync()
             //If not already read (not cached yet)
             if (!response.IsBodyRead)
             {
-                var body = await ReadBodyAsync(WebSession.ServerConnection.StreamReader, false);
+                var body = await ReadBodyAsync(false);
                 response.Body = body;
 
                 //Now set the flag to true
@@ -235,11 +253,12 @@ private async Task ReadResponseBodyAsync()
             }
         }
 
-        private async Task<byte[]> ReadBodyAsync(CustomBinaryReader reader, bool isRequest)
+        private async Task<byte[]> ReadBodyAsync(bool isRequest)
         {
             using (var bodyStream = new MemoryStream())
             {
                 var writer = new HttpWriter(bodyStream, bufferSize);
+
                 if (isRequest)
                 {
                     await CopyRequestBodyAsync(writer, TransformationMode.Uncompress);
@@ -253,21 +272,35 @@ private async Task<byte[]> ReadBodyAsync(CustomBinaryReader reader, bool isReque
             }
         }
 
+        internal async Task SyphonOutBodyAsync(bool isRequest)
+        {
+            var requestResponse = isRequest ? (RequestResponseBase)WebSession.Request : WebSession.Response;
+            if (requestResponse.IsBodyRead || !requestResponse.OriginalHasBody)
+            {
+                return;
+            }
+
+            using (var bodyStream = new MemoryStream())
+            {
+                var writer = new HttpWriter(bodyStream, bufferSize);
+                await CopyBodyAsync(isRequest, writer, TransformationMode.None, null);
+            }
+        }
+
         /// <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)
         {
-            // End the operation
             var request = WebSession.Request;
-            var reader = ProxyClient.ClientStreamReader;
 
             long contentLength = request.ContentLength;
 
             //send the request body bytes to server
             if (contentLength > 0 && hasMulipartEventSubscribers && request.IsMultipartFormData)
             {
+                var reader = GetStreamReader(true);
                 string boundary = HttpHelper.GetBoundaryFromContentType(request.ContentType);
 
                 using (var copyStream = new CopyStream(reader, writer, bufferSize))
@@ -294,21 +327,22 @@ internal async Task CopyRequestBodyAsync(HttpWriter writer, TransformationMode t
             }
             else
             {
-                await CopyBodyAsync(ProxyClient.ClientStream, reader, writer, request, transformation, OnDataSent);
+                await CopyBodyAsync(true, writer, transformation, OnDataSent);
             }
         }
 
         internal async Task CopyResponseBodyAsync(HttpWriter writer, TransformationMode transformation)
         {
-            var response = WebSession.Response;
-            var reader = WebSession.ServerConnection.StreamReader;
-
-            await CopyBodyAsync(WebSession.ServerConnection.Stream, reader, writer, response, transformation, OnDataReceived);
+            await CopyBodyAsync(false, writer, transformation, OnDataReceived);
         }
 
-        private async Task CopyBodyAsync(CustomBufferedStream stream, CustomBinaryReader reader, HttpWriter writer,
-            RequestResponseBase requestResponse, TransformationMode transformation, Action<byte[], int, int> onCopy)
+        private async Task CopyBodyAsync(bool isRequest, HttpWriter writer, TransformationMode transformation, Action<byte[], int, int> onCopy)
         {
+            var stream = GetStream(isRequest);
+            var reader = GetStreamReader(isRequest);
+
+            var requestResponse = isRequest ? (RequestResponseBase)WebSession.Request : WebSession.Response;
+
             bool isChunked = requestResponse.IsChunked;
             long contentLength = requestResponse.ContentLength;
             if (transformation == TransformationMode.None)
@@ -440,42 +474,29 @@ public async Task<string> GetRequestBodyAsString()
         /// Sets the request body
         /// </summary>
         /// <param name="body"></param>
-        public async Task SetRequestBody(byte[] body)
+        public void SetRequestBody(byte[] body)
         {
             var request = WebSession.Request;
             if (request.Locked)
             {
                 throw new Exception("You cannot call this function after request is made to server.");
             }
 
-            //syphon out the request body from client before setting the new body
-            if (!request.IsBodyRead)
-            {
-                await ReadRequestBodyAsync();
-            }
-
             request.Body = body;
-            request.UpdateContentLength();
         }
 
         /// <summary>
         /// Sets the body with the specified string
         /// </summary>
         /// <param name="body"></param>
-        public async Task SetRequestBodyString(string body)
+        public void SetRequestBodyString(string body)
         {
             if (WebSession.Request.Locked)
             {
                 throw new Exception("You cannot call this function after request is made to server.");
             }
 
-            //syphon out the request body from client before setting the new body
-            if (!WebSession.Request.IsBodyRead)
-            {
-                await ReadRequestBodyAsync();
-            }
-
-            await SetRequestBody(WebSession.Request.Encoding.GetBytes(body));
+            SetRequestBody(WebSession.Request.Encoding.GetBytes(body));
         }
 
         /// <summary>
@@ -510,47 +531,31 @@ public async Task<string> GetResponseBodyAsString()
         /// Set the response body bytes
         /// </summary>
         /// <param name="body"></param>
-        public async Task SetResponseBody(byte[] body)
+        public void SetResponseBody(byte[] body)
         {
             if (!WebSession.Request.Locked)
             {
                 throw new Exception("You cannot call this function before request is made to server.");
             }
 
             var response = WebSession.Response;
-
-            //syphon out the response body from server before setting the new body
-            if (response.Body == null)
-            {
-                await GetResponseBody();
-            }
-
             response.Body = body;
-
-            //If there is a content length header update it
-            response.UpdateContentLength();
         }
 
         /// <summary>
         /// Replace the response body with the specified string
         /// </summary>
         /// <param name="body"></param>
-        public async Task SetResponseBodyString(string body)
+        public void SetResponseBodyString(string body)
         {
             if (!WebSession.Request.Locked)
             {
                 throw new Exception("You cannot call this function before request is made to server.");
             }
 
-            //syphon out the response body from server before setting the new body
-            if (!WebSession.Response.IsBodyRead)
-            {
-                await GetResponseBody();
-            }
-
             var bodyBytes = WebSession.Response.Encoding.GetBytes(body);
 
-            await SetResponseBody(bodyBytes);
+            SetResponseBody(bodyBytes);
         }
 
         /// <summary>
@@ -651,17 +656,29 @@ public void Respond(Response response)
         {
             if (WebSession.Request.Locked)
             {
-                throw new Exception("You cannot call this function after request is made to server.");
-            }
+                if (WebSession.Response.Locked)
+                {
+                    throw new Exception("You cannot call this function after response is sent to the client.");
+                }
 
-            WebSession.Request.Locked = true;
+                response.Locked = true;
+                response.TerminateResponse = WebSession.Response.TerminateResponse;
+                WebSession.Response = response;
+            }
+            else
+            {
+                WebSession.Request.Locked = true;
 
-            response.Locked = true;
-            response.IsBodyRead = true;
+                response.Locked = true;
+                WebSession.Response = response;
 
-            WebSession.Response = response;
+                WebSession.Request.CancelRequest = true;
+            }
+        }
 
-            WebSession.Request.CancelRequest = true;
+        public void TerminateServerConnection()
+        {
+            WebSession.Response.TerminateResponse = true;
         }
 
         /// <summary>

Titanium.Web.Proxy/EventArguments/TunnelConnectEventArgs.cs

@@ -11,9 +11,8 @@ public class TunnelConnectSessionEventArgs : SessionEventArgs
         public bool IsHttpsConnect { get; internal set; }
 
         internal TunnelConnectSessionEventArgs(int bufferSize, ProxyEndPoint endPoint, ConnectRequest connectRequest, ExceptionHandler exceptionFunc) 
-            : base(bufferSize, endPoint, exceptionFunc)
+            : base(bufferSize, endPoint, exceptionFunc, connectRequest)
         {
-            WebSession.Request = connectRequest;
         }
     }
 }

Titanium.Web.Proxy/Helpers/HttpResponseWriter.cs

@@ -20,7 +20,7 @@ public HttpResponseWriter(Stream stream, int bufferSize) : base(stream, bufferSi
         public async Task WriteResponseAsync(Response response, bool flush = true)
         {
             await WriteResponseStatusAsync(response.HttpVersion, response.StatusCode, response.StatusDescription);
-            await WriteHeadersAsync(response.Headers, flush);
+            await WriteAsync(response, flush);
         }
 
         /// <summary>

Titanium.Web.Proxy/Helpers/HttpWriter.cs

@@ -5,7 +5,7 @@
 using System.Threading.Tasks;
 using StreamExtended.Helpers;
 using StreamExtended.Network;
-using Titanium.Web.Proxy.EventArguments;
+using Titanium.Web.Proxy.Compression;
 using Titanium.Web.Proxy.Http;
 using Titanium.Web.Proxy.Shared;
 
@@ -94,12 +94,9 @@ public Task WriteLineAsync(string value)
         /// <returns></returns>
         public async Task WriteHeadersAsync(HeaderCollection headers, bool flush = true)
         {
-            if (headers != null)
+            foreach (var header in headers)
             {
-                foreach (var header in headers)
-                {
-                    await header.WriteToStreamAsync(this);
-                }
+                await header.WriteToStreamAsync(this);
             }
 
             await WriteLineAsync();
@@ -265,5 +262,22 @@ private async Task CopyBytesFromStream(CustomBinaryReader reader, long count, Ac
                 onCopy?.Invoke(buffer, 0, bytesRead);
             }
         }
+
+        /// <summary>
+        /// Writes the request/response headers and body.
+        /// </summary>
+        /// <param name="requestResponse"></param>
+        /// <param name="flush"></param>
+        /// <returns></returns>
+        protected async Task WriteAsync(RequestResponseBase requestResponse, bool flush = true)
+        {
+            var body = requestResponse.CompressBodyAndUpdateContentLength();
+            await WriteHeadersAsync(requestResponse.Headers, flush);
+
+            if (body != null)
+            {
+                await WriteBodyAsync(body, requestResponse.IsChunked);
+            }
+        }
     }
 }

Titanium.Web.Proxy/Http/HttpWebClient.cs

@@ -38,7 +38,7 @@ public class HttpWebClient
         /// <summary>
         /// Web Request.
         /// </summary>
-        public Request Request { get; internal set; }
+        public Request Request { get; }
 
         /// <summary>
         /// Web Response.
@@ -56,13 +56,13 @@ public class HttpWebClient
         /// </summary>
         public bool IsHttps => Request.IsHttps;
 
-        internal HttpWebClient(int bufferSize)
+        internal HttpWebClient(int bufferSize, Request request = null, Response response = null)
         {
             this.bufferSize = bufferSize;
 
             RequestId = Guid.NewGuid();
-            Request = new Request();
-            Response = new Response();
+            Request = request ?? new Request();
+            Response = response ?? new Response();
         }
 
         /// <summary>