From ddee0e714576013726752f725be7ddf56cfb7a53 Mon Sep 17 00:00:00 2001
From: cambyzhu <cambyzhu@tencent.com>
Date: Thu, 12 Sep 2024 14:44:13 +0800
Subject: [PATCH] fix wrong mode arg of encrypt and decrypt function make BE
 crash

---
 be/src/vec/functions/function_encryption.cpp             | 9 ++++++---
 .../encryption_digest/test_encryption_function.out       | 3 +++
 .../encryption_digest/test_encryption_function.groovy    | 3 +++
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/be/src/vec/functions/function_encryption.cpp b/be/src/vec/functions/function_encryption.cpp
index c90b6a1ff60717..d017f761897133 100644
--- a/be/src/vec/functions/function_encryption.cpp
+++ b/be/src/vec/functions/function_encryption.cpp
@@ -244,8 +244,9 @@ struct EncryptionAndDecryptTwoImpl {
         if (mode_arg.size != 0) {
             if (!aes_mode_map.contains(mode_str)) {
                 all_insert_null = true;
+            } else {
+                encryption_mode = aes_mode_map.at(mode_str);
             }
-            encryption_mode = aes_mode_map.at(mode_str);
         }
         const ColumnString::Offsets* offsets_column = &column->get_offsets();
         const ColumnString::Chars* chars_column = &column->get_chars();
@@ -371,13 +372,15 @@ struct EncryptionAndDecryptMultiImpl {
             if constexpr (is_sm_mode) {
                 if (sm4_mode_map.count(mode_str) == 0) {
                     all_insert_null = true;
+                } else {
+                    encryption_mode = sm4_mode_map.at(mode_str);
                 }
-                encryption_mode = sm4_mode_map.at(mode_str);
             } else {
                 if (aes_mode_map.count(mode_str) == 0) {
                     all_insert_null = true;
+                } else {
+                    encryption_mode = aes_mode_map.at(mode_str);
                 }
-                encryption_mode = aes_mode_map.at(mode_str);
             }
         }
 
diff --git a/regression-test/data/query_p0/sql_functions/encryption_digest/test_encryption_function.out b/regression-test/data/query_p0/sql_functions/encryption_digest/test_encryption_function.out
index c652b3074558cc..3b3f06d98d8ccc 100644
--- a/regression-test/data/query_p0/sql_functions/encryption_digest/test_encryption_function.out
+++ b/regression-test/data/query_p0/sql_functions/encryption_digest/test_encryption_function.out
@@ -272,3 +272,6 @@ zhang
 -- !sql56 --
 zhang
 
+-- !sql57 --
+\N
+
diff --git a/regression-test/suites/query_p0/sql_functions/encryption_digest/test_encryption_function.groovy b/regression-test/suites/query_p0/sql_functions/encryption_digest/test_encryption_function.groovy
index 25095f46917c96..69933b6fd87db9 100644
--- a/regression-test/suites/query_p0/sql_functions/encryption_digest/test_encryption_function.groovy
+++ b/regression-test/suites/query_p0/sql_functions/encryption_digest/test_encryption_function.groovy
@@ -226,4 +226,7 @@ suite("test_encryption_function") {
     qt_sql54 """ select aes_decrypt(aes_encrypt(k,k1,k2, "AES_256_CFB"),k1,k2, "AES_256_CFB") from quantile_table2; """
     qt_sql55 """ select aes_decrypt(aes_encrypt("zhang",k1,k2, "AES_256_CFB"),k1,k2, "AES_256_CFB") from quantile_table2; """
     qt_sql56 """ select aes_decrypt(aes_encrypt("zhang",k1,k2, "AES_256_CFB"),k1,k2, "AES_256_CFB") from quantile_table2; """
+
+    //four arg (column/const) with wrong mode
+    qt_sql57 """ select sm4_decrypt(sm4_encrypt(k,"doris","abcdefghij", "SM4_128_CBC"),"doris","abcdefghij","SM4_555_CBC") from quantile_table2; """
 }