fix: SOCKS5 UDP service

HMBSbige · Aug 24, 2021 · e317674 · e317674
1 parent e1ce383
commit e317674
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/Shadowsocks.Protocol/LocalTcpServices/Socks5Service.cs b/Shadowsocks.Protocol/LocalTcpServices/Socks5Service.cs
@@ -99,6 +99,10 @@ public async ValueTask HandleAsync(IDuplexPipe pipe, CancellationToken token = d
 						Port = Socks5CreateOption.Port
 					};
 					await socks5.SendReplyAsync(Socks5Reply.Succeeded, bound, token);
+
+					// wait remote close
+					var result = await pipe.Input.ReadAsync(token);
+					Report.IfNot(result.IsCompleted);
 					break;
 				}
 				default:

diff --git a/Shadowsocks.Protocol/LocalUdpServices/Socks5UdpService.cs b/Shadowsocks.Protocol/LocalUdpServices/Socks5UdpService.cs
@@ -2,6 +2,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.Threading;
 using Shadowsocks.Protocol.ServersControllers;
+using Shadowsocks.Protocol.UdpClients;
 using Socks5.Enums;
 using Socks5.Utils;
 using System;
@@ -69,9 +70,12 @@ public async ValueTask HandleAsync(UdpReceiveResult receiveResult, UdpClient inc
 					_ => socks5UdpPacket.Address!.ToString()
 				};
 
-				var client = await _cache.GetOrCreateAsync(
-					receiveResult.RemoteEndPoint,
-					_ => _serversController.GetServerUdpAsync(target).AsTask());
+				if (!_cache.TryGetValue(receiveResult.RemoteEndPoint, out IUdpClient client))
+				{
+					client = await _serversController.GetServerUdpAsync(target);
+
+					_cache.Set(receiveResult.RemoteEndPoint, client, _cacheOptions);
+				}
 
 				_logger.LogInformation(@"Udp Send to {0} via {1}", target, client);
 				var sendBuffer = receiveResult.Buffer.AsMemory(3); //TODO Only support ss now