Core/Crypto: Allow different AES key lengths

(cherry picked from commit f8f7fb58cb10fe7c76ac67a8e425edcb652bbc66)

1 files changed, 25 insertions, 2 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()));