HTTP/2: allow to read the body

examples/Titanium.Web.Proxy.Examples.Wpf/MainWindow.xaml.cs

@@ -36,6 +36,9 @@ public partial class MainWindow : Window
         public MainWindow()
         {
             proxyServer = new ProxyServer();
+
+            proxyServer.EnableHttp2 = true;
+
             //proxyServer.CertificateManager.CertificateEngine = CertificateEngine.DefaultWindows;
 
             ////Set a password for the .pfx file
@@ -150,11 +153,9 @@ private async Task ProxyServer_BeforeRequest(object sender, SessionEventArgs e)
             SessionListItem item = null;
             await Dispatcher.InvokeAsync(() => { item = addSession(e); });
 
-            if (e.HttpClient.ConnectRequest?.TunnelType == TunnelType.Http2)
-            {
-                // GetRequestBody for HTTP/2 currently not supported
-                return;
-            }
+            //if (e.HttpClient.ConnectRequest?.TunnelType == TunnelType.Http2)
+            //{
+            //}
 
             if (e.HttpClient.Request.HasBody)
             {
@@ -174,11 +175,9 @@ await Dispatcher.InvokeAsync(() =>
                 }
             });
 
-            if (e.HttpClient.ConnectRequest?.TunnelType == TunnelType.Http2)
-            {
-                // GetRequestBody for HTTP/2 currently not supported
-                return;
-            }
+            //if (e.HttpClient.ConnectRequest?.TunnelType == TunnelType.Http2)
+            //{
+            //}
 
             if (item != null)
             {

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

@@ -89,13 +89,28 @@ private async Task readRequestBodyAsync(CancellationToken cancellationToken)
             // If not already read (not cached yet)
             if (!request.IsBodyRead)
             {
-                var body = await readBodyAsync(true, cancellationToken);
-                request.Body = body;
+                if (request.HttpVersion == HttpHeader.Version20)
+                {
+                    request.Http2BodyData = new MemoryStream();
+
+                    var tcs = new TaskCompletionSource<bool>();
+                    request.ReadHttp2BodyTaskCompletionSource = tcs;
+
+                    // signal to HTTP/2 copy frame method to continue
+                    request.ReadHttp2BeforeHandlerTaskCompletionSource.SetResult(true);
+
+                    await tcs.Task;
+                }
+                else
+                {
+                    var body = await readBodyAsync(true, cancellationToken);
+                    request.Body = body;
 
-                // Now set the flag to true
-                // So that next time we can deliver body from cache
-                request.IsBodyRead = true;
-                OnDataSent(body, 0, body.Length);
+                    // Now set the flag to true
+                    // So that next time we can deliver body from cache
+                    request.IsBodyRead = true;
+                    OnDataSent(body, 0, body.Length);
+                }
             }
         }
 
@@ -140,13 +155,28 @@ private async Task readResponseBodyAsync(CancellationToken cancellationToken)
             // If not already read (not cached yet)
             if (!response.IsBodyRead)
             {
-                var body = await readBodyAsync(false, cancellationToken);
-                response.Body = body;
+                if (response.HttpVersion == HttpHeader.Version20)
+                {
+                    response.Http2BodyData = new MemoryStream();
+
+                    var tcs = new TaskCompletionSource<bool>();
+                    response.ReadHttp2BodyTaskCompletionSource = tcs;
+
+                    // signal to HTTP/2 copy frame method to continue
+                    response.ReadHttp2BeforeHandlerTaskCompletionSource.SetResult(true);
+
+                    await tcs.Task;
+                }
+                else
+                {
+                    var body = await readBodyAsync(false, cancellationToken);
+                    response.Body = body;
 
-                // Now set the flag to true
-                // So that next time we can deliver body from cache
-                response.IsBodyRead = true;
-                OnDataReceived(body, 0, body.Length);
+                    // Now set the flag to true
+                    // So that next time we can deliver body from cache
+                    response.IsBodyRead = true;
+                    OnDataReceived(body, 0, body.Length);
+                }
             }
         }
 

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

@@ -55,25 +55,7 @@ protected SessionEventArgsBase(ProxyServer server, ProxyEndPoint endPoint,
             HttpClient = new HttpWebClient(request);
             LocalEndPoint = endPoint;
 
-            HttpClient.ProcessId = new Lazy<int>(() =>
-            {
-                if (RunTime.IsWindows)
-                {
-                    var remoteEndPoint = ClientEndPoint;
-
-                    // If client is localhost get the process id
-                    if (NetworkHelper.IsLocalIpAddress(remoteEndPoint.Address))
-                    {
-                        var ipVersion = endPoint.IpV6Enabled ? IpVersion.Ipv6 : IpVersion.Ipv4;
-                        return TcpHelper.GetProcessIdByLocalPort(ipVersion, remoteEndPoint.Port);
-                    }
-
-                    // can't access process Id of remote request from remote machine
-                    return -1;
-                }
-
-                throw new PlatformNotSupportedException();
-            });
+            HttpClient.ProcessId = new Lazy<int>(() => ProxyClient.Connection.GetProcessId(endPoint));
         }
 
         /// <summary>

src/Titanium.Web.Proxy/Http/RequestResponseBase.cs

@@ -3,6 +3,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Text;
+using System.Threading.Tasks;
 using Titanium.Web.Proxy.Compression;
 using Titanium.Web.Proxy.Extensions;
 using Titanium.Web.Proxy.Helpers;
@@ -49,6 +50,12 @@ public abstract class RequestResponseBase
         /// </summary>
         internal string OriginalContentEncoding { get; set; }
 
+        internal TaskCompletionSource<bool> ReadHttp2BeforeHandlerTaskCompletionSource;
+
+        internal TaskCompletionSource<bool> ReadHttp2BodyTaskCompletionSource;
+
+        internal MemoryStream Http2BodyData;
+
         /// <summary>
         ///     Keeps the body data after the session is finished.
         /// </summary>

src/Titanium.Web.Proxy/Http2/Http2Helper.cs

@@ -39,15 +39,14 @@ internal static async Task SendHttp2(Stream clientStream, Stream serverStream, i
             CancellationTokenSource cancellationTokenSource, Guid connectionId,
             ExceptionHandler exceptionFunc)
         {
-            var decoder = new Decoder(8192, 4096 * 16);
             var sessions = new ConcurrentDictionary<int, SessionEventArgs>();
 
             // Now async relay all server=>client & client=>server data
             var sendRelay =
-                copyHttp2FrameAsync(clientStream, serverStream, onDataSend, sessionFactory, decoder, sessions, onBeforeRequest, 
+                copyHttp2FrameAsync(clientStream, serverStream, onDataSend, sessionFactory, sessions, onBeforeRequest, 
                     bufferSize, connectionId, true, cancellationTokenSource.Token, exceptionFunc);
             var receiveRelay =
-                copyHttp2FrameAsync(serverStream, clientStream, onDataReceive, sessionFactory, decoder, sessions, onBeforeResponse, 
+                copyHttp2FrameAsync(serverStream, clientStream, onDataReceive, sessionFactory, sessions, onBeforeResponse, 
                     bufferSize, connectionId, false, cancellationTokenSource.Token, exceptionFunc);
 
             await Task.WhenAny(sendRelay, receiveRelay);
@@ -57,11 +56,13 @@ internal static async Task SendHttp2(Stream clientStream, Stream serverStream, i
         }
 
         private static async Task copyHttp2FrameAsync(Stream input, Stream output, Action<byte[], int, int> onCopy,
-            Func<SessionEventArgs> sessionFactory, Decoder decoder, ConcurrentDictionary<int, SessionEventArgs>  sessions, 
+            Func<SessionEventArgs> sessionFactory, ConcurrentDictionary<int, SessionEventArgs>  sessions, 
             Func<SessionEventArgs, Task> onBeforeRequestResponse,
             int bufferSize, Guid connectionId, bool isClient, CancellationToken cancellationToken,
             ExceptionHandler exceptionFunc)
         {
+            var decoder = new Decoder(8192, 4096 * 16);
+
             var headerBuffer = new byte[9];
             var buffer = new byte[32768];
             while (true)
@@ -88,14 +89,61 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
 
                 bool endStream = false;
 
-                //System.Diagnostics.Debug.WriteLine("CLIENT: " + isClient + ", STREAM: " + streamId + ", TYPE: " + type);
+                SessionEventArgs args = null;
+                RequestResponseBase rr = null;
+                if (type == 0 || type == 1)
+                {
+                    if (!sessions.TryGetValue(streamId, out args))
+                    {
+                        if (type == 0)
+                        {
+                            throw new ProxyHttpException("HTTP Body data received before any header frame.", null, args);
+                        }
+
+                        if (!isClient)
+                        {
+                            throw new ProxyHttpException("HTTP Response received before any Request header frame.", null, args);
+                        }
+
+                        args = sessionFactory();
+                        sessions.TryAdd(streamId, args);
+                    }
+
+                    rr = isClient ? (RequestResponseBase)args.HttpClient.Request : args.HttpClient.Response;
+                }
+
+                //System.Diagnostics.Debug.WriteLine("CONN: " + connectionId + ", CLIENT: " + isClient + ", STREAM: " + streamId + ", TYPE: " + type);
                 if (type == 0 /* data */)
                 {
                     bool endStreamFlag = (flags & (int)Http2FrameFlag.EndStream) != 0;
                     if (endStreamFlag)
                     {
                         endStream = true;
                     }
+
+                    if (rr.ReadHttp2BodyTaskCompletionSource != null)
+                    {
+                        // Get body method was called in the "before" event handler
+
+                        var data = rr.Http2BodyData;
+                        data.Write(buffer, 0, length);
+
+                        if (endStream)
+                        {
+                            rr.Body = data.ToArray();
+                            rr.IsBodyRead = true;
+
+                            var tcs = rr.ReadHttp2BodyTaskCompletionSource;
+                            rr.ReadHttp2BodyTaskCompletionSource = null;
+
+                            if (!tcs.Task.IsCompleted)
+                            {
+                                tcs.SetResult(true);
+                            }
+
+                            rr.Http2BodyData = null;
+                        }
+                    }   
                 }
                 else if (type == 1 /*headers*/)
                 {
@@ -125,13 +173,6 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
                         dataLength -= buffer[0];
                     }
 
-                    if (!sessions.TryGetValue(streamId, out var args))
-                    {
-                        // todo: remove sessions when finished, otherwise it will be a "memory leak"
-                        args = sessionFactory();
-                        sessions.TryAdd(streamId, args);
-                    }
-
                     var headerListener = new MyHeaderListener(
                         (name, value) =>
                         {
@@ -149,16 +190,18 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
 
                         if (isClient)
                         {
-                            args.HttpClient.Request.HttpVersion = HttpVersion.Version20;
-                            args.HttpClient.Request.Method = headerListener.Method;
-                            args.HttpClient.Request.OriginalUrl = headerListener.Status;
-                            args.HttpClient.Request.RequestUri = headerListener.GetUri();
+                            var request = args.HttpClient.Request;
+                            request.HttpVersion = HttpVersion.Version20;
+                            request.Method = headerListener.Method;
+                            request.OriginalUrl = headerListener.Status;
+                            request.RequestUri = headerListener.GetUri();
                         }
                         else
                         {
-                            args.HttpClient.Response.HttpVersion = HttpVersion.Version20;
+                            var response = args.HttpClient.Response;
+                            response.HttpVersion = HttpVersion.Version20;
                             int.TryParse(headerListener.Status, out int statusCode);
-                            args.HttpClient.Response.StatusCode = statusCode;
+                            response.StatusCode = statusCode;
                         }
                     }
                     catch (Exception ex)
@@ -168,13 +211,25 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
 
                     if (endHeaders)
                     {
-                        await onBeforeRequestResponse(args);
+                        var tcs = new TaskCompletionSource<bool>();
+                        rr.ReadHttp2BeforeHandlerTaskCompletionSource = tcs;
+
+                        var handler = onBeforeRequestResponse(args);
+
+                        if (handler == await Task.WhenAny(tcs.Task, handler))
+                        {
+                            rr.ReadHttp2BeforeHandlerTaskCompletionSource = null;
+                            tcs.SetResult(true);
+                        }
+
+                        rr.Locked = true;
                     }
                 }
 
                 if (!isClient && endStream)
                 {
                     sessions.TryRemove(streamId, out _);
+                    //System.Diagnostics.Debug.WriteLine("REMOVED CONN: " + connectionId + ", CLIENT: " + isClient + ", STREAM: " + streamId + ", TYPE: " + type);
                 }
 
                 // do not cancel the write operation
@@ -186,7 +241,7 @@ private static async Task copyHttp2FrameAsync(Stream input, Stream output, Actio
                     return;
                 }
 
-                /*using (var fs = new System.IO.FileStream($@"c:\11\{connectionId}.{streamId}.dat", FileMode.Append))
+                /*using (var fs = new System.IO.FileStream($@"c:\temp\{connectionId}.{streamId}.dat", FileMode.Append))
                 {
                     fs.Write(headerBuffer, 0, headerBuffer.Length);
                     fs.Write(buffer, 0, length);

src/Titanium.Web.Proxy/Network/Tcp/TcpClientConnection.cs

@@ -7,6 +7,8 @@
 using System.Net.Sockets;
 using System.Threading.Tasks;
 using Titanium.Web.Proxy.Extensions;
+using Titanium.Web.Proxy.Helpers;
+using Titanium.Web.Proxy.Models;
 
 namespace Titanium.Web.Proxy.Network.Tcp
 {
@@ -34,11 +36,42 @@ internal TcpClientConnection(ProxyServer proxyServer, TcpClient tcpClient)
 
         private readonly TcpClient tcpClient;
 
+        private int? processId;
+
         public Stream GetStream()
         {
             return tcpClient.GetStream();
         }
 
+        public int GetProcessId(ProxyEndPoint endPoint)
+        {
+            if (processId.HasValue)
+            {
+                return processId.Value;
+            }
+
+            if (RunTime.IsWindows)
+            {
+                var remoteEndPoint = (IPEndPoint)RemoteEndPoint;
+
+                // If client is localhost get the process id
+                if (NetworkHelper.IsLocalIpAddress(remoteEndPoint.Address))
+                {
+                    var ipVersion = endPoint.IpV6Enabled ? IpVersion.Ipv6 : IpVersion.Ipv4;
+                    processId = TcpHelper.GetProcessIdByLocalPort(ipVersion, remoteEndPoint.Port);
+                }
+                else
+                {
+                    // can't access process Id of remote request from remote machine
+                    processId = -1;
+                }
+
+                return processId.Value;
+            }
+
+            throw new PlatformNotSupportedException();
+        }
+
         /// <summary>
         ///     Dispose.
         /// </summary>

src/Titanium.Web.Proxy/ProxyServer.cs

@@ -149,7 +149,6 @@ public ProxyServer(string rootCertificateName, string rootCertificateIssuerName,
         ///     Enable disable HTTP/2 support.
         ///     Warning: HTTP/2 support is very limited
         ///      - only enabled when both client and server supports it (no protocol changing in proxy)
-        ///      - GetRequest/ResponseBody(AsString) methods are not supported
         ///      - cannot modify the request/response (e.g header modifications in BeforeRequest/Response events are ignored)
         /// </summary>
         public bool EnableHttp2 { get; set; } = false;