AesX86Engine improvements

3 files changed, 124 insertions, 56 deletions

diff --git a/crypto/src/crypto/Check.cs b/crypto/src/crypto/Check.cs
index aacda144f..81d07e23c 100644
--- a/crypto/src/crypto/Check.cs
+++ b/crypto/src/crypto/Check.cs
@@ -21,5 +21,19 @@ namespace Org.BouncyCastle.Crypto
             if (off > (buf.Length - len))
                 throw new OutputLengthException(msg);
         }
+
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        internal static void DataLength(Span<byte> input, int len, string msg)
+        {
+            if (input.Length < len)
+                throw new DataLengthException(msg);
+        }
+
+        internal static void OutputLength(Span<byte> output, int len, string msg)
+        {
+            if (output.Length < len)
+                throw new OutputLengthException(msg);
+        }
+#endif
     }
 }
diff --git a/crypto/src/crypto/engines/AesX86Engine.cs b/crypto/src/crypto/engines/AesX86Engine.cs
index b07e27bcf..d16019992 100644
--- a/crypto/src/crypto/engines/AesX86Engine.cs
+++ b/crypto/src/crypto/engines/AesX86Engine.cs
@@ -11,7 +11,7 @@ namespace Org.BouncyCastle.Crypto.Engines
     using Aes = System.Runtime.Intrinsics.X86.Aes;
     using Sse2 = System.Runtime.Intrinsics.X86.Sse2;
 
-    public class AesX86Engine
+    public struct AesX86Engine
         : IBlockCipher
     {
         public static bool IsSupported => Aes.IsSupported;
@@ -28,7 +28,7 @@ namespace Org.BouncyCastle.Crypto.Engines
 
                 K = new Vector128<byte>[11];
 
-                var s = Load128(key, 0);
+                var s = Load128(key.AsSpan(0, 16));
                 K[0] = s;
 
                 for (int round = 0; round < 10;)
@@ -47,8 +47,8 @@ namespace Org.BouncyCastle.Crypto.Engines
             {
                 K = new Vector128<byte>[13];
 
-                var s1 = Load128(key, 0);
-                var s2 = Load64(key, 16).ToVector128();
+                var s1 = Load128(key.AsSpan(0, 16));
+                var s2 = Load64(key.AsSpan(16, 8)).ToVector128();
                 K[0] = s1;
 
                 byte rcon = 0x01;
@@ -93,8 +93,8 @@ namespace Org.BouncyCastle.Crypto.Engines
             {
                 K = new Vector128<byte>[15];
 
-                var s1 = Load128(key, 0);
-                var s2 = Load128(key, 16);
+                var s1 = Load128(key.AsSpan(0, 16));
+                var s2 = Load128(key.AsSpan(16, 16));
                 K[0] = s1;
                 K[1] = s2;
 
@@ -140,7 +140,7 @@ namespace Org.BouncyCastle.Crypto.Engines
 
         private enum Mode { DEC_128, DEC_192, DEC_256, ENC_128, ENC_192, ENC_256, UNINITIALIZED };
 
-        private Vector128<byte>[] m_roundKeys;
+        private Vector128<byte>[] m_roundKeys = null;
         private Mode m_mode = Mode.UNINITIALIZED;
 
         public AesX86Engine()
@@ -149,11 +149,19 @@ namespace Org.BouncyCastle.Crypto.Engines
                 throw new PlatformNotSupportedException(nameof(AesX86Engine));
         }
 
-        public virtual void Init(bool forEncryption, ICipherParameters parameters)
+        public string AlgorithmName => "AES";
+
+        public bool IsPartialBlockOkay => false;
+
+        public int GetBlockSize() => 16;
+
+        public void Init(bool forEncryption, ICipherParameters parameters)
         {
             if (!(parameters is KeyParameter keyParameter))
-                throw new ArgumentException(
-                    "invalid parameter passed to AES Init - " + Platform.GetTypeName(parameters));
+            {
+                ArgumentNullException.ThrowIfNull(parameters, nameof(parameters));
+                throw new ArgumentException("invalid type: " + Platform.GetTypeName(parameters), nameof(parameters));
+            }
 
             m_roundKeys = CreateRoundKeys(keyParameter.GetKey(), forEncryption);
 
@@ -171,39 +179,50 @@ namespace Org.BouncyCastle.Crypto.Engines
             }
         }
 
-        public virtual string AlgorithmName => "AES";
-
-        public virtual bool IsPartialBlockOkay => false;
+        public int ProcessBlock(byte[] inBuf, int inOff, byte[] outBuf, int outOff)
+        {
+            Check.DataLength(inBuf, inOff, 16, "input buffer too short");
+            Check.OutputLength(outBuf, outOff, 16, "output buffer too short");
 
-        public virtual int GetBlockSize() => 16;
+            var state = Load128(inBuf.AsSpan(inOff, 16));
+            ImplRounds(ref state);
+            Store128(ref state, outBuf.AsSpan(outOff, 16));
+            return 16;
+        }
 
-        public virtual int ProcessBlock(byte[] input, int inOff, byte[] output, int outOff)
+        public int ProcessBlock(Span<byte> input, Span<byte> output)
         {
-            Check.DataLength(input, inOff, 16, "input buffer too short");
-            Check.OutputLength(output, outOff, 16, "output buffer too short");
-
-            switch (m_mode)
-            {
-            case Mode.DEC_128: Decrypt128(input, inOff, output, outOff, m_roundKeys); break;
-            case Mode.DEC_192: Decrypt192(input, inOff, output, outOff, m_roundKeys); break;
-            case Mode.DEC_256: Decrypt256(input, inOff, output, outOff, m_roundKeys); break;
-            case Mode.ENC_128: Encrypt128(input, inOff, output, outOff, m_roundKeys); break;
-            case Mode.ENC_192: Encrypt192(input, inOff, output, outOff, m_roundKeys); break;
-            case Mode.ENC_256: Encrypt256(input, inOff, output, outOff, m_roundKeys); break;
-            default: throw new InvalidOperationException("AES engine not initialised");
-            }
+            Check.DataLength(input, 16, "input buffer too short");
+            Check.OutputLength(output, 16, "output buffer too short");
 
+            var state = Load128(input);
+            ImplRounds(ref state);
+            Store128(ref state, output);
             return 16;
         }
 
-        public virtual void Reset()
+        public void Reset()
         {
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Decrypt128(byte[] input, int inOff, byte[] output, int outOff, Vector128<byte>[] roundKeys)
+        private void ImplRounds(ref Vector128<byte> state)
+        {
+            switch (m_mode)
+            {
+            case Mode.DEC_128: Decrypt128(ref state, m_roundKeys); break;
+            case Mode.DEC_192: Decrypt192(ref state, m_roundKeys); break;
+            case Mode.DEC_256: Decrypt256(ref state, m_roundKeys); break;
+            case Mode.ENC_128: Encrypt128(ref state, m_roundKeys); break;
+            case Mode.ENC_192: Encrypt192(ref state, m_roundKeys); break;
+            case Mode.ENC_256: Encrypt256(ref state, m_roundKeys); break;
+            default: throw new InvalidOperationException(nameof(AesX86Engine) + " not initialised");
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static void Decrypt128(ref Vector128<byte> state, Vector128<byte>[] roundKeys)
         {
-            var state = Load128(input, inOff);
             state = Sse2.Xor(state, roundKeys[0]);
             state = Aes.Decrypt(state, roundKeys[1]);
             state = Aes.Decrypt(state, roundKeys[2]);
@@ -215,13 +234,11 @@ namespace Org.BouncyCastle.Crypto.Engines
             state = Aes.Decrypt(state, roundKeys[8]);
             state = Aes.Decrypt(state, roundKeys[9]);
             state = Aes.DecryptLast(state, roundKeys[10]);
-            Store128(ref state, output, outOff);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Decrypt192(byte[] input, int inOff, byte[] output, int outOff, Vector128<byte>[] roundKeys)
+        private static void Decrypt192(ref Vector128<byte> state, Vector128<byte>[] roundKeys)
         {
-            var state = Load128(input, inOff);
             state = Sse2.Xor(state, roundKeys[0]);
             state = Aes.Decrypt(state, roundKeys[1]);
             state = Aes.Decrypt(state, roundKeys[2]);
@@ -235,13 +252,11 @@ namespace Org.BouncyCastle.Crypto.Engines
             state = Aes.Decrypt(state, roundKeys[10]);
             state = Aes.Decrypt(state, roundKeys[11]);
             state = Aes.DecryptLast(state, roundKeys[12]);
-            Store128(ref state, output, outOff);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Decrypt256(byte[] input, int inOff, byte[] output, int outOff, Vector128<byte>[] roundKeys)
+        private static void Decrypt256(ref Vector128<byte> state, Vector128<byte>[] roundKeys)
         {
-            var state = Load128(input, inOff);
             state = Sse2.Xor(state, roundKeys[0]);
             state = Aes.Decrypt(state, roundKeys[1]);
             state = Aes.Decrypt(state, roundKeys[2]);
@@ -257,13 +272,11 @@ namespace Org.BouncyCastle.Crypto.Engines
             state = Aes.Decrypt(state, roundKeys[12]);
             state = Aes.Decrypt(state, roundKeys[13]);
             state = Aes.DecryptLast(state, roundKeys[14]);
-            Store128(ref state, output, outOff);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Encrypt128(byte[] input, int inOff, byte[] output, int outOff, Vector128<byte>[] roundKeys)
+        private static void Encrypt128(ref Vector128<byte> state, Vector128<byte>[] roundKeys)
         {
-            var state = Load128(input, inOff);
             state = Sse2.Xor(state, roundKeys[0]);
             state = Aes.Encrypt(state, roundKeys[1]);
             state = Aes.Encrypt(state, roundKeys[2]);
@@ -275,13 +288,11 @@ namespace Org.BouncyCastle.Crypto.Engines
             state = Aes.Encrypt(state, roundKeys[8]);
             state = Aes.Encrypt(state, roundKeys[9]);
             state = Aes.EncryptLast(state, roundKeys[10]);
-            Store128(ref state, output, outOff);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Encrypt192(byte[] input, int inOff, byte[] output, int outOff, Vector128<byte>[] roundKeys)
+        private static void Encrypt192(ref Vector128<byte> state, Vector128<byte>[] roundKeys)
         {
-            var state = Load128(input, inOff);
             state = Sse2.Xor(state, roundKeys[0]);
             state = Aes.Encrypt(state, roundKeys[1]);
             state = Aes.Encrypt(state, roundKeys[2]);
@@ -295,13 +306,11 @@ namespace Org.BouncyCastle.Crypto.Engines
             state = Aes.Encrypt(state, roundKeys[10]);
             state = Aes.Encrypt(state, roundKeys[11]);
             state = Aes.EncryptLast(state, roundKeys[12]);
-            Store128(ref state, output, outOff);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Encrypt256(byte[] input, int inOff, byte[] output, int outOff, Vector128<byte>[] roundKeys)
+        private static void Encrypt256(ref Vector128<byte> state, Vector128<byte>[] roundKeys)
         {
-            var state = Load128(input, inOff);
             state = Sse2.Xor(state, roundKeys[0]);
             state = Aes.Encrypt(state, roundKeys[1]);
             state = Aes.Encrypt(state, roundKeys[2]);
@@ -317,36 +326,50 @@ namespace Org.BouncyCastle.Crypto.Engines
             state = Aes.Encrypt(state, roundKeys[12]);
             state = Aes.Encrypt(state, roundKeys[13]);
             state = Aes.EncryptLast(state, roundKeys[14]);
-            Store128(ref state, output, outOff);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static Vector128<byte> Load128(byte[] b, int n)
+        private static Vector128<byte> Load128(Span<byte> t)
         {
 #if NET7_0_OR_GREATER
-            return Vector128.Create(b, n);
+            return Vector128.Create(t);
 #else
-            return Unsafe.ReadUnaligned<Vector128<byte>>(ref b[n]);
+            if (BitConverter.IsLittleEndian && Unsafe.SizeOf<Vector128<byte>>() == 16)
+                return Unsafe.ReadUnaligned<Vector128<byte>>(ref t[0]);
+
+            return Vector128.Create(t[0], t[1], t[2], t[3], t[4], t[5], t[6], t[7], t[8], t[9], t[10], t[11], t[12],
+                t[13], t[14], t[15]);
 #endif
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static Vector64<byte> Load64(byte[] b, int n)
+        private static Vector64<byte> Load64(Span<byte> t)
         {
 #if NET7_0_OR_GREATER
-            return Vector64.Create(b, n);
+            return Vector64.Create(t);
 #else
-            return Unsafe.ReadUnaligned<Vector64<byte>>(ref b[n]);
+            if (BitConverter.IsLittleEndian && Unsafe.SizeOf<Vector64<byte>>() == 8)
+                return Unsafe.ReadUnaligned<Vector64<byte>>(ref t[0]);
+
+            return Vector64.Create(t[0], t[1], t[2], t[3], t[4], t[5], t[6], t[7]);
 #endif
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void Store128(ref Vector128<byte> s, byte[] b, int n)
+        private static void Store128(ref Vector128<byte> s, Span<byte> t)
         {
 #if NET7_0_OR_GREATER
-            Vector128.CopyTo(s, b, n);
+            Vector128.CopyTo(s, t);
 #else
-            Unsafe.WriteUnaligned(ref b[n], s);
+            if (BitConverter.IsLittleEndian && Unsafe.SizeOf<Vector128<byte>>() == 16)
+            {
+                Unsafe.WriteUnaligned(ref t[0], s);
+                return;
+            }
+
+            var u = s.AsUInt64();
+            Utilities.Pack.UInt64_To_LE(u.GetElement(0), t);
+            Utilities.Pack.UInt64_To_LE(u.GetElement(1), t.Slice(8));
 #endif
         }
     }
diff --git a/crypto/src/crypto/util/Pack.cs b/crypto/src/crypto/util/Pack.cs
index b7cfde905..9fd533cfc 100644
--- a/crypto/src/crypto/util/Pack.cs
+++ b/crypto/src/crypto/util/Pack.cs
@@ -439,5 +439,36 @@ namespace Org.BouncyCastle.Crypto.Utilities
                 bsOff += 8;
             }
         }
+
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        internal static uint LE_To_UInt32(Span<byte> s)
+        {
+            return s[0]
+                | (uint)s[1] << 8
+                | (uint)s[2] << 16
+                | (uint)s[3] << 24;
+        }
+
+        internal static ulong LE_To_UInt64(Span<byte> s)
+        {
+            uint lo = LE_To_UInt32(s);
+            uint hi = LE_To_UInt32(s.Slice(4));
+            return (ulong)hi << 32 | lo;
+        }
+
+        internal static void UInt32_To_LE(uint n, Span<byte> s)
+        {
+            s[0] = (byte)n;
+            s[1] = (byte)(n >> 8);
+            s[2] = (byte)(n >> 16);
+            s[3] = (byte)(n >> 24);
+        }
+
+        internal static void UInt64_To_LE(ulong n, Span<byte> s)
+        {
+            UInt32_To_LE((uint)n, s);
+            UInt32_To_LE((uint)(n >> 32), s.Slice(4));
+        }
+#endif
     }
 }