Optimise AESCipher to use ArrayPool

Kaioru · Oct 12, 2023 · 4647757 · 4647757
1 parent 8059b4a
commit 4647757
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/src/common/Edelstein.Common.Crypto/AESCipher.cs b/src/common/Edelstein.Common.Crypto/AESCipher.cs
@@ -1,4 +1,5 @@
-using System.Security.Cryptography;
+using System.Buffers;
+using System.Security.Cryptography;
 
 namespace Edelstein.Common.Crypto;
 
@@ -28,14 +29,15 @@ public void Transform(Span<byte> input, int remaining, uint pSrc)
         var length = 0x5B0;
         var start = 0;
 
-        var srcExp = new byte[sizeof(int) * 4];
+        var srcExpL = sizeof(int) * 4;
+        var srcExp = ArrayPool<byte>.Shared.Rent(srcExpL);
         var srcBytes = BitConverter.GetBytes(pSrc);
 
         using var crypt = _cipher.CreateEncryptor();
 
         while (remaining > 0)
         {
-            for (var i = 0; i < srcExp.Length; ++i)
+            for (var i = 0; i < srcExpL; ++i)
                 srcExp[i] = srcBytes[i % 4];
 
             if (remaining < length)
@@ -45,18 +47,17 @@ public void Transform(Span<byte> input, int remaining, uint pSrc)
             {
                 var sub = i - start;
 
-                if (sub % srcExp.Length == 0)
-                {
-                    var result = crypt.TransformFinalBlock(srcExp, 0, srcExp.Length);
-                    Array.Copy(result, srcExp, srcExp.Length);
-                }
+                if (sub % srcExpL == 0)
+                    crypt.TransformBlock(srcExp, 0, srcExpL, srcExp, 0);
 
-                input[i] ^= srcExp[sub % srcExp.Length];
+                input[i] ^= srcExp[sub % srcExpL];
             }
 
             start += length;
             remaining -= length;
             length = 0x5B4;
         }
+
+        ArrayPool<byte>.Shared.Return(srcExp);
     }
 }