Zjt/encrypt key from environment

java/tsfile/src/main/java/org/apache/tsfile/common/conf/TSFileConfig.java

@@ -156,7 +156,7 @@ public class TSFileConfig implements Serializable {
   private CompressionType compressor = CompressionType.LZ4;
 
   /** encryptKey, this should be 16 bytes String. */
-  private String encryptKey = "abcdefghijklmnop";
+  private byte[] encryptKey = "abcdefghijklmnop".getBytes(TSFileConfig.STRING_CHARSET);
 
   /** Data encryption method, default encryptType is "UNENCRYPTED". */
   private String encryptType = "UNENCRYPTED";
@@ -250,16 +250,16 @@ public void setEncryptType(String encryptType) {
     this.encryptType = encryptType;
   }
 
-  public String getEncryptKey() {
+  public byte[] getEncryptKey() {
     return this.encryptKey;
   }
 
-  public void setEncryptKey(String encryptKey) {
+  public void setEncryptKey(byte[] encryptKey) {
     this.encryptKey = encryptKey;
   }
 
-  public void setEncryptKeyFromPath(String encryptKeyPath) {
-    this.encryptKey = EncryptUtils.getEncryptKeyFromPath(encryptKeyPath);
+  public void setEncryptKeyFromToken(String token) {
+    this.encryptKey = EncryptUtils.getEncryptKeyFromToken(token);
   }
 
   public int getGroupSizeInByte() {

java/tsfile/src/main/java/org/apache/tsfile/common/conf/TSFileDescriptor.java

@@ -83,8 +83,8 @@ public void overwriteConfigByCustomSettings(Properties properties) {
     writer.setString(conf::setCompressor, "compressor");
     writer.setInt(conf::setBatchSize, "batch_size");
     writer.setString(conf::setEncryptType, "encrypt_type");
-    writer.setString(conf::setEncryptKeyFromPath, "encrypt_key_path");
     writer.setBoolean(conf::setLz4UseJni, "lz4_use_jni");
+    conf.setEncryptKeyFromToken(System.getenv("user_encrypt_token"));
   }
 
   private static class PropertiesOverWriter {

java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java

@@ -25,12 +25,17 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.crypto.Mac;
+import javax.crypto.spec.SecretKeySpec;
+
 import java.io.BufferedReader;
 import java.io.FileReader;
 import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
+import java.security.InvalidKeyException;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
 import java.util.Arrays;
 import java.util.Objects;
 
@@ -46,6 +51,12 @@ public class EncryptUtils {
 
   private static volatile EncryptParameter encryptParam;
 
+  private static final String HMAC_ALGORITHM = "HmacSHA256";
+  private static final int ITERATION_COUNT = 1024;
+  private static final int SALT_LENGTH = 16;
+  private static final int INT_SIZE = 4;
+  private static final int dkLen = 16;
+
   public static String getNormalKeyStr() {
     if (normalKeyStr == null) {
       synchronized (EncryptUtils.class) {
@@ -106,6 +117,96 @@ public static String getEncryptKeyFromPath(String path) {
     }
   }
 
+  public static byte[] getEncryptKeyFromToken(String token) {
+    if (token == null || token.trim().isEmpty()) {
+      return defaultKey.getBytes();
+    }
+    byte[] salt = generateSalt();
+    try {
+      return deriveKeyInternal(token.getBytes(), salt, ITERATION_COUNT, dkLen);
+    } catch (NoSuchAlgorithmException | InvalidKeyException e) {
+      throw new EncryptException("Error deriving key from token", e);
+    }
+  }
+
+  private static byte[] deriveKeyInternal(byte[] password, byte[] salt, int c, int dkLen)
+      throws NoSuchAlgorithmException, InvalidKeyException {
+
+    int hLen = getPRFLength();
+
+    if (dkLen < 1) {
+      throw new EncryptException("main key's dkLen must be positive integer: " + dkLen);
+    }
+    if ((long) dkLen > (long) (Math.pow(2, 32) - 1) * hLen) {
+      throw new EncryptException("main key's dkLen is too long: " + dkLen);
+    }
+
+    int n = (int) Math.ceil((double) dkLen / hLen);
+    int r = dkLen - (n - 1) * hLen;
+
+    byte[] blocks = new byte[n * hLen];
+
+    for (int i = 1; i <= n; i++) {
+      byte[] block = F(password, salt, c, i);
+      System.arraycopy(block, 0, blocks, (i - 1) * hLen, hLen);
+    }
+
+    return Arrays.copyOf(blocks, dkLen);
+  }
+
+  /** main function F */
+  private static byte[] F(byte[] password, byte[] salt, int c, int i)
+      throws NoSuchAlgorithmException, InvalidKeyException {
+
+    // U1 = PRF(P, S || INT(i))
+    byte[] input = concatenate(salt, intToBigEndian(i));
+    byte[] U = prf(password, input);
+    byte[] result = U.clone();
+
+    // U2 to Uc
+    for (int j = 2; j <= c; j++) {
+      U = prf(password, U);
+      xorBytes(result, U);
+    }
+
+    return result;
+  }
+
+  /** PRF implementation (HMAC-SHA256) */
+  private static byte[] prf(byte[] key, byte[] data)
+      throws NoSuchAlgorithmException, InvalidKeyException {
+    Mac hmac = Mac.getInstance(HMAC_ALGORITHM);
+    hmac.init(new SecretKeySpec(key, HMAC_ALGORITHM));
+    return hmac.doFinal(data);
+  }
+
+  private static int getPRFLength() throws NoSuchAlgorithmException {
+    return Mac.getInstance(HMAC_ALGORITHM).getMacLength();
+  }
+
+  private static byte[] generateSalt() {
+    byte[] salt = new byte[SALT_LENGTH];
+    new SecureRandom().nextBytes(salt);
+    return salt;
+  }
+
+  private static byte[] intToBigEndian(int i) {
+    return new byte[] {(byte) (i >>> 24), (byte) (i >>> 16), (byte) (i >>> 8), (byte) i};
+  }
+
+  private static void xorBytes(byte[] result, byte[] input) {
+    for (int i = 0; i < result.length; i++) {
+      result[i] ^= input[i];
+    }
+  }
+
+  private static byte[] concatenate(byte[] a, byte[] b) {
+    byte[] output = new byte[a.length + b.length];
+    System.arraycopy(a, 0, output, 0, a.length);
+    System.arraycopy(b, 0, output, a.length, b.length);
+    return output;
+  }
+
   public static byte[] hexStringToByteArray(String hexString) {
     int len = hexString.length();
     byte[] byteArray = new byte[len / 2];
@@ -139,11 +240,10 @@ public static String getNormalKeyStr(TSFileConfig conf) {
           "SHA-256 algorithm not found while using SHA-256 to generate data key", e);
     }
     md.update("IoTDB is the best".getBytes());
-    md.update(conf.getEncryptKey().getBytes());
+    md.update(conf.getEncryptKey());
     byte[] data_key = Arrays.copyOfRange(md.digest(), 0, 16);
     data_key =
-        IEncryptor.getEncryptor(conf.getEncryptType(), conf.getEncryptKey().getBytes())
-            .encrypt(data_key);
+        IEncryptor.getEncryptor(conf.getEncryptType(), conf.getEncryptKey()).encrypt(data_key);
 
     StringBuilder valueStr = new StringBuilder();
 
@@ -180,7 +280,7 @@ public static EncryptParameter getEncryptParameter(TSFileConfig conf) {
             "SHA-256 algorithm not found while using SHA-256 to generate data key", e);
       }
       md.update("IoTDB is the best".getBytes());
-      md.update(conf.getEncryptKey().getBytes());
+      md.update(conf.getEncryptKey());
       dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
     } else {
       encryptType = "org.apache.tsfile.encrypt.UNENCRYPTED";
@@ -227,7 +327,7 @@ public static IEncrypt getEncrypt(TSFileConfig conf) {
             "SHA-256 algorithm not found while using SHA-256 to generate data key", e);
       }
       md.update("IoTDB is the best".getBytes());
-      md.update(conf.getEncryptKey().getBytes());
+      md.update(conf.getEncryptKey());
       dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
     } else {
       encryptType = "org.apache.tsfile.encrypt.UNENCRYPTED";

java/tsfile/src/main/java/org/apache/tsfile/file/metadata/TsFileMetadata.java

@@ -157,7 +157,7 @@ public static TsFileMetadata deserializeFrom(
         IDecryptor decryptor =
             IDecryptor.getDecryptor(
                 propertiesMap.get("encryptType"),
-                TSFileDescriptor.getInstance().getConfig().getEncryptKey().getBytes());
+                TSFileDescriptor.getInstance().getConfig().getEncryptKey());
         String str = propertiesMap.get("encryptKey");
         fileMetaData.dataEncryptKey = decryptor.decrypt(EncryptUtils.getSecondKeyFromStr(str));
         fileMetaData.encryptType = propertiesMap.get("encryptType");

java/tsfile/src/main/java/org/apache/tsfile/write/TsFileWriter.java

@@ -205,10 +205,10 @@ protected TsFileWriter(TsFileIOWriter fileWriter, Schema schema, TSFileConfig co
             "SHA-256 algorithm not found while using SHA-256 to generate data key", e);
       }
       md.update("IoTDB is the best".getBytes());
-      md.update(config.getEncryptKey().getBytes());
+      md.update(config.getEncryptKey());
       dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
       encryptKey =
-          IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey().getBytes())
+          IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey())
               .encrypt(dataEncryptKey);
     } else {
       encryptLevel = "0";

java/tsfile/src/main/java/org/apache/tsfile/write/v4/AbstractTableModelTsFileWriter.java

@@ -117,10 +117,10 @@ protected AbstractTableModelTsFileWriter(File file, long chunkGroupSizeThreshold
             "SHA-256 algorithm not found while using SHA-256 to generate data key", e);
       }
       md.update("IoTDB is the best".getBytes());
-      md.update(config.getEncryptKey().getBytes());
+      md.update(config.getEncryptKey());
       dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
       encryptKey =
-          IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey().getBytes())
+          IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey())
               .encrypt(dataEncryptKey);
     } else {
       encryptLevel = "0";