Reuse AES key schedule when IV varies; avoid full OpenSSL re-init per row in non-GCM modes

Author: filimonov

src/Functions/FunctionsAES.h

@@ -305,10 +305,12 @@ class FunctionEncrypt : public IFunction
         const bool iv_is_constant = iv_column && isColumnConst(*iv_column);
         /// Avoid re-initializing cipher with the same key/IV on every row (expensive with OpenSSL 3).
         const bool can_reuse_context = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant && (!iv_column || iv_is_constant);
+        /// Reuse only key schedule even if IV changes per row (still cheaper than full init).
+        const bool can_reuse_key_schedule = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant;
 
         StringRef const_key_value{};
         StringRef const_iv_value{};
-        if (can_reuse_context)
+        if (can_reuse_key_schedule)
         {
             const_key_value = key_holder.setKey(key_size, key_column->getDataAt(0));
             if constexpr (mode != CipherMode::MySQLCompatibility)
@@ -317,7 +319,7 @@ class FunctionEncrypt : public IFunction
                     throw Exception(ErrorCodes::BAD_ARGUMENTS, "Invalid key size {} != expected size {}", const_key_value.size, key_size);
             }
 
-            if (iv_column)
+            if (iv_column && iv_is_constant)
             {
                 const_iv_value = iv_column->getDataAt(0);
                 if (const_iv_value.size == 0)
@@ -347,6 +349,22 @@ class FunctionEncrypt : public IFunction
                         reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
                     onError("EVP_EncryptInit_ex");
             }
+            else if (can_reuse_key_schedule)
+            {
+                key_value = const_key_value;
+                if (iv_column)
+                {
+                    iv_value = iv_column->getDataAt(row_idx);
+                    if (iv_value.size == 0)
+                        iv_value.data = nullptr;
+                }
+
+                validateIV<mode>(iv_value, iv_size);
+
+                if (EVP_EncryptInit_ex(evp_ctx, nullptr, nullptr, nullptr,
+                        reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
+                    onError("EVP_EncryptInit_ex");
+            }
             else
             {
                 key_value = key_holder.setKey(key_size, key_column->getDataAt(row_idx));
@@ -379,7 +397,7 @@ class FunctionEncrypt : public IFunction
             // Avoid extra work on empty ciphertext/plaintext for some ciphers
             if (!(input_value.size == 0 && block_size == 1 && mode != CipherMode::RFC5116_AEAD_AES_GCM))
             {
-                if (!can_reuse_context)
+                if (!can_reuse_context && !can_reuse_key_schedule)
                 {
                     // 1: Init CTX
                     if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
@@ -636,10 +654,12 @@ class FunctionDecrypt : public IFunction
         const bool iv_is_constant = iv_column && isColumnConst(*iv_column);
         /// Avoid re-initializing cipher with the same key/IV on every row (expensive with OpenSSL 3).
         const bool can_reuse_context = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant && (!iv_column || iv_is_constant);
+        /// Reuse only key schedule even if IV changes per row.
+        const bool can_reuse_key_schedule = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant;
 
         StringRef const_key_value{};
         StringRef const_iv_value{};
-        if (can_reuse_context)
+        if (can_reuse_key_schedule)
         {
             const_key_value = key_holder.setKey(key_size, key_column->getDataAt(0));
             if constexpr (mode != CipherMode::MySQLCompatibility)
@@ -648,7 +668,7 @@ class FunctionDecrypt : public IFunction
                     throw Exception(ErrorCodes::BAD_ARGUMENTS, "Invalid key size {} != expected size {}", const_key_value.size, key_size);
             }
 
-            if (iv_column)
+            if (iv_column && iv_is_constant)
             {
                 const_iv_value = iv_column->getDataAt(0);
 
@@ -681,6 +701,24 @@ class FunctionDecrypt : public IFunction
                         reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
                     onError("EVP_DecryptInit_ex");
             }
+            else if (can_reuse_key_schedule)
+            {
+                key_value = const_key_value;
+                if (iv_column)
+                {
+                    iv_value = iv_column->getDataAt(row_idx);
+
+                    /// If the length is zero (empty string is passed) it should be treat as no IV.
+                    if (iv_value.size == 0)
+                        iv_value.data = nullptr;
+                }
+
+                validateIV<mode>(iv_value, iv_size);
+
+                if (EVP_DecryptInit_ex(evp_ctx, nullptr, nullptr, nullptr,
+                        reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
+                    onError("EVP_DecryptInit_ex");
+            }
             else
             {
                 key_value = key_holder.setKey(key_size, key_column->getDataAt(row_idx));
@@ -729,7 +767,7 @@ class FunctionDecrypt : public IFunction
             /// This makes sense for default implementation for NULLs.
             if (input_value.size > 0)
             {
-                if (!can_reuse_context)
+                if (!can_reuse_context && !can_reuse_key_schedule)
                 {
                     // 1: Init CTX
                     if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)