Mixed implementation for startWith + PeekBytesAsync with the use of bufferpool (to get the best of the two worlds)

Author: StephaneGraziano

src/StreamExtended/Network/CopyStream.cs

@@ -55,9 +55,9 @@ public byte PeekByteFromBuffer(int index)
             return reader.PeekByteAsync(index, cancellationToken);
         }
 
-        public Task<byte[]> PeekBytesAsync(int index, int size, CancellationToken cancellationToken = default(CancellationToken))
+        public Task<int> PeekBytesAsync(byte[] buffer, int offset, int index, int size, CancellationToken cancellationToken = default(CancellationToken))
         {
-            return reader.PeekBytesAsync(index, size, cancellationToken);
+            return reader.PeekBytesAsync(buffer, offset, index, size, cancellationToken);
         }
 
         public void Flush()

src/StreamExtended/Network/CustomBufferedPeekStream.cs

@@ -88,12 +88,14 @@ byte ICustomStreamReader.PeekByteFromBuffer(int index)
         /// <summary>
         /// Peeks bytes asynchronous.
         /// </summary>
+        /// <param name="buffer">The buffer to copy.</param>
+        /// <param name="offset">The offset where copying.</param>
         /// <param name="index">The index.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        Task<byte[]> ICustomStreamReader.PeekBytesAsync(int index, int size, CancellationToken cancellationToken)
+        Task<int> ICustomStreamReader.PeekBytesAsync(byte[] buffer, int offset, int index, int size, CancellationToken cancellationToken)
         {
-            return baseStream.PeekBytesAsync(index, size, cancellationToken);
+            return baseStream.PeekBytesAsync(buffer, offset, index, size, cancellationToken);
         }
 
         /// <summary>

src/StreamExtended/Network/CustomBufferedStream.cs

@@ -252,10 +252,12 @@ public override int ReadByte()
         /// <summary>
         /// Peeks bytes asynchronous.
         /// </summary>
+        /// <param name="buffer">The buffer to copy.</param>
+        /// <param name="offset">The offset where copying.</param>
         /// <param name="index">The index.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public async Task<byte[]> PeekBytesAsync(int index, int size, CancellationToken cancellationToken = default(CancellationToken))
+        public async Task<int> PeekBytesAsync(byte[] buffer, int offset, int index, int size, CancellationToken cancellationToken = default(CancellationToken))
         {
             if (Available <= index)
             {
@@ -270,12 +272,12 @@ public override int ReadByte()
 
             if (Available <= (index + size))
             {
-                return null;
+                return -1;
             }
 
-            var vRet = new byte[size];
-            Array.Copy(streamBuffer, index, vRet, 0, size);
-            return vRet;
+            Buffer.BlockCopy(streamBuffer, index, buffer, offset, size);
+
+            return size;
         }
 
         /// <summary>

src/StreamExtended/Network/ICustomStreamReader.cs

@@ -40,10 +40,12 @@ public interface ICustomStreamReader
         /// <summary>
         /// Peeks bytes asynchronous.
         /// </summary>
+        /// <param name="buffer">The buffer to copy.</param>
+        /// <param name="offset">The offset where copying.</param>
         /// <param name="index">The index.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        Task<byte[]> PeekBytesAsync(int index, int size, CancellationToken cancellationToken = default(CancellationToken));
+        Task<int> PeekBytesAsync(byte[] buffer, int offset, int index, int size, CancellationToken cancellationToken = default(CancellationToken));
 
         byte ReadByteFromBuffer();
 

src/Titanium.Web.Proxy/ExplicitClientHandler.cs

@@ -49,7 +49,7 @@ private async Task handleClient(ExplicitProxyEndPoint endPoint, TcpClientConnect
                 TunnelConnectSessionEventArgs connectArgs = null;
 
                 // Client wants to create a secure tcp tunnel (probably its a HTTPS or Websocket request)
-                if (await HttpHelper.IsConnectMethod(clientStream, cancellationToken) == 1)
+                if (await HttpHelper.IsConnectMethod(clientStream, BufferPool, BufferSize, cancellationToken) == 1)
                 {
                     // read the first line HTTP command
                     string httpCmd = await clientStream.ReadLineAsync(cancellationToken);
@@ -220,7 +220,7 @@ await clientStreamWriter.WriteResponseAsync(connectArgs.HttpClient.Response,
                                 $"Couldn't authenticate host '{connectHostname}' with certificate '{certName}'.", e, connectArgs);
                         }
 
-                        if (await HttpHelper.IsConnectMethod(clientStream) == -1)
+                        if (await HttpHelper.IsConnectMethod(clientStream, BufferPool, BufferSize, cancellationToken) == -1)
                         {
                             decryptSsl = false;
                         }
@@ -292,7 +292,7 @@ await TcpHelper.SendRaw(clientStream, connection.Stream, BufferPool, BufferSize,
                     }
                 }
 
-                if (connectArgs != null && await HttpHelper.IsPriMethod(clientStream) == 1)
+                if (connectArgs != null && await HttpHelper.IsPriMethod(clientStream, BufferPool, BufferSize, cancellationToken) == 1)
                 {
                     // todo
                     string httpCmd = await clientStream.ReadLineAsync(cancellationToken);

src/Titanium.Web.Proxy/Helpers/HttpHelper.cs

@@ -3,6 +3,7 @@
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
+using StreamExtended;
 using StreamExtended.Network;
 using Titanium.Web.Proxy.Extensions;
 using Titanium.Web.Proxy.Http;
@@ -122,19 +123,19 @@ internal static string GetWildCardDomainName(string hostname)
         /// </summary>
         /// <param name="clientStreamReader">The client stream reader.</param>
         /// <returns>1: when CONNECT, 0: when valid HTTP method, -1: otherwise</returns>
-        internal static Task<int> IsConnectMethod(ICustomStreamReader clientStreamReader, CancellationToken cancellationToken = default(CancellationToken))
+        internal static Task<int> IsConnectMethod(ICustomStreamReader clientStreamReader, IBufferPool bufferPool, int bufferSize, CancellationToken cancellationToken = default(CancellationToken))
         {
-            return startsWith(clientStreamReader, "CONNECT", cancellationToken);
+            return startsWith(clientStreamReader, bufferPool, bufferSize, "CONNECT", cancellationToken);
         }
 
         /// <summary>
         ///     Determines whether is pri method (HTTP/2).
         /// </summary>
         /// <param name="clientStreamReader">The client stream reader.</param>
         /// <returns>1: when PRI, 0: when valid HTTP method, -1: otherwise</returns>
-        internal static Task<int> IsPriMethod(ICustomStreamReader clientStreamReader, CancellationToken cancellationToken = default(CancellationToken))
+        internal static Task<int> IsPriMethod(ICustomStreamReader clientStreamReader, IBufferPool bufferPool, int bufferSize, CancellationToken cancellationToken = default(CancellationToken))
         {
-            return startsWith(clientStreamReader, "PRI", cancellationToken);
+            return startsWith(clientStreamReader, bufferPool, bufferSize, "PRI", cancellationToken);
         }
 
         /// <summary>
@@ -145,22 +146,46 @@ internal static string GetWildCardDomainName(string hostname)
         /// <returns>
         ///     1: when starts with the given string, 0: when valid HTTP method, -1: otherwise
         /// </returns>
-        private static async Task<int> startsWith(ICustomStreamReader clientStreamReader, string expectedStart, CancellationToken cancellationToken = default(CancellationToken))
+        private static async Task<int> startsWith(ICustomStreamReader clientStreamReader, IBufferPool bufferPool, int bufferSize, string expectedStart, CancellationToken cancellationToken = default(CancellationToken))
         {
             int iRet = -1;
-            int lengthToCheck = 10;
-
-            var vBuffer = await clientStreamReader.PeekBytesAsync(0, lengthToCheck, cancellationToken);
-            if (vBuffer != null)
+            const int lengthToCheck = 10;
+            byte[] buffer = null;
+            try
             {
-                var httpMethod = defaultEncoding.GetString(vBuffer);
+                buffer = bufferPool.GetBuffer(Math.Max(bufferSize, lengthToCheck));
 
-                if (httpMethod.StartsWith(expectedStart))
-                    iRet = 1;
-                else if (Regex.Match(httpMethod, @"^[a-z]{3,} ", RegexOptions.IgnoreCase).Success) //valid HTTP requests start by at least 3 letters + space
-                    iRet = 0;
-            }
+                int peeked = await clientStreamReader.PeekBytesAsync(buffer, 0, 0, lengthToCheck, cancellationToken);
+
+                if (peeked > 0)
+                {
+                    bool isExpected = true;
+
+                    for (int i = 0; i < lengthToCheck; i++)
+                    {
+                        int b = buffer[i];
 
+                        if (b == ' ' && i > 2)
+                            return isExpected ? 1 : 0;
+                        else
+                        {
+                            char ch = (char)b;
+                            if (!char.IsLetter(ch))
+                                return -1;
+                            else if (i >= expectedStart.Length || ch != expectedStart[i])
+                                isExpected = false;                            
+                        }
+                    }
+
+                    // only letters
+                    iRet = isExpected ? 1 : 0;
+                }
+            }
+            finally
+            {
+                bufferPool.ReturnBuffer(buffer);
+                buffer = null;
+            }
             return iRet;
         }
     }