test: Add SOCKS5 UDP test

.github/workflows/CI.yml

@@ -38,7 +38,7 @@ jobs:
 
       - name: Test
         shell: pwsh
-        run: dotnet test -c Release --filter ClassName!~Socks5
+        run: dotnet test -c Release
 
   build:
     name: Build

Shadowsocks.Protocol/ListenServices/UdpListenService.cs

@@ -9,6 +9,7 @@
 
 namespace Shadowsocks.Protocol.ListenServices
 {
+	//TODO Remake
 	public class UdpListenService : IListenService
 	{
 		private readonly ILogger<UdpListenService> _logger;

Shadowsocks.Protocol/UdpClients/ShadowsocksUdpClient.cs

@@ -9,6 +9,7 @@
 
 namespace Shadowsocks.Protocol.UdpClients
 {
+	//TODO Remake
 	public class ShadowsocksUdpClient : IUdpClient
 	{
 		private readonly ILogger _logger;

Socks5/Clients/Socks5Client.cs

@@ -126,29 +126,37 @@ public async ValueTask<ServerBound> UdpAssociateAsync(IPAddress address, ushort
 			return bound;
 		}
 
-		//TODO .NET6.0
-		public async Task<Socks5UdpReceivePacket> ReceiveAsync()
+		public async Task<Socks5UdpReceivePacket> ReceiveAsync(CancellationToken token = default)
 		{
 			Verify.Operation(Status is Status.Established && _udpClient is not null, @"Socks5 is not established.");
 
-			var res = await _udpClient.ReceiveAsync();
+			var buffer = ArrayPool<byte>.Shared.Rent(0x10000);
+			try
+			{
+				var length = await _udpClient.Client.ReceiveAsync(buffer, SocketFlags.None, token);
 
-			return Unpack.Udp(res.Buffer);
+				return Unpack.Udp(buffer.AsMemory(0, length));
+			}
+			finally
+			{
+				ArrayPool<byte>.Shared.Return(buffer);
+			}
 		}
 
-		public Task<int> SendUdpAsync(ReadOnlyMemory<byte> data, string dst, ushort dstPort)
+		public Task<int> SendUdpAsync(ReadOnlyMemory<byte> data, string dst, ushort dstPort, CancellationToken token = default)
 		{
-			return SendUdpAsync(data, dst, default, dstPort);
+			return SendUdpAsync(data, dst, default, dstPort, token);
 		}
 
-		public Task<int> SendUdpAsync(ReadOnlyMemory<byte> data, IPAddress dstAddress, ushort dstPort)
+		public Task<int> SendUdpAsync(ReadOnlyMemory<byte> data, IPAddress dstAddress, ushort dstPort, CancellationToken token = default)
 		{
-			return SendUdpAsync(data, default, dstAddress, dstPort);
+			return SendUdpAsync(data, default, dstAddress, dstPort, token);
 		}
 
 		private async Task<int> SendUdpAsync(
 			ReadOnlyMemory<byte> data,
-			string? dst, IPAddress? dstAddress, ushort dstPort)
+			string? dst, IPAddress? dstAddress, ushort dstPort,
+			CancellationToken token = default)
 		{
 			Verify.Operation(Status is Status.Established && _udpClient is not null, @"Socks5 is not established.");
 
@@ -157,7 +165,7 @@ private async Task<int> SendUdpAsync(
 			{
 				var length = Pack.Udp(buffer, dst, dstAddress, dstPort, data.Span);
 
-				return await _udpClient.SendAsync(buffer, length);
+				return await _udpClient.Client.SendAsync(buffer.AsMemory(0, length), SocketFlags.None, token);
 			}
 			finally
 			{

Socks5/Models/Socks5UdpReceivePacket.cs

@@ -4,13 +4,18 @@
 
 namespace Socks5.Models
 {
+	// +----+------+------+----------+----------+----------+
+	// |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
+	// +----+------+------+----------+----------+----------+
+	// | 2  |  1   |  1   | Variable |    2     | Variable |
+	// +----+------+------+----------+----------+----------+
 	public struct Socks5UdpReceivePacket
 	{
-		public Memory<byte> Data;
 		public byte Fragment;
 		public AddressType Type;
 		public IPAddress? Address;
 		public string? Domain;
 		public ushort Port;
+		public Memory<byte> Data;
 	}
 }

Socks5/Servers/Socks5Server.cs → Socks5/Servers/SimpleSocks5Server.cs

@@ -11,7 +11,10 @@
 
 namespace Socks5.Servers
 {
-	public class Socks5Server
+	/// <summary>
+	/// A simple SOCKS5 server for test use only
+	/// </summary>
+	public class SimpleSocks5Server
 	{
 		public TcpListener TcpListener { get; }
 
@@ -23,10 +26,18 @@ public class Socks5Server
 			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;
 
-		public Socks5Server(IPEndPoint bindEndPoint, UsernamePassword? credential = null)
+		public SimpleSocks5Server(IPEndPoint bindEndPoint, UsernamePassword? credential = null)
 		{
 			_credential = credential;
 			TcpListener = new TcpListener(bindEndPoint);
@@ -93,8 +104,7 @@ private async ValueTask HandleAsync(TcpClient rec, CancellationToken token)
 						}
 						case Command.UdpAssociate:
 						{
-							//TODO
-							await service.SendReplyAsync(Socks5Reply.CommandNotSupported, ReplyTcpBound, token);
+							await service.SendReplyAsync(Socks5Reply.Succeeded, ReplyUdpBound, token);
 							break;
 						}
 						default:

Socks5/Servers/SimpleSocks5UdpServer.cs

@@ -0,0 +1,107 @@
+using Microsoft;
+using Microsoft.VisualStudio.Threading;
+using Socks5.Enums;
+using Socks5.Utils;
+using System;
+using System.Buffers;
+using System.Net;
+using System.Net.Sockets;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Socks5.Servers
+{
+	/// <summary>
+	/// Just for test use only
+	/// </summary>
+	public class SimpleSocks5UdpServer
+	{
+		public UdpClient UdpListener { get; }
+
+		private readonly CancellationTokenSource _cts;
+
+		public SimpleSocks5UdpServer(IPEndPoint bindEndPoint)
+		{
+			UdpListener = new UdpClient(bindEndPoint);
+
+			_cts = new CancellationTokenSource();
+		}
+
+		public async ValueTask StartAsync()
+		{
+			try
+			{
+				while (!_cts.IsCancellationRequested)
+				{
+					//TODO .NET6.0
+					var message = await UdpListener.ReceiveAsync();
+
+					HandleAsync(message, _cts.Token).Forget();
+				}
+			}
+			catch (Exception)
+			{
+				Stop();
+			}
+		}
+
+		private async ValueTask HandleAsync(UdpReceiveResult result, CancellationToken token)
+		{
+			if (token.IsCancellationRequested)
+			{
+				return;
+			}
+
+			var socks5UdpPacket = Unpack.Udp(result.Buffer);
+			if (socks5UdpPacket.Fragment is not 0x00)
+			{
+				return; // Ignore
+			}
+
+			UdpClient client;
+			if (socks5UdpPacket.Type is AddressType.Domain)
+			{
+				Assumes.NotNull(socks5UdpPacket.Domain);
+				client = new UdpClient(socks5UdpPacket.Domain, socks5UdpPacket.Port);
+			}
+			else
+			{
+				Assumes.NotNull(socks5UdpPacket.Address);
+				client = new UdpClient(socks5UdpPacket.Address.AddressFamily);
+				client.Connect(socks5UdpPacket.Address, socks5UdpPacket.Port);
+			}
+
+			await client.Client.SendAsync(socks5UdpPacket.Data, SocketFlags.None, token);
+
+			var headerLength = result.Buffer.Length - socks5UdpPacket.Data.Length;
+			var receiveBuffer = ArrayPool<byte>.Shared.Rent(0x10000);
+			try
+			{
+				var receiveLength = await client.Client.ReceiveAsync(receiveBuffer.AsMemory(headerLength), SocketFlags.None, token);
+				result.Buffer.AsSpan(0, headerLength).CopyTo(receiveBuffer);
+
+				//TODO .NET6.0
+				await UdpListener.Client.SendToAsync(
+					new ArraySegment<byte>(receiveBuffer, 0, headerLength + receiveLength),
+					SocketFlags.None,
+					result.RemoteEndPoint);
+			}
+			finally
+			{
+				ArrayPool<byte>.Shared.Return(receiveBuffer);
+			}
+		}
+
+		public void Stop()
+		{
+			try
+			{
+				UdpListener.Dispose();
+			}
+			finally
+			{
+				_cts.Cancel();
+			}
+		}
+	}
+}

Socks5/Utils/Socks5TestUtils.cs

@@ -124,9 +124,9 @@ public static async ValueTask<bool> Socks5UdpAssociateAsync(
 				buffer[offset++] = 0x00;
 				buffer[offset++] = 0x01;
 
-				await client.SendUdpAsync(buffer.AsMemory(0, offset), targetHost, targetPort);
+				await client.SendUdpAsync(buffer.AsMemory(0, offset), targetHost, targetPort, token);
 
-				var res = await client.ReceiveAsync();
+				var res = await client.ReceiveAsync(token);
 
 				return res.Data.Span[..2].SequenceEqual(buffer.AsSpan(0, 2));
 			}

Socks5/Utils/Unpack.cs

@@ -7,6 +7,7 @@
 using System.Buffers.Binary;
 using System.Collections.Generic;
 using System.Net;
+using System.Runtime.InteropServices;
 using System.Text;
 
 namespace Socks5.Utils
@@ -117,32 +118,34 @@ public static int DestinationAddress(AddressType type, ReadOnlySpan<byte> bytes,
 			return offset;
 		}
 
-		public static Socks5UdpReceivePacket Udp(byte[] buffer)
+		public static Socks5UdpReceivePacket Udp(ReadOnlyMemory<byte> buffer)
 		{
 			// +----+------+------+----------+----------+----------+
 			// |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
 			// +----+------+------+----------+----------+----------+
 			// | 2  |  1   |  1   | Variable |    2     | Variable |
 			// +----+------+------+----------+----------+----------+
 
-			Requires.Range(buffer.LongLength >= 7, nameof(buffer));
+			var span = buffer.Span;
+			Requires.Range(buffer.Length >= 7, nameof(buffer));
 
 			var res = new Socks5UdpReceivePacket();
 
-			if (buffer[0] != Constants.Rsv || buffer[1] != Constants.Rsv)
+			if (span[0] is not Constants.Rsv || span[1] is not Constants.Rsv)
 			{
-				throw new ProtocolErrorException($@"Protocol failed, RESERVED is not 0x0000: 0x{buffer[0]:X2}{buffer[1]:X2}.");
+				throw new ProtocolErrorException($@"Protocol failed, RESERVED is not 0x0000: 0x{span[0]:X2}{span[1]:X2}.");
 			}
 
-			res.Fragment = buffer[2];
+			res.Fragment = span[2];
 
-			res.Type = (AddressType)buffer[3];
+			res.Type = (AddressType)span[3];
+			Requires.Defined(res.Type, nameof(res.Type));
 
 			var offset = 4;
-			offset += DestinationAddress(res.Type, buffer.AsSpan(offset), out res.Address, out res.Domain);
+			offset += DestinationAddress(res.Type, span[offset..], out res.Address, out res.Domain);
 
-			res.Port = BinaryPrimitives.ReadUInt16BigEndian(buffer.AsSpan(offset));
-			res.Data = buffer.AsMemory(offset + 2);
+			res.Port = BinaryPrimitives.ReadUInt16BigEndian(span[offset..]);
+			res.Data = MemoryMarshal.AsMemory(buffer[(offset + 2)..]);
 
 			return res;
 		}

UnitTest/Socks5Test.cs

@@ -1,5 +1,6 @@
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using Microsoft.VisualStudio.Threading;
+using Socks5.Enums;
 using Socks5.Models;
 using Socks5.Servers;
 using Socks5.Utils;
@@ -20,7 +21,7 @@ public async Task ConnectTestAsync()
 				UserName = @"114514！",
 				Password = @"1919810￥"
 			};
-			var server = new Socks5Server(serverEndpoint, userPass);
+			var server = new SimpleSocks5Server(serverEndpoint, userPass);
 			server.StartAsync().Forget();
 			try
 			{
@@ -42,12 +43,39 @@ public async Task ConnectTestAsync()
 		[TestMethod]
 		public async Task UdpAssociateTestAsync()
 		{
-			var option = new Socks5CreateOption
+			var serverEndpoint = new IPEndPoint(IPAddress.Loopback, 0);
+			var userPass = new UsernamePassword
+			{
+				UserName = @"114514！",
+				Password = @"1919810￥"
+			};
+			var server = new SimpleSocks5Server(serverEndpoint, userPass);
+			var udpServer = new SimpleSocks5UdpServer(serverEndpoint);
+			udpServer.StartAsync().Forget();
+			server.ReplyUdpBound = new ServerBound
 			{
-				Address = IPAddress.Loopback,
-				Port = 23333
+				Type = AddressType.IPv4,
+				Address = ((IPEndPoint)udpServer.UdpListener.Client.LocalEndPoint!).Address,
+				Domain = default,
+				Port = (ushort)((IPEndPoint)udpServer.UdpListener.Client.LocalEndPoint!).Port,
 			};
-			Assert.IsTrue(await Socks5TestUtils.Socks5UdpAssociateAsync(option));
+			server.StartAsync().Forget();
+			try
+			{
+				var port = (ushort)((IPEndPoint)server.TcpListener.LocalEndpoint).Port;
+				var option = new Socks5CreateOption
+				{
+					Address = IPAddress.Loopback,
+					Port = port,
+					UsernamePassword = userPass
+				};
+				Assert.IsTrue(await Socks5TestUtils.Socks5UdpAssociateAsync(option));
+			}
+			finally
+			{
+				server.Stop();
+				udpServer.Stop();
+			}
 		}
 	}
 }