From 0c6d53968116c39c137acb7d53d2ef1f355728e5 Mon Sep 17 00:00:00 2001
From: Bruce Wayne <admin@bige0.com>
Date: Thu, 16 Sep 2021 17:39:03 +0800
Subject: [PATCH] refactor: Socks5 UDP server

---
 Socks5/Servers/SimpleSocks5Server.cs    | 24 +++++++++++++++---------
 Socks5/Servers/SimpleSocks5UdpServer.cs | 24 ++++++++++++++++--------
 UnitTest/Socks5Test.cs                  | 11 -----------
 3 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/Socks5/Servers/SimpleSocks5Server.cs b/Socks5/Servers/SimpleSocks5Server.cs
index ad0d873..80834ef 100644
--- a/Socks5/Servers/SimpleSocks5Server.cs
+++ b/Socks5/Servers/SimpleSocks5Server.cs
@@ -26,14 +26,6 @@ public class SimpleSocks5Server
 			Port = IPEndPoint.MinPort,
 		};
 
-		public ServerBound ReplyUdpBound { get; set; } = new()
-		{
-			Type = AddressType.IPv4,
-			Address = IPAddress.Any,
-			Domain = default,
-			Port = IPEndPoint.MinPort,
-		};
-
 		private readonly UsernamePassword? _credential;
 		private readonly CancellationTokenSource _cts;
 
@@ -102,11 +94,25 @@ private async ValueTask HandleAsync(Socket socket, CancellationToken token)
 					}
 					case Command.UdpAssociate:
 					{
-						await service.SendReplyAsync(Socks5Reply.Succeeded, ReplyUdpBound, token);
+						var remote = (IPEndPoint)socket.RemoteEndPoint!;
+						var local = new IPEndPoint(((IPEndPoint)TcpListener.LocalEndpoint).Address, IPEndPoint.MinPort);
+						using var udpServer = new SimpleSocks5UdpServer(local, remote);
+						udpServer.StartAsync().Forget();
+
+						ServerBound replyUdpBound = new()
+						{
+							Type = AddressType.IPv4,
+							Address = local.Address,
+							Domain = default,
+							Port = (ushort)((IPEndPoint)udpServer.UdpListener.Client.LocalEndPoint!).Port,
+						};
+
+						await service.SendReplyAsync(Socks5Reply.Succeeded, replyUdpBound, token);
 
 						// wait remote close
 						var result = await pipe.Input.ReadAsync(token);
 						Report.IfNot(result.IsCompleted);
+
 						break;
 					}
 					default:
diff --git a/Socks5/Servers/SimpleSocks5UdpServer.cs b/Socks5/Servers/SimpleSocks5UdpServer.cs
index 706e47f..be61f2a 100644
--- a/Socks5/Servers/SimpleSocks5UdpServer.cs
+++ b/Socks5/Servers/SimpleSocks5UdpServer.cs
@@ -1,5 +1,4 @@
 using Microsoft;
-using Microsoft.VisualStudio.Threading;
 using Socks5.Enums;
 using Socks5.Utils;
 using System;
@@ -14,14 +13,21 @@ namespace Socks5.Servers
 	/// <summary>
 	/// Just for test use only
 	/// </summary>
-	public class SimpleSocks5UdpServer
+	public class SimpleSocks5UdpServer : IDisposable
 	{
 		public UdpClient UdpListener { get; }
+		private readonly IPEndPoint _remoteEndpoint;
 
 		private readonly CancellationTokenSource _cts;
 
-		public SimpleSocks5UdpServer(IPEndPoint bindEndPoint)
+		private const int MaxUdpSize = 0x10000;
+
+		public SimpleSocks5UdpServer(IPEndPoint bindEndPoint, IPEndPoint remoteEndpoint)
 		{
+			_remoteEndpoint = remoteEndpoint;
+			Requires.NotNull(bindEndPoint, nameof(bindEndPoint));
+			Requires.NotNull(remoteEndpoint, nameof(remoteEndpoint));
+
 			UdpListener = new UdpClient(bindEndPoint);
 
 			_cts = new CancellationTokenSource();
@@ -35,13 +41,15 @@ public async ValueTask StartAsync()
 				{
 					//TODO .NET6.0
 					var message = await UdpListener.ReceiveAsync();
-
-					HandleAsync(message, _cts.Token).Forget();
+					if (Equals(message.RemoteEndPoint.Address, _remoteEndpoint.Address))
+					{
+						await HandleAsync(message, _cts.Token);
+					}
 				}
 			}
 			catch (Exception)
 			{
-				Stop();
+				Dispose();
 			}
 		}
 
@@ -74,7 +82,7 @@ private async ValueTask HandleAsync(UdpReceiveResult result, CancellationToken t
 			await client.Client.SendAsync(socks5UdpPacket.Data, SocketFlags.None, token);
 
 			var headerLength = result.Buffer.Length - socks5UdpPacket.Data.Length;
-			var receiveBuffer = ArrayPool<byte>.Shared.Rent(0x10000);
+			var receiveBuffer = ArrayPool<byte>.Shared.Rent(MaxUdpSize);
 			try
 			{
 				var receiveLength = await client.Client.ReceiveAsync(receiveBuffer.AsMemory(headerLength), SocketFlags.None, token);
@@ -93,7 +101,7 @@ await UdpListener.Client.SendToAsync(
 			}
 		}
 
-		public void Stop()
+		public void Dispose()
 		{
 			try
 			{
diff --git a/UnitTest/Socks5Test.cs b/UnitTest/Socks5Test.cs
index bde6840..18c87bd 100644
--- a/UnitTest/Socks5Test.cs
+++ b/UnitTest/Socks5Test.cs
@@ -1,6 +1,5 @@
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using Microsoft.VisualStudio.Threading;
-using Socks5.Enums;
 using Socks5.Models;
 using Socks5.Servers;
 using Socks5.Utils;
@@ -50,15 +49,6 @@ public async Task UdpAssociateTestAsync()
 				Password = @"1919810￥"
 			};
 			var server = new SimpleSocks5Server(serverEndpoint, userPass);
-			var udpServer = new SimpleSocks5UdpServer(serverEndpoint);
-			udpServer.StartAsync().Forget();
-			server.ReplyUdpBound = new ServerBound
-			{
-				Type = AddressType.IPv4,
-				Address = ((IPEndPoint)udpServer.UdpListener.Client.LocalEndPoint!).Address,
-				Domain = default,
-				Port = (ushort)((IPEndPoint)udpServer.UdpListener.Client.LocalEndPoint!).Port,
-			};
 			server.StartAsync().Forget();
 			try
 			{
@@ -74,7 +64,6 @@ public async Task UdpAssociateTestAsync()
 			finally
 			{
 				server.Stop();
-				udpServer.Stop();
 			}
 		}
 	}