From 1a2f601cd0a7a555dc90cb1374fcab66c56e2f66 Mon Sep 17 00:00:00 2001
From: Matthew Endsley <mendsley@gmail.com>
Date: Thu, 29 Jul 2021 22:23:14 -0700
Subject: [PATCH] Allow the default AES implementation to be overridden

---
 Hazel/Crypto/AesGcm.cs         | 18 +++----------
 Hazel/Crypto/CryptoProvider.cs | 36 +++++++++++++++++++++++++
 Hazel/Crypto/DefaultAes.cs     | 49 ++++++++++++++++++++++++++++++++++
 Hazel/Crypto/IAes.cs           | 27 +++++++++++++++++++
 Hazel/Hazel.csproj             |  3 +++
 5 files changed, 119 insertions(+), 14 deletions(-)
 create mode 100644 Hazel/Crypto/CryptoProvider.cs
 create mode 100644 Hazel/Crypto/DefaultAes.cs
 create mode 100644 Hazel/Crypto/IAes.cs

diff --git a/Hazel/Crypto/AesGcm.cs b/Hazel/Crypto/AesGcm.cs
index 51ec281..bfbbc01 100644
--- a/Hazel/Crypto/AesGcm.cs
+++ b/Hazel/Crypto/AesGcm.cs
@@ -22,7 +22,7 @@ namespace Hazel.Crypto
 
         private const int TagSize = 16;
 
-        private readonly ICryptoTransform encryptor_;
+        private readonly IAes encryptor_;
 
         private readonly ByteSpan hashSubkey_;
         private readonly ByteSpan blockJ_;
@@ -43,17 +43,7 @@ namespace Hazel.Crypto
             }
 
             // Create the AES block cipher
-            using (Aes aes = Aes.Create())
-            {
-                aes.KeySize = 128;
-                aes.KeySize = 128;
-                aes.BlockSize = 128;
-                aes.Mode = CipherMode.ECB;
-                aes.Padding = PaddingMode.Zeros;
-                aes.Key = key.ToArray();
-
-                this.encryptor_ = aes.CreateEncryptor();
-            }
+            this.encryptor_ = CryptoProvider.CreateAes(key);
 
             // Allocate scratch space
             ByteSpan scratchSpace = new byte[96];
@@ -65,7 +55,7 @@ namespace Hazel.Crypto
             this.blockScratch_ = scratchSpace.Slice(80, 16);
 
             // Create the GHASH subkey by encrypting the 0-block
-            this.encryptor_.TransformBlock(this.hashSubkey_.GetUnderlyingArray(), this.hashSubkey_.Offset, this.hashSubkey_.Length, this.hashSubkey_.GetUnderlyingArray(), this.hashSubkey_.Offset);
+            this.encryptor_.EncryptBlock(this.hashSubkey_, this.hashSubkey_);
         }
 
         /// <summary>
@@ -235,7 +225,7 @@ namespace Hazel.Crypto
                 ++counter;
 
                 // CIPH[k](CB[i])
-                this.encryptor_.TransformBlock(counterBlock.GetUnderlyingArray(), counterBlock.Offset, 16, this.blockScratch_.GetUnderlyingArray(), this.blockScratch_.Offset);
+                this.encryptor_.EncryptBlock(counterBlock.Slice(0, 16), this.blockScratch_);
 
                 // Y[i] = X[i] xor CIPH[k](CB[i])
                 for (int jj = 0; jj != 16 && writeIndex < data.Length; ++jj, ++writeIndex)
diff --git a/Hazel/Crypto/CryptoProvider.cs b/Hazel/Crypto/CryptoProvider.cs
new file mode 100644
index 0000000..2c56c70
--- /dev/null
+++ b/Hazel/Crypto/CryptoProvider.cs
@@ -0,0 +1,36 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Hazel.Crypto
+{
+    public static class CryptoProvider
+    {
+        public delegate IAes CreateAesOverrideDelegate(ByteSpan key);
+
+        /// <summary>
+        /// Override the default AES creation function
+        /// </summary>
+        public static CreateAesOverrideDelegate OverrideCreateAes = null;
+
+        /// <summary>
+        /// Create a new AES cipher
+        /// </summary>
+        /// <param name="key">Encrtyption key</param>
+        public static IAes CreateAes(ByteSpan key)
+        {
+            if (OverrideCreateAes != null)
+            {
+                IAes result = OverrideCreateAes(key);
+                if (null != result)
+                {
+                    return result;
+                }
+            }
+
+            return new DefaultAes(key);
+        }
+    }
+}
diff --git a/Hazel/Crypto/DefaultAes.cs b/Hazel/Crypto/DefaultAes.cs
new file mode 100644
index 0000000..da72fb8
--- /dev/null
+++ b/Hazel/Crypto/DefaultAes.cs
@@ -0,0 +1,49 @@
+using System;
+using System.Security.Cryptography;
+
+namespace Hazel.Crypto
+{
+    /// <summary>
+    /// AES provider using the default System.Security.Cryptography implementation
+    /// </summary>
+    public class DefaultAes : IAes
+    {
+        private readonly ICryptoTransform encryptor_;
+
+        /// <summary>
+        /// Create a new default instance of the AES block cipher
+        /// </summary>
+        /// <param name="key">Encryption key</param>
+        public DefaultAes(ByteSpan key)
+        {
+            // Create the AES block cipher
+            using (Aes aes = Aes.Create())
+            {
+                aes.KeySize = key.Length * 8;
+                aes.BlockSize = aes.KeySize;
+                aes.Mode = CipherMode.ECB;
+                aes.Padding = PaddingMode.Zeros;
+                aes.Key = key.ToArray();
+
+                this.encryptor_ = aes.CreateEncryptor();
+            }
+        }
+
+        /// <inheritdoc/>
+        public void Dispose()
+        {
+            this.encryptor_.Dispose();
+        }
+
+        /// <inheritdoc/>
+        public int EncryptBlock(ByteSpan inputSpan, ByteSpan outputSpan)
+        {
+            if (inputSpan.Length != outputSpan.Length)
+            {
+                throw new ArgumentException($"ouputSpan length ({outputSpan.Length}) does not match inputSpan length ({inputSpan.Length})", nameof(outputSpan));
+            }
+
+            return this.encryptor_.TransformBlock(inputSpan.GetUnderlyingArray(), inputSpan.Offset, inputSpan.Length, outputSpan.GetUnderlyingArray(), outputSpan.Offset);
+        }
+    }
+}
diff --git a/Hazel/Crypto/IAes.cs b/Hazel/Crypto/IAes.cs
new file mode 100644
index 0000000..6c494cd
--- /dev/null
+++ b/Hazel/Crypto/IAes.cs
@@ -0,0 +1,27 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Hazel.Crypto
+{
+    /// <summary>
+    /// AES encryption interface
+    /// </summary>
+    public interface IAes : IDisposable
+    {
+        /// <summary>
+        /// Encrypts the specified region of the input byte array and copies
+        /// the resulting transform to the specified region of the output
+        /// array.
+        /// </summary>
+        /// <param name="inputSpan">The input for which to encrypt</param>
+        /// <param name="outputSpan">
+        /// The otput to which to write the encrypted data. This span can
+        /// overlap with `inputSpan`.
+        /// </param>
+        /// <returns>The number of bytes written</returns>
+        int EncryptBlock(ByteSpan inputSpan, ByteSpan outputSpan);
+    }
+}
diff --git a/Hazel/Hazel.csproj b/Hazel/Hazel.csproj
index c94b972..6b0b91d 100644
--- a/Hazel/Hazel.csproj
+++ b/Hazel/Hazel.csproj
@@ -75,6 +75,9 @@
     <Compile Include="ConnectionState.cs" />
     <Compile Include="Crypto\AesGcm.cs" />
     <Compile Include="Crypto\Const.cs" />
+    <Compile Include="Crypto\CryptoProvider.cs" />
+    <Compile Include="Crypto\DefaultAes.cs" />
+    <Compile Include="Crypto\IAes.cs" />
     <Compile Include="Crypto\Sha256Stream.cs" />
     <Compile Include="Crypto\SpanCryptoExtensions.cs" />
     <Compile Include="Crypto\X25519.cs" />
-- 
2.39.5