Core/Crypto: Allow different AES key lengths

2 files changed, 32 insertions, 3 deletions

diff --git a/src/common/Cryptography/AES.cpp b/src/common/Cryptography/AES.cpp
index 29ccfd0cf06..38a20af506b 100644
--- a/src/common/Cryptography/AES.cpp
+++ b/src/common/Cryptography/AES.cpp
@@ -19,10 +19,26 @@
 #include "Errors.h"
 #include <limits>
 
-Trinity::Crypto::AES::AES(bool encrypting) : _ctx(EVP_CIPHER_CTX_new()), _encrypting(encrypting)
+Trinity::Crypto::AES::AES(bool encrypting, size_t keySizeBits /*= 128*/) : _ctx(EVP_CIPHER_CTX_new()), _encrypting(encrypting)
 {
     EVP_CIPHER_CTX_init(_ctx);
-    int status = EVP_CipherInit_ex(_ctx, EVP_aes_128_gcm(), nullptr, nullptr, nullptr, _encrypting ? 1 : 0);
+    EVP_CIPHER const* cipher = nullptr;
+    switch (keySizeBits)
+    {
+        case 128:
+            cipher = EVP_aes_128_gcm();
+            break;
+        case 192:
+            cipher = EVP_aes_192_gcm();
+            break;
+        case 256:
+            cipher = EVP_aes_256_gcm();
+            break;
+        default:
+            ASSERT(false, "Invalid AES key size " SZFMTD, keySizeBits);
+    }
+
+    int status = EVP_CipherInit_ex(_ctx, cipher, nullptr, nullptr, nullptr, _encrypting ? 1 : 0);
     ASSERT(status);
 }
 
@@ -37,6 +53,13 @@ void Trinity::Crypto::AES::Init(Key const& key)
     ASSERT(status);
 }
 
+void Trinity::Crypto::AES::Init(std::span<uint8 const> key)
+{
+    ASSERT(key.size() == size_t(EVP_CIPHER_CTX_get_key_length(_ctx)));
+    int status = EVP_CipherInit_ex(_ctx, nullptr, nullptr, key.data(), nullptr, -1);
+    ASSERT(status);
+}
+
 bool Trinity::Crypto::AES::Process(IV const& iv, uint8* data, size_t length, Tag& tag)
 {
     ASSERT(length <= static_cast<size_t>(std::numeric_limits<int>::max()));
diff --git a/src/common/Cryptography/AES.h b/src/common/Cryptography/AES.h
index b7ff64c11d0..db98df24d7c 100644
--- a/src/common/Cryptography/AES.h
+++ b/src/common/Cryptography/AES.h
@@ -20,6 +20,7 @@
 
 #include "Define.h"
 #include <array>
+#include <span>
 #include <openssl/evp.h>
 
 namespace Trinity::Crypto
@@ -35,10 +36,15 @@ namespace Trinity::Crypto
         using Key = std::array<uint8, KEY_SIZE_BYTES>;
         using Tag = uint8[TAG_SIZE_BYTES];
 
-        AES(bool encrypting);
+        AES(bool encrypting, size_t keySizeBits = 128);
+        AES(AES const&) = delete;
+        AES(AES&&) = delete;
+        AES& operator=(AES const&) = delete;
+        AES& operator=(AES&&) = delete;
         ~AES();
 
         void Init(Key const& key);
+        void Init(std::span<uint8 const> key);
 
         bool Process(IV const& iv, uint8* data, size_t length, Tag& tag);
         bool ProcessNoIntegrityCheck(IV const& iv, uint8* data, size_t partialLength);