Improve relaying logic

src/HttpToSocks5Proxy/Helpers.cs

@@ -74,5 +74,21 @@ public static AddressType GetAddressType(string hostname)
             }
             return AddressType.DomainName;
         }
+        public static void TryDispose(this Socket socket)
+        {
+            if (socket == null)
+                return;
+
+            try
+            {
+                socket.Shutdown(SocketShutdown.Both);
+            }
+            catch { }
+            try
+            {
+                socket.Close();
+            }
+            catch { }
+        }
     }
 }

src/HttpToSocks5Proxy/HttpToSocks5Proxy.cs

@@ -3,7 +3,6 @@
 using System.Linq;
 using System.Net;
 using System.Net.Sockets;
-using System.Threading.Tasks;
 using System.Collections.Generic;
 using MihaZupan.Enums;
 
@@ -78,8 +77,8 @@ public HttpToSocks5Proxy(ProxyInfo[] proxyList, int internalServerPort = 0)
         {
             if (internalServerPort < 0 || internalServerPort > 65535) throw new ArgumentOutOfRangeException(nameof(internalServerPort));
             if (proxyList == null) throw new ArgumentNullException(nameof(proxyList));
-            if (proxyList.Length == 0) throw new ArgumentException("proxyList is empty");
-            if (proxyList.Any(p => p == null)) throw new ArgumentNullException("Proxy in proxyList is null");
+            if (proxyList.Length == 0) throw new ArgumentException("proxyList is empty", nameof(proxyList));
+            if (proxyList.Any(p => p == null)) throw new ArgumentNullException(nameof(proxyList), "Proxy in proxyList is null");
 
             ProxyList = proxyList;
             InternalServerPort = internalServerPort;
@@ -99,9 +98,22 @@ public HttpToSocks5Proxy(ProxyInfo[] proxyList, int internalServerPort = 0)
         private void OnAcceptCallback(IAsyncResult AR)
         {
             if (Stopped) return;
-            Socket clientSocket = InternalServerSocket.EndAccept(AR);
-            InternalServerSocket.BeginAccept(new AsyncCallback(OnAcceptCallback), null);
-            HandleRequest(clientSocket);
+
+            Socket clientSocket = null;
+            try
+            {
+                clientSocket = InternalServerSocket.EndAccept(AR);
+            }
+            catch { }
+
+            try
+            {
+                InternalServerSocket.BeginAccept(new AsyncCallback(OnAcceptCallback), null);
+            }
+            catch { StopInternalServer(); }
+
+            if (clientSocket != null)
+                HandleRequest(clientSocket);
         }
         private void HandleRequest(Socket clientSocket)
         {
@@ -181,13 +193,12 @@ private void HandleRequest(Socket clientSocket)
             {
                 if (success)
                 {
-                    RelayData(socks5Socket, clientSocket);
-                    RelayData(clientSocket, socks5Socket);
+                    SocketRelay.RelayBiDirectionally(socks5Socket, clientSocket);
                 }
                 else
                 {
-                    TryDisposeSocket(clientSocket);
-                    TryDisposeSocket(socks5Socket);
+                    clientSocket.TryDispose();
+                    socks5Socket.TryDispose();
                 }
             }
         }
@@ -365,46 +376,9 @@ private static void SendString(Socket socket, string text)
             => socket.Send(Encoding.UTF8.GetBytes(text));
         private static void SendError(Socket socket, SocketConnectionResult error, string httpVersion = "HTTP/1.1 ")
             => SendString(socket, ErrorResponseBuilder.Build(error, httpVersion));
-        private static void RelayData(Socket source, Socket target)
-        {
-            Task.Run(() =>
-            {
-                try
-                {
-                    // ToDo: Async socket relaying to avoid hijacking two Tasks for the duration of the underlying connection
-                    int read;
-                    byte[] buffer = new byte[81920];
-                    while ((read = source.Receive(buffer, 0, buffer.Length, SocketFlags.None)) > 0)
-                    {
-                        target.Send(buffer, 0, read, SocketFlags.None);
-                    }
-                }
-                finally
-                {
-                    TryDisposeSocket(source);
-                    TryDisposeSocket(target);
-                }
-            });
-        }
 
         private static Socket CreateSocket()
             => new Socket(SocketType.Stream, ProtocolType.Tcp);
-        private static void TryDisposeSocket(Socket socket)
-        {
-            if (socket == null)
-                return;
-
-            try
-            {
-                socket.Shutdown(SocketShutdown.Both);
-            }
-            catch { }
-            try
-            {
-                socket.Close();
-            }
-            catch { }
-        }
 
         private bool Stopped = false;
         public void StopInternalServer()

src/HttpToSocks5Proxy/SocketRelay.cs

@@ -0,0 +1,145 @@
+using System.Net.Sockets;
+using System.Threading.Tasks;
+
+namespace MihaZupan
+{
+    internal class SocketRelay
+    {
+        private SocketAsyncEventArgs RecSAEA, SendSAEA;
+        private Socket Source, Target;
+        private byte[] Buffer;
+        private int SendingOffset;
+        private int Received;
+        private int StackDepth;
+        private const int StackDepthCutoff = 256;
+
+        public SocketRelay Other;
+
+        public SocketRelay(Socket source, Socket target)
+        {
+            Source = source;
+            Target = target;
+            Buffer = new byte[81920];
+            RecSAEA = new SocketAsyncEventArgs()
+            {
+                UserToken = this
+            };
+            SendSAEA = new SocketAsyncEventArgs()
+            {
+                UserToken = this
+            };
+            RecSAEA.SetBuffer(Buffer, 0, Buffer.Length);
+            SendSAEA.SetBuffer(Buffer, 0, Buffer.Length);
+            RecSAEA.Completed += OnSaeaReceiveCompleted;
+            SendSAEA.Completed += OnSaeaSendCompleted;
+        }
+
+        private void OnCleanup()
+        {
+            if (Other is null)
+                return;
+
+            Source.TryDispose();
+            Target.TryDispose();
+            try { RecSAEA?.Dispose(); } catch { }
+            try { SendSAEA?.Dispose(); } catch { }
+
+            Source = Target = null;
+            RecSAEA = SendSAEA = null;
+            Buffer = null;
+
+            Other?.OnCleanup();
+            Other = null;
+        }
+
+        public void StartRelaying()
+        {
+            try
+            {
+                if (!Source.ReceiveAsync(RecSAEA))
+                {
+                    OnReceived();
+                }
+            }
+            catch
+            {
+                OnCleanup();
+            }
+        }
+        private void OnReceived()
+        {
+            try
+            {
+                SendingOffset = 0;
+                Received = RecSAEA.BytesTransferred;
+
+                SendSAEA.SetBuffer(Buffer, 0, Received);
+
+                if (!Target.SendAsync(SendSAEA))
+                {
+                    OnSent();
+                }
+            }
+            catch
+            {
+                OnCleanup();
+            }
+        }
+        private void OnSent()
+        {
+            try
+            {
+                SendingOffset += SendSAEA.BytesTransferred;
+
+                while (SendingOffset != Received)
+                {
+                    SendSAEA.SetBuffer(Buffer, SendingOffset, Received - SendingOffset);
+
+                    if (Target.SendAsync(SendSAEA))
+                        return;
+
+                    SendingOffset += SendSAEA.BytesTransferred;
+                }
+
+                if (++StackDepth == StackDepthCutoff)
+                {
+                    StackDepth = 0;
+                    Task.Run(StartRelaying);
+                }
+                else
+                {
+                    StartRelaying();
+                }
+            }
+            catch
+            {
+                OnCleanup();
+            }
+        }
+
+        private static void OnSaeaReceiveCompleted(object _, SocketAsyncEventArgs saea)
+        {
+            var relay = saea.UserToken as SocketRelay;
+            relay.StackDepth = 0;
+            relay.OnReceived();
+        }
+        private static void OnSaeaSendCompleted(object _, SocketAsyncEventArgs saea)
+        {
+            var relay = saea.UserToken as SocketRelay;
+            relay.StackDepth = 0;
+            relay.OnSent();
+        }
+
+        public static void RelayBiDirectionally(Socket s1, Socket s2)
+        {
+            var relayOne = new SocketRelay(s1, s2);
+            var relayTwo = new SocketRelay(s2, s1);
+
+            relayOne.Other = relayTwo;
+            relayTwo.Other = relayOne;
+
+            Task.Run(relayOne.StartRelaying);
+            Task.Run(relayTwo.StartRelaying);
+        }
+    }
+}