Skip to content

Commit 0e4f172

Browse files
filimonovfilimonov
committed
Optimize AES GCM key reuse: keep key schedule for constant key, reinit only IV/IVLEN per row to cut OpenSSL 3 re-init overhead; fix redundant init and ensure AAD handled in reuse path
1 parent 1025c95 commit 0e4f172

File tree

1 file changed

+111
-41
lines changed

1 file changed

+111
-41
lines changed

‎src/Functions/FunctionsAES.h‎

Lines changed: 111 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -307,10 +307,12 @@ class FunctionEncrypt : public IFunction
307307
const bool can_reuse_context = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant && (!iv_column || iv_is_constant);
308308
/// Reuse only key schedule even if IV changes per row (still cheaper than full init).
309309
const bool can_reuse_key_schedule = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant;
310+
/// GCM: reuse key schedule, still re-set IVLEN/IV per row.
311+
const bool can_reuse_gcm_key = (mode == CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant;
310312

311313
StringRef const_key_value{};
312314
StringRef const_iv_value{};
313-
if (can_reuse_key_schedule)
315+
if (can_reuse_key_schedule || can_reuse_gcm_key)
314316
{
315317
const_key_value = key_holder.setKey(key_size, key_column->getDataAt(0));
316318
if constexpr (mode != CipherMode::MySQLCompatibility)
@@ -328,10 +330,20 @@ class FunctionEncrypt : public IFunction
328330
validateIV<mode>(const_iv_value, iv_size);
329331
}
330332

331-
if (EVP_EncryptInit_ex(evp_ctx, evp_cipher, nullptr,
332-
reinterpret_cast<const unsigned char*>(const_key_value.data),
333-
reinterpret_cast<const unsigned char*>(const_iv_value.data)) != 1)
334-
onError("EVP_EncryptInit_ex");
333+
if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
334+
{
335+
if (EVP_EncryptInit_ex(evp_ctx, evp_cipher, nullptr,
336+
reinterpret_cast<const unsigned char*>(const_key_value.data),
337+
nullptr) != 1)
338+
onError("EVP_EncryptInit_ex");
339+
}
340+
else
341+
{
342+
if (EVP_EncryptInit_ex(evp_ctx, evp_cipher, nullptr,
343+
reinterpret_cast<const unsigned char*>(const_key_value.data),
344+
reinterpret_cast<const unsigned char*>(const_iv_value.data)) != 1)
345+
onError("EVP_EncryptInit_ex");
346+
}
335347
}
336348

337349
for (size_t row_idx = 0; row_idx < input_rows_count; ++row_idx)
@@ -349,6 +361,27 @@ class FunctionEncrypt : public IFunction
349361
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
350362
onError("EVP_EncryptInit_ex");
351363
}
364+
else if (can_reuse_gcm_key)
365+
{
366+
key_value = const_key_value;
367+
if (iv_column)
368+
{
369+
iv_value = iv_column->getDataAt(row_idx);
370+
if (iv_value.size == 0)
371+
iv_value.data = nullptr;
372+
}
373+
374+
if (iv_value.size == 0)
375+
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Invalid IV size {} != expected size {}", iv_value.size, iv_size);
376+
377+
if (EVP_CIPHER_CTX_ctrl(evp_ctx, EVP_CTRL_AEAD_SET_IVLEN, safe_cast<int>(iv_value.size), nullptr) != 1)
378+
onError("EVP_CIPHER_CTX_ctrl");
379+
380+
if (EVP_EncryptInit_ex(evp_ctx, nullptr, nullptr,
381+
nullptr,
382+
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
383+
onError("EVP_EncryptInit_ex");
384+
}
352385
else if (can_reuse_key_schedule)
353386
{
354387
key_value = const_key_value;
@@ -397,32 +430,22 @@ class FunctionEncrypt : public IFunction
397430
// Avoid extra work on empty ciphertext/plaintext for some ciphers
398431
if (!(input_value.size == 0 && block_size == 1 && mode != CipherMode::RFC5116_AEAD_AES_GCM))
399432
{
400-
if (!can_reuse_context && !can_reuse_key_schedule)
433+
if (!can_reuse_context && !can_reuse_key_schedule && !can_reuse_gcm_key)
401434
{
402435
// 1: Init CTX
403436
if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
404437
{
405-
// 1.a.1: Init CTX with custom IV length and optionally with AAD
406-
if (EVP_EncryptInit_ex(evp_ctx, evp_cipher, nullptr, nullptr, nullptr) != 1)
407-
onError("EVP_EncryptInit_ex");
408-
409-
if (EVP_CIPHER_CTX_ctrl(evp_ctx, EVP_CTRL_AEAD_SET_IVLEN, safe_cast<int>(iv_value.size), nullptr) != 1)
410-
onError("EVP_CIPHER_CTX_ctrl");
438+
// 1.a.1: Init CTX with custom IV length
439+
if (EVP_EncryptInit_ex(evp_ctx, evp_cipher, nullptr, nullptr, nullptr) != 1)
440+
onError("EVP_EncryptInit_ex");
411441

412-
if (EVP_EncryptInit_ex(evp_ctx, nullptr, nullptr,
413-
reinterpret_cast<const unsigned char*>(key_value.data),
414-
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
415-
onError("EVP_EncryptInit_ex");
442+
if (EVP_CIPHER_CTX_ctrl(evp_ctx, EVP_CTRL_AEAD_SET_IVLEN, safe_cast<int>(iv_value.size), nullptr) != 1)
443+
onError("EVP_CIPHER_CTX_ctrl");
416444

417-
// 1.a.2 Set AAD
418-
if (aad_column)
419-
{
420-
const auto aad_data = aad_column->getDataAt(row_idx);
421-
int tmp_len = 0;
422-
if (aad_data.size != 0 && EVP_EncryptUpdate(evp_ctx, nullptr, &tmp_len,
423-
reinterpret_cast<const unsigned char *>(aad_data.data), safe_cast<int>(aad_data.size)) != 1)
424-
onError("EVP_EncryptUpdate");
425-
}
445+
if (EVP_EncryptInit_ex(evp_ctx, nullptr, nullptr,
446+
reinterpret_cast<const unsigned char*>(key_value.data),
447+
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
448+
onError("EVP_EncryptInit_ex");
426449
}
427450
else
428451
{
@@ -436,6 +459,18 @@ class FunctionEncrypt : public IFunction
436459
}
437460
}
438461

462+
if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
463+
{
464+
if (aad_column)
465+
{
466+
const auto aad_data = aad_column->getDataAt(row_idx);
467+
int tmp_len = 0;
468+
if (aad_data.size != 0 && EVP_EncryptUpdate(evp_ctx, nullptr, &tmp_len,
469+
reinterpret_cast<const unsigned char *>(aad_data.data), safe_cast<int>(aad_data.size)) != 1)
470+
onError("EVP_EncryptUpdate");
471+
}
472+
}
473+
439474
int output_len = 0;
440475
// 2: Feed the data to the cipher
441476
if (EVP_EncryptUpdate(evp_ctx,
@@ -656,10 +691,11 @@ class FunctionDecrypt : public IFunction
656691
const bool can_reuse_context = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant && (!iv_column || iv_is_constant);
657692
/// Reuse only key schedule even if IV changes per row.
658693
const bool can_reuse_key_schedule = (mode != CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant;
694+
const bool can_reuse_gcm_key = (mode == CipherMode::RFC5116_AEAD_AES_GCM) && key_is_constant;
659695

660696
StringRef const_key_value{};
661697
StringRef const_iv_value{};
662-
if (can_reuse_key_schedule)
698+
if (can_reuse_key_schedule || can_reuse_gcm_key)
663699
{
664700
const_key_value = key_holder.setKey(key_size, key_column->getDataAt(0));
665701
if constexpr (mode != CipherMode::MySQLCompatibility)
@@ -679,10 +715,20 @@ class FunctionDecrypt : public IFunction
679715
validateIV<mode>(const_iv_value, iv_size);
680716
}
681717

682-
if (EVP_DecryptInit_ex(evp_ctx, evp_cipher, nullptr,
683-
reinterpret_cast<const unsigned char*>(const_key_value.data),
684-
reinterpret_cast<const unsigned char*>(const_iv_value.data)) != 1)
685-
onError("EVP_DecryptInit_ex");
718+
if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
719+
{
720+
if (EVP_DecryptInit_ex(evp_ctx, evp_cipher, nullptr,
721+
reinterpret_cast<const unsigned char*>(const_key_value.data),
722+
nullptr) != 1)
723+
onError("EVP_DecryptInit_ex");
724+
}
725+
else
726+
{
727+
if (EVP_DecryptInit_ex(evp_ctx, evp_cipher, nullptr,
728+
reinterpret_cast<const unsigned char*>(const_key_value.data),
729+
reinterpret_cast<const unsigned char*>(const_iv_value.data)) != 1)
730+
onError("EVP_DecryptInit_ex");
731+
}
686732
}
687733

688734
for (size_t row_idx = 0; row_idx < input_rows_count; ++row_idx)
@@ -701,6 +747,28 @@ class FunctionDecrypt : public IFunction
701747
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
702748
onError("EVP_DecryptInit_ex");
703749
}
750+
else if (can_reuse_gcm_key)
751+
{
752+
key_value = const_key_value;
753+
if (iv_column)
754+
{
755+
iv_value = iv_column->getDataAt(row_idx);
756+
757+
/// If the length is zero (empty string is passed) it should be treat as no IV.
758+
if (iv_value.size == 0)
759+
iv_value.data = nullptr;
760+
}
761+
762+
if (iv_value.size == 0)
763+
throw Exception(ErrorCodes::BAD_ARGUMENTS, "Invalid IV size {} != expected size {}", iv_value.size, iv_size);
764+
765+
if (EVP_CIPHER_CTX_ctrl(evp_ctx, EVP_CTRL_AEAD_SET_IVLEN, safe_cast<int>(iv_value.size), nullptr) != 1)
766+
onError("EVP_CIPHER_CTX_ctrl");
767+
768+
if (EVP_DecryptInit_ex(evp_ctx, nullptr, nullptr, nullptr,
769+
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
770+
onError("EVP_DecryptInit_ex");
771+
}
704772
else if (can_reuse_key_schedule)
705773
{
706774
key_value = const_key_value;
@@ -767,7 +835,7 @@ class FunctionDecrypt : public IFunction
767835
/// This makes sense for default implementation for NULLs.
768836
if (input_value.size > 0)
769837
{
770-
if (!can_reuse_context && !can_reuse_key_schedule)
838+
if (!can_reuse_context && !can_reuse_key_schedule && !can_reuse_gcm_key)
771839
{
772840
// 1: Init CTX
773841
if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
@@ -784,16 +852,6 @@ class FunctionDecrypt : public IFunction
784852
reinterpret_cast<const unsigned char*>(key_value.data),
785853
reinterpret_cast<const unsigned char*>(iv_value.data)) != 1)
786854
onError("EVP_DecryptInit_ex");
787-
788-
// 1.a.2: Set AAD if present
789-
if (aad_column)
790-
{
791-
StringRef aad_data = aad_column->getDataAt(row_idx);
792-
int tmp_len = 0;
793-
if (aad_data.size != 0 && EVP_DecryptUpdate(evp_ctx, nullptr, &tmp_len,
794-
reinterpret_cast<const unsigned char *>(aad_data.data), safe_cast<int>(aad_data.size)) != 1)
795-
onError("EVP_DecryptUpdate");
796-
}
797855
}
798856
else
799857
{
@@ -808,6 +866,18 @@ class FunctionDecrypt : public IFunction
808866
}
809867

810868
// 2: Feed the data to the cipher
869+
if constexpr (mode == CipherMode::RFC5116_AEAD_AES_GCM)
870+
{
871+
if (aad_column)
872+
{
873+
StringRef aad_data = aad_column->getDataAt(row_idx);
874+
int tmp_len = 0;
875+
if (aad_data.size != 0 && EVP_DecryptUpdate(evp_ctx, nullptr, &tmp_len,
876+
reinterpret_cast<const unsigned char *>(aad_data.data), safe_cast<int>(aad_data.size)) != 1)
877+
onError("EVP_DecryptUpdate");
878+
}
879+
}
880+
811881
int output_len = 0;
812882
if (EVP_DecryptUpdate(evp_ctx,
813883
reinterpret_cast<unsigned char*>(decrypted), &output_len,

0 commit comments

Comments
 (0)