From bff6760a9b0d6e6306783b6890b03a155c6fa5ad Mon Sep 17 00:00:00 2001 From: Student Main Date: Sat, 25 Jul 2020 16:32:19 +0800 Subject: [PATCH] async await style udp operation --- .../Controller/Service/UDPListener.cs | 12 +--- .../Controller/Service/UDPRelay.cs | 61 +++++++++---------- 2 files changed, 31 insertions(+), 42 deletions(-) diff --git a/shadowsocks-csharp/Controller/Service/UDPListener.cs b/shadowsocks-csharp/Controller/Service/UDPListener.cs index 52f375a9..0ed972fe 100644 --- a/shadowsocks-csharp/Controller/Service/UDPListener.cs +++ b/shadowsocks-csharp/Controller/Service/UDPListener.cs @@ -13,20 +13,14 @@ namespace Shadowsocks.Controller { public interface IDatagramService { - [Obsolete] - bool Handle(byte[] firstPacket, int length, Socket socket, object state); - - public abstract bool Handle(CachedNetworkStream stream, object state); + public abstract Task< bool> Handle(Memory packet, Socket socket, EndPoint client); void Stop(); } public abstract class DatagramService : IDatagramService { - [Obsolete] - public abstract bool Handle(byte[] firstPacket, int length, Socket socket, object state); - - public abstract bool Handle(CachedNetworkStream stream, object state); + public abstract Task Handle(Memory packet, Socket socket, EndPoint client); public virtual void Stop() { } } @@ -101,7 +95,7 @@ namespace Shadowsocks.Controller var len = result.ReceivedBytes; foreach (IDatagramService service in _services) { - if (service.Handle(buffer, len, _udpSocket, result.RemoteEndPoint)) + if (await service.Handle(new Memory(buffer)[..len], _udpSocket, result.RemoteEndPoint)) { break; } diff --git a/shadowsocks-csharp/Controller/Service/UDPRelay.cs b/shadowsocks-csharp/Controller/Service/UDPRelay.cs index 3c0e5afa..3f8e6f2c 100644 --- a/shadowsocks-csharp/Controller/Service/UDPRelay.cs +++ b/shadowsocks-csharp/Controller/Service/UDPRelay.cs @@ -1,8 +1,10 @@ using System; +using System.Buffers; using System.Collections.Generic; using System.Net; using System.Net.Sockets; using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; using System.Threading.Tasks; using NLog; using Shadowsocks.Controller.Strategy; @@ -26,28 +28,17 @@ namespace Shadowsocks.Controller this._controller = controller; } - // TODO: UDP is datagram protocol not stream protocol - public override bool Handle(CachedNetworkStream stream, object state) - { - byte[] fp = new byte[256]; - int len = stream.ReadFirstBlock(fp); - return Handle(fp, len, stream.Socket, state); - } - - [Obsolete] - public override bool Handle(byte[] firstPacket, int length, Socket socket, object state) + public override async Task Handle(Memory packet, Socket socket, EndPoint client) { if (socket.ProtocolType != ProtocolType.Udp) { return false; } - if (length < 4) + if (packet.Length < 4) { return false; } - // UDPListener.UDPState udpState = (UDPListener.UDPState)state; - IPEndPoint remoteEndPoint = (IPEndPoint)state; - // IPEndPoint remoteEndPoint = (IPEndPoint)socket.RemoteEndPoint; + IPEndPoint remoteEndPoint = (IPEndPoint)client; UDPHandler handler = _cache.get(remoteEndPoint); if (handler == null) { @@ -55,14 +46,14 @@ namespace Shadowsocks.Controller handler.Receive(); _cache.add(remoteEndPoint, handler); } - handler.Send(firstPacket, length); + await handler.SendAsync(packet); return true; } public class UDPHandler { private static Logger logger = LogManager.GetCurrentClassLogger(); - + private static MemoryPool pool = MemoryPool.Shared; private Socket _local; private Socket _remote; @@ -92,8 +83,7 @@ namespace Shadowsocks.Controller _localEndPoint = localEndPoint; // TODO async resolving - IPAddress ipAddress; - bool parsed = IPAddress.TryParse(server.server, out ipAddress); + bool parsed = IPAddress.TryParse(server.server, out IPAddress ipAddress); if (!parsed) { IPHostEntry ipHostInfo = Dns.GetHostEntry(server.server); @@ -104,15 +94,19 @@ namespace Shadowsocks.Controller _remote.Bind(new IPEndPoint(ListenAddress, 0)); } - public void Send(byte[] data, int length) + public async Task SendAsync(ReadOnlyMemory data) { IEncryptor encryptor = EncryptorFactory.GetEncryptor(_server.method, _server.password); - var slicedData = data.AsSpan(0, length); - byte[] dataOut = new byte[slicedData.Length + 1000]; - var dataToSend = slicedData[3..]; - int outlen = encryptor.EncryptUDP(slicedData[3..], dataOut); + using IMemoryOwner mem = pool.Rent(data.Length + 1000); + + // byte[] dataOut = new byte[slicedData.Length + 1000]; + int outlen = encryptor.EncryptUDP(data.Span[3..], mem.Memory.Span); logger.Debug(_localEndPoint, _remoteEndPoint, outlen, "UDP Relay up"); - _remote?.SendTo(dataOut, outlen, SocketFlags.None, _remoteEndPoint); + if (!MemoryMarshal.TryGetArray(mem.Memory[..outlen], out ArraySegment outData)) + { + throw new InvalidOperationException("Can't extract underly array segment"); + }; + await _remote?.SendToAsync(outData, SocketFlags.None, _remoteEndPoint); } public async Task ReceiveAsync() @@ -121,22 +115,23 @@ namespace Shadowsocks.Controller logger.Debug($"++++++Receive Server Port, size:" + _buffer.Length); try { - while (true) { - var result = await _remote.ReceiveFromAsync(_buffer, SocketFlags.None, remoteEndPoint); int bytesRead = result.ReceivedBytes; - byte[] dataOut = new byte[bytesRead]; - int outlen; - IEncryptor encryptor = EncryptorFactory.GetEncryptor(_server.method, _server.password); - outlen = encryptor.DecryptUDP(dataOut, _buffer.AsSpan(0, bytesRead)); - byte[] sendBuf = new byte[outlen + 3]; - Array.Copy(dataOut, 0, sendBuf, 3, outlen); + using IMemoryOwner owner = pool.Rent(bytesRead + 3); + Memory o = owner.Memory; + IEncryptor encryptor = EncryptorFactory.GetEncryptor(_server.method, _server.password); + int outlen = encryptor.DecryptUDP(o.Span[3..], _buffer.AsSpan(0, bytesRead)); logger.Debug(_remoteEndPoint, _localEndPoint, outlen, "UDP Relay down"); - await _local?.SendToAsync(sendBuf, SocketFlags.None, _localEndPoint); + if (!MemoryMarshal.TryGetArray(o[..(outlen + 3)], out ArraySegment data)) + { + throw new InvalidOperationException("Can't extract underly array segment"); + }; + await _local?.SendToAsync(data, SocketFlags.None, _localEndPoint); + } } catch (Exception e)