refactor: Better ShadowsocksUdpClient

HMBSbige · Sep 17, 2021 · 5d4f9c6 · 5d4f9c6
1 parent 0c6d539
commit 5d4f9c6
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 31 deletions.
diff --git a/Shadowsocks.Protocol/ListenServices/TcpListenService.cs b/Shadowsocks.Protocol/ListenServices/TcpListenService.cs
@@ -41,20 +41,20 @@ public async ValueTask StartAsync()
 			try
 			{
 				TCPListener.Start();
-				_logger.LogInformation(@"{LoggerHeader} {$Local} Start", LoggerHeader, TCPListener.LocalEndpoint);
+				_logger.LogInformation(@"{LoggerHeader} {Local} Start", LoggerHeader, TCPListener.LocalEndpoint);
 
 				while (!_cts.IsCancellationRequested)
 				{
 					var socket = await TCPListener.AcceptSocketAsync();
 					socket.NoDelay = true;
 
-					_logger.LogInformation(@"{LoggerHeader} {$Remote} => {$Local}", LoggerHeader, socket.RemoteEndPoint, socket.LocalEndPoint);
+					_logger.LogInformation(@"{LoggerHeader} {Remote} => {Local}", LoggerHeader, socket.RemoteEndPoint, socket.LocalEndPoint);
 					HandleAsync(socket, _cts.Token).Forget();
 				}
 			}
 			catch (Exception ex)
 			{
-				_logger.LogError(ex, @"{LoggerHeader} {$Local} Stop!", LoggerHeader, TCPListener.LocalEndpoint);
+				_logger.LogError(ex, @"{LoggerHeader} {Local} Stop!", LoggerHeader, TCPListener.LocalEndpoint);
 				Stop();
 			}
 		}
@@ -100,7 +100,7 @@ private async Task HandleAsync(Socket socket, CancellationToken token)
 			finally
 			{
 				socket.FullClose();
-				_logger.LogInformation(@"{LoggerHeader} {$Remote} disconnected", LoggerHeader, remoteEndPoint);
+				_logger.LogInformation(@"{LoggerHeader} {Remote} disconnected", LoggerHeader, remoteEndPoint);
 			}
 		}
 

diff --git a/Shadowsocks.Protocol/ListenServices/UdpListenService.cs b/Shadowsocks.Protocol/ListenServices/UdpListenService.cs
@@ -38,15 +38,15 @@ public async ValueTask StartAsync()
 			try
 			{
 				UdpListener.Client.Bind(_local);
-				_logger.LogInformation(@"{LoggerHeader} {$Local} Start", LoggerHeader, UdpListener.Client.LocalEndPoint);
+				_logger.LogInformation(@"{LoggerHeader} {Local} Start", LoggerHeader, UdpListener.Client.LocalEndPoint);
 
 				while (!_cts.IsCancellationRequested)
 				{
 					//TODO .NET6.0
 					var message = await UdpListener.ReceiveAsync();
 #if DEBUG
 					_logger.LogDebug(
-						@"{LoggerHeader}: {ReceiveLength} bytes {Remote} => {$Local}",
+						@"{LoggerHeader}: {ReceiveLength} bytes {Remote} => {Local}",
 						LoggerHeader,
 						message.Buffer.Length,
 						message.RemoteEndPoint,
@@ -58,7 +58,7 @@ public async ValueTask StartAsync()
 			}
 			catch (Exception ex)
 			{
-				_logger.LogError(ex, @"{$Local} Stop!", UdpListener.Client.LocalEndPoint);
+				_logger.LogError(ex, @"{Local} Stop!", UdpListener.Client.LocalEndPoint);
 				Stop();
 			}
 		}

diff --git a/Shadowsocks.Protocol/LocalUdpServices/Socks5UdpService.cs b/Shadowsocks.Protocol/LocalUdpServices/Socks5UdpService.cs
@@ -1,6 +1,5 @@
 using Microsoft.Extensions.Caching.Memory;
 using Microsoft.Extensions.Logging;
-using Microsoft.VisualStudio.Threading;
 using Shadowsocks.Protocol.ServersControllers;
 using Shadowsocks.Protocol.UdpClients;
 using Socks5.Enums;
@@ -30,9 +29,9 @@ public TimeSpan? SlidingExpiration
 			.RegisterPostEvictionCallback(
 				(key, value, reason, state) =>
 				{
-					if (value is System.IAsyncDisposable disposable)
+					if (value is IDisposable disposable)
 					{
-						disposable.DisposeAsync().Forget();
+						disposable.Dispose();
 					}
 				}
 			);

diff --git a/Shadowsocks.Protocol/UdpClients/IUdpClient.cs b/Shadowsocks.Protocol/UdpClients/IUdpClient.cs
@@ -1,14 +1,11 @@
 using System;
-using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Shadowsocks.Protocol.UdpClients
 {
-	public interface IUdpClient : IAsyncDisposable
+	public interface IUdpClient : IDisposable
 	{
-		UdpClient Client { get; }
-
 		ValueTask<int> ReceiveAsync(Memory<byte> buffer, CancellationToken cancellationToken = default);
 		ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default);
 	}

diff --git a/Shadowsocks.Protocol/UdpClients/ShadowsocksUdpClient.cs b/Shadowsocks.Protocol/UdpClients/ShadowsocksUdpClient.cs
@@ -1,4 +1,5 @@
 using Microsoft;
+using Pipelines.Extensions;
 using Shadowsocks.Crypto;
 using Shadowsocks.Protocol.Models;
 using System;
@@ -13,39 +14,51 @@ public class ShadowsocksUdpClient : IUdpClient
 	{
 		private readonly ShadowsocksServerInfo _serverInfo;
 
-		private const int MaxUDPSize = 0x10000;
+		private const int MaxUdpSize = 0x10000;
 
-		public UdpClient Client { get; }
+		private readonly Socket _client;
 
-		public ShadowsocksUdpClient(ShadowsocksServerInfo serverInfo, bool isIPv6 = false)
+		public ShadowsocksUdpClient(ShadowsocksServerInfo serverInfo)
 		{
+			Requires.NotNull(serverInfo, nameof(serverInfo));
+			Requires.NotNullAllowStructs(serverInfo.Method, nameof(serverInfo));
+			Requires.NotNullAllowStructs(serverInfo.Password, nameof(serverInfo));
+			Requires.NotNullAllowStructs(serverInfo.Address, nameof(serverInfo));
+
 			_serverInfo = serverInfo;
 
-			Client = new UdpClient(0, isIPv6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork);
+			_client = new Socket(SocketType.Dgram, ProtocolType.Udp);
+			_client.Connect(serverInfo.Address, serverInfo.Port);
 		}
 
 		public async ValueTask<int> ReceiveAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
 		{
-			//TODO .NET6.0
-			var res = await Client.ReceiveAsync();
+			var encBuffer = ArrayPool<byte>.Shared.Rent(MaxUdpSize);
+			try
+			{
+				var res = await _client.ReceiveAsync(encBuffer, SocketFlags.None, cancellationToken);
 
-			using var decryptor = ShadowsocksCrypto.Create(_serverInfo.Method!, _serverInfo.Password!);
+				using var decryptor = ShadowsocksCrypto.Create(_serverInfo.Method!, _serverInfo.Password!);
 
-			return decryptor.DecryptUDP(res.Buffer, buffer.Span);
+				return decryptor.DecryptUDP(encBuffer.AsSpan(0, res), buffer.Span);
+			}
+			finally
+			{
+				ArrayPool<byte>.Shared.Return(encBuffer);
+			}
 		}
 
 		public async ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
 		{
-			using var encryptor = ShadowsocksCrypto.Create(_serverInfo.Method!, _serverInfo.Password!);
-			var encBuffer = ArrayPool<byte>.Shared.Rent(MaxUDPSize);
+			var encBuffer = ArrayPool<byte>.Shared.Rent(MaxUdpSize);
 			try
 			{
+				using var encryptor = ShadowsocksCrypto.Create(_serverInfo.Method!, _serverInfo.Password!);
 				var length = encryptor.EncryptUDP(buffer.Span, encBuffer);
 
-				//TODO .NET6.0
-				var sendLength = await Client.SendAsync(encBuffer, length, _serverInfo.Address, _serverInfo.Port);
+				var sendLength = await _client.SendAsync(encBuffer.AsMemory(0, length), SocketFlags.None, cancellationToken);
 				Report.IfNot(sendLength == length, @"Send Udp {0}/{1}", sendLength, length);
-				return sendLength == length ? buffer.Length : 0;
+				return sendLength == length ? buffer.Length : default;
 			}
 			finally
 			{
@@ -58,11 +71,9 @@ public async ValueTask<int> SendAsync(ReadOnlyMemory<byte> buffer, CancellationT
 			return _serverInfo.ToString();
 		}
 
-		public ValueTask DisposeAsync()
+		public void Dispose()
 		{
-			Client.Dispose();
-
-			return default;
+			_client.FullClose();
 		}
 	}
 }