#823 Pass SessionEventArgs to HttpStream

justcoding121 · Apr 23, 2021 · 285e023 · 285e023
1 parent ad04876
commit 285e023
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 15 deletions.
diff --git a/src/Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs b/src/Titanium.Web.Proxy/EventArguments/SessionEventArgs.cs
@@ -228,7 +228,7 @@ internal async Task SyphonOutBodyAsync(bool isRequest, CancellationToken cancell
 
             var reader = isRequest ? (HttpStream)ClientStream : HttpClient.Connection.Stream;
 
-            await reader.CopyBodyAsync(requestResponse, true, NullWriter.Instance, TransformationMode.None, null, cancellationToken);
+            await reader.CopyBodyAsync(requestResponse, true, NullWriter.Instance, TransformationMode.None, isRequest, this, cancellationToken);
         }
 
         /// <summary>
@@ -270,13 +270,13 @@ internal async Task CopyRequestBodyAsync(IHttpStreamWriter writer, Transformatio
             }
             else
             {
-                await reader.CopyBodyAsync(request, false, writer, transformation, OnDataSent, cancellationToken);
+                await reader.CopyBodyAsync(request, false, writer, transformation, true, this, cancellationToken);
             }
         }
 
         private async Task copyResponseBodyAsync(IHttpStreamWriter writer, TransformationMode transformation, CancellationToken cancellationToken)
         {
-            await HttpClient.Connection.Stream.CopyBodyAsync(HttpClient.Response, false, writer, transformation, OnDataReceived, cancellationToken);
+            await HttpClient.Connection.Stream.CopyBodyAsync(HttpClient.Response, false, writer, transformation, false, this, cancellationToken);
         }
 
         /// <summary>

diff --git a/src/Titanium.Web.Proxy/Helpers/HttpStream.cs b/src/Titanium.Web.Proxy/Helpers/HttpStream.cs
@@ -1003,14 +1003,14 @@ internal ValueTask WriteBodyAsync(byte[] data, bool isChunked, CancellationToken
             return WriteAsync(data, cancellationToken: cancellationToken);
         }
 
-        public async Task CopyBodyAsync(RequestResponseBase requestResponse, bool useOriginalHeaderValues, IHttpStreamWriter writer, TransformationMode transformation, Action<byte[], int, int>? onCopy, CancellationToken cancellationToken)
+        public async Task CopyBodyAsync(RequestResponseBase requestResponse, bool useOriginalHeaderValues, IHttpStreamWriter writer, TransformationMode transformation, bool isRequest, SessionEventArgs args, CancellationToken cancellationToken)
         {
             bool isChunked = useOriginalHeaderValues ? requestResponse.OriginalIsChunked : requestResponse.IsChunked;
             long contentLength = useOriginalHeaderValues ? requestResponse.OriginalContentLength : requestResponse.ContentLength;
 
             if (transformation == TransformationMode.None)
             {
-                await CopyBodyAsync(writer, isChunked, contentLength, onCopy, cancellationToken);
+                await CopyBodyAsync(writer, isChunked, contentLength, isRequest, args, cancellationToken);
                 return;
             }
 
@@ -1029,7 +1029,7 @@ public async Task CopyBodyAsync(RequestResponseBase requestResponse, bool useOri
             try
             {
                 var http = new HttpStream(server, s, bufferPool, cancellationToken, true);
-                await http.CopyBodyAsync(writer, false, -1, onCopy, cancellationToken);
+                await http.CopyBodyAsync(writer, false, -1, isRequest, args, cancellationToken);
             }
             finally
             {
@@ -1051,12 +1051,13 @@ public async Task CopyBodyAsync(RequestResponseBase requestResponse, bool useOri
         /// <param name="cancellationToken"></param>
         /// <returns></returns>
         public Task CopyBodyAsync(IHttpStreamWriter writer, bool isChunked, long contentLength,
-            Action<byte[], int, int>? onCopy, CancellationToken cancellationToken)
+            bool isRequest,
+            SessionEventArgs args, CancellationToken cancellationToken)
         {
             // For chunked request we need to read data as they arrive, until we reach a chunk end symbol
             if (isChunked)
             {
-                return copyBodyChunkedAsync(writer, onCopy, cancellationToken);
+                return copyBodyChunkedAsync(writer, isRequest, args, cancellationToken);
             }
 
             // http 1.0 or the stream reader limits the stream
@@ -1066,7 +1067,7 @@ public Task CopyBodyAsync(IHttpStreamWriter writer, bool isChunked, long content
             }
 
             // If not chunked then its easy just read the amount of bytes mentioned in content length header
-            return copyBytesToStream(writer, contentLength, onCopy, cancellationToken);
+            return copyBytesToStream(writer, contentLength, isRequest, args, cancellationToken);
         }
 
         /// <summary>
@@ -1095,7 +1096,7 @@ private async ValueTask writeBodyChunkedAsync(byte[] data, CancellationToken can
         /// <param name="onCopy"></param>
         /// <param name="cancellationToken"></param>
         /// <returns></returns>
-        private async Task copyBodyChunkedAsync(IHttpStreamWriter writer, Action<byte[], int, int>? onCopy, CancellationToken cancellationToken)
+        private async Task copyBodyChunkedAsync(IHttpStreamWriter writer, bool isRequest, SessionEventArgs args, CancellationToken cancellationToken)
         {
             while (true)
             {
@@ -1120,7 +1121,7 @@ private async Task copyBodyChunkedAsync(IHttpStreamWriter writer, Action<byte[],
 
                 if (chunkSize != 0)
                 {
-                    await copyBytesToStream(writer, chunkSize, onCopy, cancellationToken);
+                    await copyBytesToStream(writer, chunkSize, isRequest, args, cancellationToken);
                 }
 
                 await writer.WriteLineAsync(cancellationToken);
@@ -1143,7 +1144,7 @@ private async Task copyBodyChunkedAsync(IHttpStreamWriter writer, Action<byte[],
         /// <param name="onCopy"></param>
         /// <param name="cancellationToken"></param>
         /// <returns></returns>
-        private async Task copyBytesToStream(IHttpStreamWriter writer, long count, Action<byte[], int, int>? onCopy,
+        private async Task copyBytesToStream(IHttpStreamWriter writer, long count, bool isRequest, SessionEventArgs args,
             CancellationToken cancellationToken)
         {
             var buffer = bufferPool.GetBuffer();
@@ -1170,7 +1171,10 @@ private async Task copyBytesToStream(IHttpStreamWriter writer, long count, Actio
 
                     await writer.WriteAsync(buffer, 0, bytesRead, cancellationToken);
 
-                    onCopy?.Invoke(buffer, 0, bytesRead);
+                    if (isRequest)
+                        args.OnDataSent(buffer, 0, bytesRead);
+                    else
+                        args.OnDataReceived(buffer, 0, bytesRead);
                 }
             }
             finally

diff --git a/src/Titanium.Web.Proxy/ResponseHandler.cs b/src/Titanium.Web.Proxy/ResponseHandler.cs
@@ -118,7 +118,7 @@ await handleHttpSessionRequest(args, null, args.ClientConnection.NegotiatedAppli
                     // Copy body if exists
                     var serverStream = args.HttpClient.Connection.Stream;
                     await serverStream.CopyBodyAsync(response, false, clientStream, TransformationMode.None,
-                        args.OnDataReceived, cancellationToken);
+                        false, args, cancellationToken);
                 }
             }
 

diff --git a/src/Titanium.Web.Proxy/StreamExtended/Network/IHttpStreamReader.cs b/src/Titanium.Web.Proxy/StreamExtended/Network/IHttpStreamReader.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Titanium.Web.Proxy.EventArguments;
 
 namespace Titanium.Web.Proxy.StreamExtended.Network
 {
@@ -11,6 +12,6 @@ public interface IHttpStreamReader : ILineStream
         Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken);
 
         Task CopyBodyAsync(IHttpStreamWriter writer, bool isChunked, long contentLength,
-            Action<byte[], int, int>? onCopy, CancellationToken cancellationToken);
+           bool isRequest, SessionEventArgs args, CancellationToken cancellationToken);
     }
 }