Optimise packet decoder and encoder to use ArrayPool

Kaioru · Oct 12, 2023 · 45e8a57 · 45e8a57
1 parent 80dfb76
commit 45e8a57
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 26 deletions.
diff --git a/src/common/Edelstein.Common.Crypto/AESCipher.cs b/src/common/Edelstein.Common.Crypto/AESCipher.cs
@@ -23,9 +23,8 @@ public AESCipher(ReadOnlySpan<byte> userKey)
         _cipher.Mode = CipherMode.ECB;
     }
 
-    public void Transform(Span<byte> input, uint pSrc)
+    public void Transform(Span<byte> input, int remaining, uint pSrc)
     {
-        var remaining = input.Length;
         var length = 0x5B0;
         var start = 0;
 

diff --git a/src/common/Edelstein.Common.Crypto/ShandaCipher.cs b/src/common/Edelstein.Common.Crypto/ShandaCipher.cs
@@ -2,9 +2,8 @@
 
 public class ShandaCipher
 {
-    public static void EncryptTransform(Span<byte> input)
+    public static void EncryptTransform(Span<byte> input, int size)
     {
-        var size = input.Length;
         for (var i = 0; i < 3; i++)
         {
             byte a = 0;
@@ -23,7 +22,7 @@ public static void EncryptTransform(Span<byte> input)
             }
 
             a = 0;
-            for (var j = input.Length; j > 0; j--)
+            for (var j = size; j > 0; j--)
             {
                 c = input[j - 1];
                 c = RollLeft(c, 4);
@@ -37,9 +36,8 @@ public static void EncryptTransform(Span<byte> input)
         }
     }
 
-    public static void DecryptTransform(Span<byte> input)
+    public static void DecryptTransform(Span<byte> input, int size)
     {
-        var size = input.Length;
         for (var i = 0; i < 3; i++)
         {
             byte a;

diff --git a/src/common/Edelstein.Common.Network.DotNetty/Codecs/NettyPacketDecoder.cs b/src/common/Edelstein.Common.Network.DotNetty/Codecs/NettyPacketDecoder.cs
@@ -1,4 +1,5 @@
-using DotNetty.Buffers;
+using System.Buffers;
+using DotNetty.Buffers;
 using DotNetty.Codecs;
 using DotNetty.Transport.Channels;
 using Edelstein.Common.Crypto;
@@ -70,9 +71,9 @@ protected override void Decode(IChannelHandlerContext context, IByteBuffer input
                     return;
                 }
 
-                var buffer = new byte[_length];
+                var buffer = ArrayPool<byte>.Shared.Rent(_length);
 
-                input.ReadBytes(buffer);
+                input.ReadBytes(buffer, 0, _length);
                 Checkpoint(NettyPacketState.DecodingHeader);
 
                 if (_length < 0x2) return;
@@ -87,8 +88,8 @@ protected override void Decode(IChannelHandlerContext context, IByteBuffer input
 
                     if (socket.IsDataEncrypted)
                     {
-                        _aesCipher.Transform(buffer, seqRecv);
-                        ShandaCipher.DecryptTransform(buffer);
+                        _aesCipher.Transform(buffer, _length, seqRecv);
+                        ShandaCipher.DecryptTransform(buffer, _length);
                     }
 
                     socket.SeqRecv = _igCipher.Hash(seqRecv, 4, 0);

diff --git a/src/common/Edelstein.Common.Network.DotNetty/Codecs/NettyPacketEncoder.cs b/src/common/Edelstein.Common.Network.DotNetty/Codecs/NettyPacketEncoder.cs
@@ -1,4 +1,5 @@
-using DotNetty.Buffers;
+using System.Buffers;
+using DotNetty.Buffers;
 using DotNetty.Codecs;
 using DotNetty.Transport.Channels;
 using Edelstein.Common.Crypto;
@@ -32,31 +33,33 @@ IByteBuffer output
     {
         var socket = context.Channel.GetAttribute(NettyAttributes.SocketKey).Get();
         var dataLen = (short)message.Buffer.Length;
-        var buffer = message.Buffer.ToArray();
-
+        var buffer = ArrayPool<byte>.Shared.Rent(dataLen);
+
+        Array.Copy(message.Buffer, buffer, dataLen);
+
         if (socket != null)
         {
             var seqSend = socket.SeqSend;
             var rawSeq = (short)(seqSend >> 16 ^ -(_version.Major + 1));
 
             if (socket.IsDataEncrypted)
             {
-                dataLen ^= rawSeq;
-
-                ShandaCipher.EncryptTransform(buffer);
-                _aesCipher.Transform(buffer, seqSend);
+                ShandaCipher.EncryptTransform(buffer, dataLen);
+                _aesCipher.Transform(buffer, dataLen, seqSend);
             }
 
             output.WriteShortLE(rawSeq);
-            output.WriteShortLE(dataLen);
-            output.WriteBytes(buffer);
+            output.WriteShortLE(dataLen ^ rawSeq);
+            output.WriteBytes(buffer, 0, dataLen);
 
             socket.SeqSend = _igCipher.Hash(seqSend, 4, 0);
         }
         else
         {
             output.WriteShortLE(dataLen);
-            output.WriteBytes(buffer);
+            output.WriteBytes(buffer, 0, dataLen);
         }
+
+        ArrayPool<byte>.Shared.Return(buffer);
     }
 }
diff --git a/src/common/Edelstein.Common.Network.DotNetty/Handlers/NettyTransportAcceptorHandler.cs b/src/common/Edelstein.Common.Network.DotNetty/Handlers/NettyTransportAcceptorHandler.cs
@@ -1,4 +1,5 @@
-using DotNetty.Transport.Channels;
+using System.Buffers;
+using DotNetty.Transport.Channels;
 using Edelstein.Common.Utilities.Packets;
 using Edelstein.Protocol.Network;
 using Edelstein.Protocol.Network.Transports;
@@ -68,13 +69,15 @@ public override void ChannelInactive(IChannelHandlerContext context)
     public override void ChannelRead(IChannelHandlerContext context, object message)
     {
         var adapter = context.Channel.GetAttribute(NettyAttributes.AdapterKey).Get();
-        adapter?.OnPacket((IPacket)message);
+        var packet = (IPacket)message;
+
+        adapter?.OnPacket(packet);
+        ArrayPool<byte>.Shared.Return(packet.Buffer);
     }
 
     public override void ExceptionCaught(IChannelHandlerContext context, Exception exception)
     {
         var adapter = context.Channel.GetAttribute(NettyAttributes.AdapterKey).Get();
-
         adapter?.OnException(exception);
     }
 }
diff --git a/src/common/Edelstein.Common.Network.DotNetty/Handlers/NettyTransportConnectorHandler.cs b/src/common/Edelstein.Common.Network.DotNetty/Handlers/NettyTransportConnectorHandler.cs
@@ -1,4 +1,5 @@
-using DotNetty.Transport.Channels;
+using System.Buffers;
+using DotNetty.Transport.Channels;
 using Edelstein.Common.Utilities.Packets;
 using Edelstein.Protocol.Network;
 using Edelstein.Protocol.Network.Transports;
@@ -54,6 +55,8 @@ public override void ChannelRead(IChannelHandlerContext context, object message)
 
             _ = _sockets.Insert(newSocket);
         }
+
+        ArrayPool<byte>.Shared.Return(packet.Buffer);
     }
 
     public override void ChannelInactive(IChannelHandlerContext context)