[feature][core] 优化Encrypt和HMAC函数。新增Encrypt单元测试feature/udf-encrypt

7 files changed, 309 insertions, 57 deletions

diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/Encrypt.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/Encrypt.java
index cc05397..b20ff18 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/Encrypt.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/Encrypt.java
@@ -1,6 +1,9 @@
 package com.geedgenetworks.core.udf;
 
-import cn.hutool.core.util.ArrayUtil;
+import cn.hutool.core.util.URLUtil;
+import cn.hutool.json.JSONArray;
+import cn.hutool.json.JSONObject;
+import cn.hutool.json.JSONUtil;
 import com.alibaba.fastjson2.JSON;
 import com.geedgenetworks.common.Constants;
 import com.geedgenetworks.common.Event;
@@ -15,13 +18,17 @@ import com.geedgenetworks.core.pojo.KmsKey;
 import com.geedgenetworks.core.udf.encrypt.EncryptionAlgorithm;
 import com.geedgenetworks.core.utils.*;
 import com.geedgenetworks.utils.StringUtil;
-import io.github.jopenlibs.vault.VaultException;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.apache.flink.configuration.Configuration;
 
+import java.io.IOException;
+import java.net.URI;
 import java.util.Arrays;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 @Slf4j
 public class Encrypt implements ScalarFunction {
@@ -31,7 +38,7 @@ public class Encrypt implements ScalarFunction {
     private String identifier;
     private String defaultVal;
     private String type;
-    private transient SingleValueMap.Data<LoadIntervalDataUtil<String[]>> sensitiveFieldsData;
+    private transient SingleValueMap.Data<LoadIntervalDataUtil<Set<String>>> sensitiveFieldsData;
     private transient SingleValueMap.Data<LoadIntervalDataUtil<KmsKey>> kmsKeyData;
     private transient EncryptionAlgorithm encryptionAlgorithm;
 
@@ -44,35 +51,37 @@ public class Encrypt implements ScalarFunction {
         this.lookupFieldName = udfContext.getLookup_fields().get(0);
         this.outputFieldName = udfContext.getOutput_fields().get(0);
         this.identifier = udfContext.getParameters().get("identifier").toString();
-
         Configuration configuration = (Configuration) runtimeContext.getExecutionConfig().getGlobalJobParameters();
         CommonConfig commonConfig = JSON.parseObject(configuration.toMap().get(Constants.SYSPROP_GROOTSTREAM_CONFIG), CommonConfig.class);
-        Map<String, KmsConfig> kmsConfigs = commonConfig.getKmsConfig();
-        if (kmsConfigs.isEmpty()) {
-            throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Global parameter kms type is not null!");
-        } else if (kmsConfigs.size() > 1) {
-            throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Global parameter kms type is repeated!");
-        }
-        KmsConfig kmsConfig = kmsConfigs.values().iterator().next();
+        KmsConfig kmsConfig = commonConfig.getKmsConfig().get(configuration.toMap().get(Constants.SYSPROP_KMS_TYPE_CONFIG));
         SSLConfig sslConfig = commonConfig.getSslConfig();
         Map<String, String> propertiesConfig = commonConfig.getPropertiesConfig();
         type = kmsConfig.getType();
         try {
-            encryptionAlgorithm = EncryptionAlgorithmUtils.getEncryptionAlgorithm(identifier);
+            encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm(identifier);
             if (encryptionAlgorithm == null) {
                 throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Parameters identifier is illegal!");
             }
-            kmsKeyData = SingleValueMap.acquireData("kmsKeyData",
-                    () -> LoadIntervalDataUtil.newInstance(() -> getKmsKey(kmsConfig, sslConfig, identifier),
-                            LoadIntervalDataOptions.defaults("kmsKeyData", 60000)), LoadIntervalDataUtil::stop);
-            sensitiveFieldsData = SingleValueMap.acquireData("sensitiveFields",
-                    () -> LoadIntervalDataUtil.newInstance(() -> getEncryptFields(propertiesConfig.get("projection.encrypt.schema.registry.uri")),
-                            LoadIntervalDataOptions.defaults("sensitiveFields", 60000)), LoadIntervalDataUtil::stop);
-            KmsKey kmsKey = kmsKeyData.getData().data();
-            if (encryptionAlgorithm.getSecretKeyLength() == kmsKey.getKeyData().length) {
+            if (!type.equals(KmsUtils.KMS_TYPE_LOCAL)) {
+                kmsKeyData = SingleValueMap.acquireData("kmsKeyData",
+                        () -> LoadIntervalDataUtil.newInstance(() -> KmsUtils.getVaultKey(kmsConfig, sslConfig, identifier),
+                                LoadIntervalDataOptions.defaults("kmsKeyData", Integer.parseInt(propertiesConfig.getOrDefault(Constants.SYSPROP_ENCRYPT_KMS_KEY_SCHEDULER_INTERVAL_NAME, "5")) * 60000L)),
+                        LoadIntervalDataUtil::stop);
+                KmsKey kmsKey = kmsKeyData.getData().data();
+                if (kmsKey == null) {
+                    throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Initialization UDF Encrypt failed!");
+                }
+                if (encryptionAlgorithm.getSecretKeyLength() != kmsKey.getKeyData().length) {
+                    throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Global parameter kms secret Key requires " + encryptionAlgorithm.getSecretKeyLength() + " bytes!");
+                }
                 encryptionAlgorithm.setKmsKey(kmsKey);
-            } else {
-                throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Global parameter kms secret Key requires " + encryptionAlgorithm.getSecretKeyLength() + " bytes!");
+            }
+            sensitiveFieldsData = SingleValueMap.acquireData("sensitiveFields",
+                    () -> LoadIntervalDataUtil.newInstance(() -> getSensitiveFields(propertiesConfig.get("projection.encrypt.schema.registry.uri")),
+                            LoadIntervalDataOptions.defaults("sensitiveFields", Integer.parseInt(propertiesConfig.getOrDefault(Constants.SYSPROP_ENCRYPT_SENSITIVE_FIELDS_SCHEDULER_INTERVAL_NAME, "5")) * 60000L)),
+                    LoadIntervalDataUtil::stop);
+            if (sensitiveFieldsData.getData().data() == null) {
+                throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, "Initialization UDF Encrypt failed!");
             }
         } catch (Exception e) {
             throw new GrootStreamRuntimeException(CommonErrorCode.UNSUPPORTED_OPERATION, "Initialization UDF Encrypt failed!", e);
@@ -82,11 +91,13 @@ public class Encrypt implements ScalarFunction {
     @Override
     public Event evaluate(Event event) {
         try {
-            KmsKey kmsKey = kmsKeyData.getData().data();
-            if (kmsKey.getKeyVersion() != encryptionAlgorithm.getKmsKey().getKeyVersion() || !Arrays.equals(kmsKey.getKeyData(), encryptionAlgorithm.getKmsKey().getKeyData())) {
-                encryptionAlgorithm.setKmsKey(kmsKey);
+            if (!type.equals(KmsUtils.KMS_TYPE_LOCAL)) {
+                KmsKey kmsKey = kmsKeyData.getData().data();
+                if (kmsKey.getKeyVersion() != encryptionAlgorithm.getKmsKey().getKeyVersion() || !Arrays.equals(kmsKey.getKeyData(), encryptionAlgorithm.getKmsKey().getKeyData())) {
+                    encryptionAlgorithm.setKmsKey(kmsKey);
+                }
             }
-            if (ArrayUtil.contains(sensitiveFieldsData.getData().data(), lookupFieldName) && event.getExtractedFields().containsKey(lookupFieldName)) {
+            if (sensitiveFieldsData.getData().data().contains(lookupFieldName) && event.getExtractedFields().containsKey(lookupFieldName)) {
                 String value = (String) event.getExtractedFields().get(lookupFieldName);
                 if (StringUtil.isNotBlank(value)) {
                     String encryptResult = encryptionAlgorithm.encrypt(value);
@@ -136,24 +147,18 @@ public class Encrypt implements ScalarFunction {
         }
     }
 
-    private KmsKey getKmsKey(KmsConfig kmsConfig, SSLConfig sslConfig, String identifier) throws VaultException {
-        KmsKey kmsKey = null;
-        if (KmsUtils.KMS_TYPE_VAULT.equals(kmsConfig.getType())) {
-            kmsKey = KmsUtils.getVaultKey(kmsConfig, sslConfig, identifier);
-        } else if (KmsUtils.KMS_TYPE_LOCAL.equals(kmsConfig.getType())) {
-            kmsKey = new KmsKey(kmsConfig.getSecretKey().getBytes(), 1);
+    public Set<String> getSensitiveFields(String url) throws IOException {
+        Set<String> sensitiveFieldsSet;
+        String sensitiveFieldsStr = HttpClientPoolUtil.getInstance().httpGet(URI.create(URLUtil.normalize(url)));
+        JSONObject sensitiveFieldsJson = JSONUtil.parseObj(sensitiveFieldsStr);
+        if (sensitiveFieldsJson.getInt("status", 500) == 200) {
+            JSONArray sensitiveFieldsJsonArr = sensitiveFieldsJson.getJSONArray("data");
+            sensitiveFieldsSet = IntStream.range(0, sensitiveFieldsJsonArr.size())
+                    .mapToObj(sensitiveFieldsJsonArr::getStr)
+                    .collect(Collectors.toSet());
+        } else {
+            throw new GrootStreamRuntimeException(CommonErrorCode.UNSUPPORTED_OPERATION, "Get encrypt fields error! Error message: " + sensitiveFieldsStr);
         }
-        return kmsKey;
-    }
-
-    private String[] getEncryptFields(String url) {
-        String[] encryptFields = new String[]{"phone_number", "server_ip"};
-//        try {
-//            String s = HttpClientPoolUtil.getInstance().httpGet(URI.create(URLUtil.normalize(url)));
-//            encryptFields = s.split(",");
-//        } catch (Exception e) {
-//            log.error("Get encrypt fields error! " + e.getMessage());
-//        }
-        return encryptFields;
+        return sensitiveFieldsSet;
     }
 }
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES128GCM96Algorithm.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES128GCM96Algorithm.java
index 74be5a8..db4369e 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES128GCM96Algorithm.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES128GCM96Algorithm.java
@@ -15,12 +15,14 @@ public class AES128GCM96Algorithm implements EncryptionAlgorithm {
     private static final int GCM_TAG_LENGTH = 128;
     private static final int GCM_96_NONCE_LENGTH = 12;
     private static final int SECRET_KEY_LENGTH = 16;
+    private static final byte[] DEFAULT_SECRET_KEY = ".geedgenetworks.".getBytes();
 
     private final Cipher cipher;
     private KmsKey kmsKey;
 
     public AES128GCM96Algorithm() throws Exception {
         this.cipher = Cipher.getInstance(TRANSFORMATION);
+        this.kmsKey = new KmsKey(DEFAULT_SECRET_KEY, 1);
     }
 
     @Override
@@ -66,13 +68,12 @@ public class AES128GCM96Algorithm implements EncryptionAlgorithm {
         String decryptedString = "";
         try {
             byte[] nonce = RandomUtil.randomBytes(GCM_96_NONCE_LENGTH);
-            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
-            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
-            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kmsKey.getKeyData(), ALGORITHM), gcmSpec);
             byte[] combined = Base64.getDecoder().decode(content);
             byte[] encryptedBytes = new byte[combined.length - GCM_96_NONCE_LENGTH];
             System.arraycopy(combined, 0, nonce, 0, GCM_96_NONCE_LENGTH);
             System.arraycopy(combined, GCM_96_NONCE_LENGTH, encryptedBytes, 0, encryptedBytes.length);
+            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
+            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kmsKey.getKeyData(), ALGORITHM), gcmSpec);
             byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
             decryptedString = new String(decryptedBytes);
         } catch (Exception e) {
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES256GCM96Algorithm.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES256GCM96Algorithm.java
index 64d88d9..dec7e01 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES256GCM96Algorithm.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/AES256GCM96Algorithm.java
@@ -15,12 +15,14 @@ public class AES256GCM96Algorithm implements EncryptionAlgorithm {
     private static final int GCM_TAG_LENGTH = 128;
     private static final int GCM_96_NONCE_LENGTH = 12;
     private static final int SECRET_KEY_LENGTH = 32;
+    private static final byte[] DEFAULT_SECRET_KEY = ".........geedgenetworks.........".getBytes();
 
     private final Cipher cipher;
     private KmsKey kmsKey;
 
     public AES256GCM96Algorithm() throws Exception {
         this.cipher = Cipher.getInstance(TRANSFORMATION);
+        this.kmsKey = new KmsKey(DEFAULT_SECRET_KEY, 1);
     }
 
     @Override
@@ -66,13 +68,12 @@ public class AES256GCM96Algorithm implements EncryptionAlgorithm {
         String decryptedString = "";
         try {
             byte[] nonce = RandomUtil.randomBytes(GCM_96_NONCE_LENGTH);
-            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
-            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
-            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kmsKey.getKeyData(), ALGORITHM), gcmSpec);
             byte[] combined = Base64.getDecoder().decode(content);
             byte[] encryptedBytes = new byte[combined.length - GCM_96_NONCE_LENGTH];
             System.arraycopy(combined, 0, nonce, 0, GCM_96_NONCE_LENGTH);
             System.arraycopy(combined, GCM_96_NONCE_LENGTH, encryptedBytes, 0, encryptedBytes.length);
+            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
+            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kmsKey.getKeyData(), ALGORITHM), gcmSpec);
             byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
             decryptedString = new String(decryptedBytes);
         } catch (Exception e) {
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/SM4GCM96Algorithm.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/SM4GCM96Algorithm.java
index 3c13820..e13cb40 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/SM4GCM96Algorithm.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/encrypt/SM4GCM96Algorithm.java
@@ -15,12 +15,14 @@ public class SM4GCM96Algorithm implements EncryptionAlgorithm {
     private static final int GCM_TAG_LENGTH = 128;
     private static final int GCM_96_NONCE_LENGTH = 12;
     private static final int SECRET_KEY_LENGTH = 16;
+    private static final byte[] DEFAULT_SECRET_KEY = ".geedgenetworks.".getBytes();
 
     private final Cipher cipher;
     private KmsKey kmsKey;
 
     public SM4GCM96Algorithm() throws Exception {
         this.cipher = Cipher.getInstance(TRANSFORMATION);
+        this.kmsKey = new KmsKey(DEFAULT_SECRET_KEY, 1);
     }
 
     @Override
@@ -66,13 +68,12 @@ public class SM4GCM96Algorithm implements EncryptionAlgorithm {
         String decryptedString = "";
         try {
             byte[] nonce = RandomUtil.randomBytes(GCM_96_NONCE_LENGTH);
-            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
-            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
-            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kmsKey.getKeyData(), ALGORITHM), gcmSpec);
             byte[] combined = Base64.getDecoder().decode(content);
             byte[] encryptedBytes = new byte[combined.length - GCM_96_NONCE_LENGTH];
             System.arraycopy(combined, 0, nonce, 0, GCM_96_NONCE_LENGTH);
             System.arraycopy(combined, GCM_96_NONCE_LENGTH, encryptedBytes, 0, encryptedBytes.length);
+            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
+            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(kmsKey.getKeyData(), ALGORITHM), gcmSpec);
             byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
             decryptedString = new String(decryptedBytes);
         } catch (Exception e) {
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/EncryptionAlgorithmUtils.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/EncryptionAlgorithmUtils.java
index 7041c73..0a0fe33 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/utils/EncryptionAlgorithmUtils.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/EncryptionAlgorithmUtils.java
@@ -15,7 +15,7 @@ public final class EncryptionAlgorithmUtils {
     public static final String ALGORITHM_AES_256_GCM96_NAME = "aes-256-gcm96";
     public static final String ALGORITHM_SM4_GCM96_NAME = "sm4-gcm96";
 
-    public static EncryptionAlgorithm getEncryptionAlgorithm(String identifier) throws Exception {
+    public static EncryptionAlgorithm createEncryptionAlgorithm(String identifier) throws Exception {
         switch (identifier) {
             case ALGORITHM_AES_128_GCM96_NAME:
                 return new AES128GCM96Algorithm();
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KmsUtils.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KmsUtils.java
index 8e6a345..9519dd5 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/utils/KmsUtils.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KmsUtils.java
@@ -21,7 +21,7 @@ public class KmsUtils {
     public static final String KMS_TYPE_LOCAL = "local";
     public static final String KMS_TYPE_VAULT = "vault";
 
-    public static KmsKey getVaultKey(KmsConfig kmsConfig, SSLConfig sslConfig, String identifier) throws VaultException {
+    public static KmsKey getVaultKey(KmsConfig kmsConfig, SSLConfig sslConfig, String identifier) throws Exception {
         Vault vault = getVaultClient(kmsConfig, sslConfig);
         String exportKeyPath;
         if (EncryptionAlgorithmUtils.ALGORITHM_SM4_GCM96_NAME.equals(identifier)) {
@@ -34,8 +34,7 @@ public class KmsUtils {
             JsonObject keys = exportResponse.getDataObject().get("keys").asObject();
             return new KmsKey(Base64.getDecoder().decode(StrUtil.trim(keys.get(keys.size() + "").asString(), '"')), keys.size());
         } else {
-            log.error("Get kms key error! code: {} body: {}", exportResponse.getRestResponse().getStatus(), new String(exportResponse.getRestResponse().getBody()));
-            return null;
+            throw new RuntimeException("Get vault key error! code: " + exportResponse.getRestResponse().getStatus() + " body: " + new String(exportResponse.getRestResponse().getBody()));
         }
     }
 
diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/simple/EncryptFunctionTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/simple/EncryptFunctionTest.java
new file mode 100644
index 0000000..a83d853
--- /dev/null
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/simple/EncryptFunctionTest.java
@@ -0,0 +1,245 @@
+package com.geedgenetworks.core.udf.test.simple;
+
+import cn.hutool.core.util.RandomUtil;
+import com.alibaba.fastjson2.JSON;
+import com.geedgenetworks.common.Constants;
+import com.geedgenetworks.common.Event;
+import com.geedgenetworks.common.config.CommonConfig;
+import com.geedgenetworks.common.config.KmsConfig;
+import com.geedgenetworks.common.config.SSLConfig;
+import com.geedgenetworks.common.exception.GrootStreamRuntimeException;
+import com.geedgenetworks.common.udf.UDFContext;
+import com.geedgenetworks.core.pojo.KmsKey;
+import com.geedgenetworks.core.udf.Encrypt;
+import com.geedgenetworks.core.udf.encrypt.EncryptionAlgorithm;
+import com.geedgenetworks.core.utils.EncryptionAlgorithmUtils;
+import com.geedgenetworks.core.utils.HttpClientPoolUtil;
+import com.geedgenetworks.core.utils.KmsUtils;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.functions.RuntimeContext;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.metrics.groups.OperatorMetricGroup;
+import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.mockito.*;
+
+import java.io.IOException;
+import java.security.Security;
+import java.util.*;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+public class EncryptFunctionTest {
+    private static UDFContext udfContext;
+    private static MockedStatic<HttpClientPoolUtil> httpClientPoolUtilMockedStatic;
+    private static final String DATA = "13812345678";
+
+    @BeforeAll
+    public static void setUp() throws IOException {
+        Security.addProvider(new BouncyCastleProvider());
+        udfContext = new UDFContext();
+        udfContext.setLookup_fields(Collections.singletonList("phone_number"));
+        udfContext.setOutput_fields(Collections.singletonList("phone_number"));
+        httpClientPoolUtilMockedStatic = mockSensitiveFields();
+    }
+
+    @AfterAll
+    public static void after() {
+        httpClientPoolUtilMockedStatic.close();
+    }
+
+    @Test
+    public void testEncryptByVault() throws Exception {
+        String secretKey = RandomUtil.randomString(32);
+        MockedStatic<KmsUtils> kmsUtilsMockedStatic = Mockito.mockStatic(KmsUtils.class);
+        Mockito.when(KmsUtils.getVaultKey(ArgumentMatchers.any(), ArgumentMatchers.any(), ArgumentMatchers.any())).thenReturn(new KmsKey(secretKey.getBytes(), 1));
+        RuntimeContext runtimeContext = mockVaultRuntimeContext();
+        Map<String, Object> map = new HashMap<>();
+        map.put("identifier", EncryptionAlgorithmUtils.ALGORITHM_AES_256_GCM96_NAME);
+        udfContext.setParameters(map);
+        Encrypt encrypt = new Encrypt();
+        encrypt.open(runtimeContext, udfContext);
+        Event event = new Event();
+        Map<String, Object> extractedFields = new HashMap<>();
+        extractedFields.put("phone_number", DATA);
+        event.setExtractedFields(extractedFields);
+        Event result = encrypt.evaluate(event);
+        EncryptionAlgorithm encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm(EncryptionAlgorithmUtils.ALGORITHM_AES_256_GCM96_NAME);
+        assertNotNull(encryptionAlgorithm);
+        encryptionAlgorithm.setKmsKey(new KmsKey(secretKey.getBytes(), 1));
+        String encrypted = result.getExtractedFields().get("phone_number").toString();
+        assertTrue(encrypted.contains("vault:v1:"));
+        String decrypted = encryptionAlgorithm.decrypt(encrypted.split(":")[2]);
+        assertEquals(DATA, decrypted);
+        encrypt.close();
+        kmsUtilsMockedStatic.close();
+    }
+
+    @Test
+    public void testEncryptByLocal() throws Exception {
+        byte[] secretKey = ".........geedgenetworks.........".getBytes();
+        RuntimeContext runtimeContext = mockLocalRuntimeContext();
+        Map<String, Object> map = new HashMap<>();
+        map.put("identifier", EncryptionAlgorithmUtils.ALGORITHM_AES_256_GCM96_NAME);
+        udfContext.setParameters(map);
+        Encrypt encrypt = new Encrypt();
+        encrypt.open(runtimeContext, udfContext);
+        Event event = new Event();
+        Map<String, Object> extractedFields = new HashMap<>();
+        extractedFields.put("phone_number", DATA);
+        event.setExtractedFields(extractedFields);
+        Event result = encrypt.evaluate(event);
+        EncryptionAlgorithm encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm(EncryptionAlgorithmUtils.ALGORITHM_AES_256_GCM96_NAME);
+        assertNotNull(encryptionAlgorithm);
+        encryptionAlgorithm.setKmsKey(new KmsKey(secretKey, 1));
+        String decrypted = encryptionAlgorithm.decrypt((String) result.getExtractedFields().get("phone_number"));
+        assertEquals(DATA, decrypted);
+        encrypt.close();
+    }
+
+    @Test
+    public void testEncryptByIdentifier() {
+        Map<String, Object> map = new HashMap<>();
+        map.put("identifier", EncryptionAlgorithmUtils.ALGORITHM_AES_256_GCM96_NAME);
+        udfContext.setParameters(map);
+        Encrypt encrypt1 = new Encrypt();
+        assertDoesNotThrow(() -> encrypt1.open(mockLocalRuntimeContext(), udfContext));
+        encrypt1.close();
+
+        Encrypt encrypt2 = new Encrypt();
+        map.put("identifier", EncryptionAlgorithmUtils.ALGORITHM_AES_128_GCM96_NAME);
+        udfContext.setParameters(map);
+        assertDoesNotThrow(() -> encrypt2.open(mockLocalRuntimeContext(), udfContext));
+        encrypt2.close();
+
+        Encrypt encrypt3 = new Encrypt();
+        map.put("identifier", EncryptionAlgorithmUtils.ALGORITHM_SM4_GCM96_NAME);
+        udfContext.setParameters(map);
+        assertDoesNotThrow(() -> encrypt3.open(mockLocalRuntimeContext(), udfContext));
+        encrypt3.close();
+    }
+
+    @Test
+    public void testEncryptionAlgorithm() throws Exception {
+        EncryptionAlgorithm encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm(EncryptionAlgorithmUtils.ALGORITHM_AES_128_GCM96_NAME);
+        assertNotNull(encryptionAlgorithm);
+        encryptionAlgorithm.setKmsKey(new KmsKey("aaaaaaaaaaaaaaaa".getBytes(), 1));
+        String encryptData = encryptionAlgorithm.encrypt(DATA);
+        String decryptData = encryptionAlgorithm.decrypt(encryptData);
+        assertEquals(DATA, decryptData);
+
+        encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm(EncryptionAlgorithmUtils.ALGORITHM_AES_256_GCM96_NAME);
+        assertNotNull(encryptionAlgorithm);
+        encryptionAlgorithm.setKmsKey(new KmsKey("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".getBytes(), 1));
+        encryptData = encryptionAlgorithm.encrypt(DATA);
+        decryptData = encryptionAlgorithm.decrypt(encryptData);
+        assertEquals(DATA, decryptData);
+
+        encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm(EncryptionAlgorithmUtils.ALGORITHM_SM4_GCM96_NAME);
+        assertNotNull(encryptionAlgorithm);
+        encryptionAlgorithm.setKmsKey(new KmsKey("aaaaaaaaaaaaaaaa".getBytes(), 1));
+        encryptData = encryptionAlgorithm.encrypt(DATA);
+        decryptData = encryptionAlgorithm.decrypt(encryptData);
+        assertEquals(DATA, decryptData);
+
+        encryptionAlgorithm = EncryptionAlgorithmUtils.createEncryptionAlgorithm("sm4");
+        assertNull(encryptionAlgorithm);
+    }
+
+    @Test
+    public void testEncryptError() {
+        RuntimeContext runtimeContext = mockLocalRuntimeContext();
+        Encrypt encrypt = new Encrypt();
+        udfContext.setParameters(null);
+        assertThrows(GrootStreamRuntimeException.class, () -> encrypt.open(runtimeContext, udfContext));
+
+        Map<String, Object> map = new HashMap<>();
+        udfContext.setParameters(map);
+        assertThrows(GrootStreamRuntimeException.class, () -> encrypt.open(runtimeContext, udfContext));
+
+        map.put("identifier", "aes");
+        udfContext.setParameters(map);
+        assertThrows(GrootStreamRuntimeException.class, () -> encrypt.open(runtimeContext, udfContext));
+    }
+
+    static RuntimeContext mockLocalRuntimeContext() {
+        RuntimeContext runtimeContext = Mockito.mock(RuntimeContext.class);
+        ExecutionConfig executionConfig = Mockito.mock(ExecutionConfig.class);
+        Mockito.when(runtimeContext.getExecutionConfig()).thenReturn(executionConfig);
+        MetricGroup metricGroup = Mockito.mock(OperatorMetricGroup.class);
+        Mockito.when(runtimeContext.getMetricGroup()).thenReturn(metricGroup);
+        Mockito.when(metricGroup.addGroup(Mockito.anyString())).thenReturn(metricGroup);
+        Mockito.when(metricGroup.counter(Mockito.anyString())).thenReturn(new SimpleCounter());
+        Configuration configuration = new Configuration();
+        CommonConfig commonConfig = new CommonConfig();
+        Map<String, KmsConfig> kmsConfigs = new HashMap<>();
+        KmsConfig kmsConfig = new KmsConfig();
+        kmsConfig.setType(KmsUtils.KMS_TYPE_LOCAL);
+        kmsConfigs.put(KmsUtils.KMS_TYPE_LOCAL, kmsConfig);
+        kmsConfig = new KmsConfig();
+        kmsConfig.setType(KmsUtils.KMS_TYPE_VAULT);
+        kmsConfigs.put(KmsUtils.KMS_TYPE_VAULT, kmsConfig);
+        SSLConfig sslConfig = new SSLConfig();
+        sslConfig.setSkipVerification(true);
+        Map<String, String> propertiesConfig = new HashMap<>();
+        propertiesConfig.put("projection.encrypt.schema.registry.uri", "127.0.0.1:9999/v1/schema/session_record?option=encrypt_fields");
+        commonConfig.setKmsConfig(kmsConfigs);
+        commonConfig.setSslConfig(sslConfig);
+        commonConfig.setPropertiesConfig(propertiesConfig);
+        configuration.setString(Constants.SYSPROP_GROOTSTREAM_CONFIG, JSON.toJSONString(commonConfig));
+        configuration.setString(Constants.SYSPROP_KMS_TYPE_CONFIG, KmsUtils.KMS_TYPE_LOCAL);
+        Mockito.when(executionConfig.getGlobalJobParameters()).thenReturn(configuration);
+        return runtimeContext;
+    }
+
+    static RuntimeContext mockVaultRuntimeContext() {
+        RuntimeContext runtimeContext = Mockito.mock(RuntimeContext.class);
+        ExecutionConfig executionConfig = Mockito.mock(ExecutionConfig.class);
+        Mockito.when(runtimeContext.getExecutionConfig()).thenReturn(executionConfig);
+        MetricGroup metricGroup = Mockito.mock(OperatorMetricGroup.class);
+        Mockito.when(runtimeContext.getMetricGroup()).thenReturn(metricGroup);
+        Mockito.when(metricGroup.addGroup(Mockito.anyString())).thenReturn(metricGroup);
+        Mockito.when(metricGroup.counter(Mockito.anyString())).thenReturn(new SimpleCounter());
+        Configuration configuration = new Configuration();
+        CommonConfig commonConfig = new CommonConfig();
+        Map<String, KmsConfig> kmsConfigs = new HashMap<>();
+        KmsConfig kmsConfig = new KmsConfig();
+        kmsConfig.setType(KmsUtils.KMS_TYPE_VAULT);
+        kmsConfigs.put(KmsUtils.KMS_TYPE_VAULT, kmsConfig);
+        kmsConfig = new KmsConfig();
+        kmsConfig.setType(KmsUtils.KMS_TYPE_LOCAL);
+        kmsConfigs.put(KmsUtils.KMS_TYPE_LOCAL, kmsConfig);
+        SSLConfig sslConfig = new SSLConfig();
+        sslConfig.setSkipVerification(true);
+        Map<String, String> propertiesConfig = new HashMap<>();
+        propertiesConfig.put("projection.encrypt.schema.registry.uri", "127.0.0.1:9999/v1/schema/session_record?option=encrypt_fields");
+        commonConfig.setKmsConfig(kmsConfigs);
+        commonConfig.setSslConfig(sslConfig);
+        commonConfig.setPropertiesConfig(propertiesConfig);
+        configuration.setString(Constants.SYSPROP_GROOTSTREAM_CONFIG, JSON.toJSONString(commonConfig));
+        configuration.setString(Constants.SYSPROP_KMS_TYPE_CONFIG, KmsUtils.KMS_TYPE_VAULT);
+        Mockito.when(executionConfig.getGlobalJobParameters()).thenReturn(configuration);
+        return runtimeContext;
+    }
+
+    static MockedStatic<HttpClientPoolUtil> mockSensitiveFields() throws IOException {
+        String sensitiveFieldsStr = "{\n" +
+                "    \"status\": 200,\n" +
+                "    \"success\": true,\n" +
+                "    \"message\": \"Success\",\n" +
+                "    \"data\": [\n" +
+                "        \"phone_number\",\n" +
+                "        \"server_ip\"\n" +
+                "    ]\n" +
+                "}";
+        HttpClientPoolUtil instance = Mockito.mock(HttpClientPoolUtil.class);
+        Mockito.when(instance.httpGet(ArgumentMatchers.any())).thenReturn(sensitiveFieldsStr);
+        MockedStatic<HttpClientPoolUtil> httpClientPoolUtilMockedStatic = Mockito.mockStatic(HttpClientPoolUtil.class);
+        Mockito.when(HttpClientPoolUtil.getInstance()).thenReturn(instance);
+        return httpClientPoolUtilMockedStatic;
+    }
+}