Skip to content

Commit 7859080

Browse files
zhujt20zhujt20
authored
Zjt/encrypt key from environment (#512)
* use environment variable to generate main encrypt key * delete chinese annotation * delete chinese annotation
1 parent a71b8d4 commit 7859080

File tree

6 files changed

+116
-16
lines changed

6 files changed

+116
-16
lines changed

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ public class TSFileConfig implements Serializable {
156156
private CompressionType compressor = CompressionType.LZ4;
157157

158158
/** encryptKey, this should be 16 bytes String. */
159-
private String encryptKey = "abcdefghijklmnop";
159+
private byte[] encryptKey = "abcdefghijklmnop".getBytes(TSFileConfig.STRING_CHARSET);
160160

161161
/** Data encryption method, default encryptType is "UNENCRYPTED". */
162162
private String encryptType = "UNENCRYPTED";
@@ -250,16 +250,16 @@ public void setEncryptType(String encryptType) {
250250
this.encryptType = encryptType;
251251
}
252252

253-
public String getEncryptKey() {
253+
public byte[] getEncryptKey() {
254254
return this.encryptKey;
255255
}
256256

257-
public void setEncryptKey(String encryptKey) {
257+
public void setEncryptKey(byte[] encryptKey) {
258258
this.encryptKey = encryptKey;
259259
}
260260

261-
public void setEncryptKeyFromPath(String encryptKeyPath) {
262-
this.encryptKey = EncryptUtils.getEncryptKeyFromPath(encryptKeyPath);
261+
public void setEncryptKeyFromToken(String token) {
262+
this.encryptKey = EncryptUtils.getEncryptKeyFromToken(token);
263263
}
264264

265265
public int getGroupSizeInByte() {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,8 @@ public void overwriteConfigByCustomSettings(Properties properties) {
8383
writer.setString(conf::setCompressor, "compressor");
8484
writer.setInt(conf::setBatchSize, "batch_size");
8585
writer.setString(conf::setEncryptType, "encrypt_type");
86-
writer.setString(conf::setEncryptKeyFromPath, "encrypt_key_path");
8786
writer.setBoolean(conf::setLz4UseJni, "lz4_use_jni");
87+
conf.setEncryptKeyFromToken(System.getenv("user_encrypt_token"));
8888
}
8989

9090
private static class PropertiesOverWriter {

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

Lines changed: 105 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,17 @@
2525
import org.slf4j.Logger;
2626
import org.slf4j.LoggerFactory;
2727

28+
import javax.crypto.Mac;
29+
import javax.crypto.spec.SecretKeySpec;
30+
2831
import java.io.BufferedReader;
2932
import java.io.FileReader;
3033
import java.io.IOException;
3134
import java.lang.reflect.InvocationTargetException;
35+
import java.security.InvalidKeyException;
3236
import java.security.MessageDigest;
3337
import java.security.NoSuchAlgorithmException;
38+
import java.security.SecureRandom;
3439
import java.util.Arrays;
3540
import java.util.Objects;
3641

@@ -46,6 +51,12 @@ public class EncryptUtils {
4651

4752
private static volatile EncryptParameter encryptParam;
4853

54+
private static final String HMAC_ALGORITHM = "HmacSHA256";
55+
private static final int ITERATION_COUNT = 1024;
56+
private static final int SALT_LENGTH = 16;
57+
private static final int INT_SIZE = 4;
58+
private static final int dkLen = 16;
59+
4960
public static String getNormalKeyStr() {
5061
if (normalKeyStr == null) {
5162
synchronized (EncryptUtils.class) {
@@ -106,6 +117,96 @@ public static String getEncryptKeyFromPath(String path) {
106117
}
107118
}
108119

120+
public static byte[] getEncryptKeyFromToken(String token) {
121+
if (token == null || token.trim().isEmpty()) {
122+
return defaultKey.getBytes();
123+
}
124+
byte[] salt = generateSalt();
125+
try {
126+
return deriveKeyInternal(token.getBytes(), salt, ITERATION_COUNT, dkLen);
127+
} catch (NoSuchAlgorithmException | InvalidKeyException e) {
128+
throw new EncryptException("Error deriving key from token", e);
129+
}
130+
}
131+
132+
private static byte[] deriveKeyInternal(byte[] password, byte[] salt, int c, int dkLen)
133+
throws NoSuchAlgorithmException, InvalidKeyException {
134+
135+
int hLen = getPRFLength();
136+
137+
if (dkLen < 1) {
138+
throw new EncryptException("main key's dkLen must be positive integer: " + dkLen);
139+
}
140+
if ((long) dkLen > (long) (Math.pow(2, 32) - 1) * hLen) {
141+
throw new EncryptException("main key's dkLen is too long: " + dkLen);
142+
}
143+
144+
int n = (int) Math.ceil((double) dkLen / hLen);
145+
int r = dkLen - (n - 1) * hLen;
146+
147+
byte[] blocks = new byte[n * hLen];
148+
149+
for (int i = 1; i <= n; i++) {
150+
byte[] block = F(password, salt, c, i);
151+
System.arraycopy(block, 0, blocks, (i - 1) * hLen, hLen);
152+
}
153+
154+
return Arrays.copyOf(blocks, dkLen);
155+
}
156+
157+
/** main function F */
158+
private static byte[] F(byte[] password, byte[] salt, int c, int i)
159+
throws NoSuchAlgorithmException, InvalidKeyException {
160+
161+
// U1 = PRF(P, S || INT(i))
162+
byte[] input = concatenate(salt, intToBigEndian(i));
163+
byte[] U = prf(password, input);
164+
byte[] result = U.clone();
165+
166+
// U2 to Uc
167+
for (int j = 2; j <= c; j++) {
168+
U = prf(password, U);
169+
xorBytes(result, U);
170+
}
171+
172+
return result;
173+
}
174+
175+
/** PRF implementation (HMAC-SHA256) */
176+
private static byte[] prf(byte[] key, byte[] data)
177+
throws NoSuchAlgorithmException, InvalidKeyException {
178+
Mac hmac = Mac.getInstance(HMAC_ALGORITHM);
179+
hmac.init(new SecretKeySpec(key, HMAC_ALGORITHM));
180+
return hmac.doFinal(data);
181+
}
182+
183+
private static int getPRFLength() throws NoSuchAlgorithmException {
184+
return Mac.getInstance(HMAC_ALGORITHM).getMacLength();
185+
}
186+
187+
private static byte[] generateSalt() {
188+
byte[] salt = new byte[SALT_LENGTH];
189+
new SecureRandom().nextBytes(salt);
190+
return salt;
191+
}
192+
193+
private static byte[] intToBigEndian(int i) {
194+
return new byte[] {(byte) (i >>> 24), (byte) (i >>> 16), (byte) (i >>> 8), (byte) i};
195+
}
196+
197+
private static void xorBytes(byte[] result, byte[] input) {
198+
for (int i = 0; i < result.length; i++) {
199+
result[i] ^= input[i];
200+
}
201+
}
202+
203+
private static byte[] concatenate(byte[] a, byte[] b) {
204+
byte[] output = new byte[a.length + b.length];
205+
System.arraycopy(a, 0, output, 0, a.length);
206+
System.arraycopy(b, 0, output, a.length, b.length);
207+
return output;
208+
}
209+
109210
public static byte[] hexStringToByteArray(String hexString) {
110211
int len = hexString.length();
111212
byte[] byteArray = new byte[len / 2];
@@ -139,11 +240,10 @@ public static String getNormalKeyStr(TSFileConfig conf) {
139240
"SHA-256 algorithm not found while using SHA-256 to generate data key", e);
140241
}
141242
md.update("IoTDB is the best".getBytes());
142-
md.update(conf.getEncryptKey().getBytes());
243+
md.update(conf.getEncryptKey());
143244
byte[] data_key = Arrays.copyOfRange(md.digest(), 0, 16);
144245
data_key =
145-
IEncryptor.getEncryptor(conf.getEncryptType(), conf.getEncryptKey().getBytes())
146-
.encrypt(data_key);
246+
IEncryptor.getEncryptor(conf.getEncryptType(), conf.getEncryptKey()).encrypt(data_key);
147247

148248
StringBuilder valueStr = new StringBuilder();
149249

@@ -180,7 +280,7 @@ public static EncryptParameter getEncryptParameter(TSFileConfig conf) {
180280
"SHA-256 algorithm not found while using SHA-256 to generate data key", e);
181281
}
182282
md.update("IoTDB is the best".getBytes());
183-
md.update(conf.getEncryptKey().getBytes());
283+
md.update(conf.getEncryptKey());
184284
dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
185285
} else {
186286
encryptType = "org.apache.tsfile.encrypt.UNENCRYPTED";
@@ -227,7 +327,7 @@ public static IEncrypt getEncrypt(TSFileConfig conf) {
227327
"SHA-256 algorithm not found while using SHA-256 to generate data key", e);
228328
}
229329
md.update("IoTDB is the best".getBytes());
230-
md.update(conf.getEncryptKey().getBytes());
330+
md.update(conf.getEncryptKey());
231331
dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
232332
} else {
233333
encryptType = "org.apache.tsfile.encrypt.UNENCRYPTED";

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ public static TsFileMetadata deserializeFrom(
157157
IDecryptor decryptor =
158158
IDecryptor.getDecryptor(
159159
propertiesMap.get("encryptType"),
160-
TSFileDescriptor.getInstance().getConfig().getEncryptKey().getBytes());
160+
TSFileDescriptor.getInstance().getConfig().getEncryptKey());
161161
String str = propertiesMap.get("encryptKey");
162162
fileMetaData.dataEncryptKey = decryptor.decrypt(EncryptUtils.getSecondKeyFromStr(str));
163163
fileMetaData.encryptType = propertiesMap.get("encryptType");

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -205,10 +205,10 @@ protected TsFileWriter(TsFileIOWriter fileWriter, Schema schema, TSFileConfig co
205205
"SHA-256 algorithm not found while using SHA-256 to generate data key", e);
206206
}
207207
md.update("IoTDB is the best".getBytes());
208-
md.update(config.getEncryptKey().getBytes());
208+
md.update(config.getEncryptKey());
209209
dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
210210
encryptKey =
211-
IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey().getBytes())
211+
IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey())
212212
.encrypt(dataEncryptKey);
213213
} else {
214214
encryptLevel = "0";

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,10 +117,10 @@ protected AbstractTableModelTsFileWriter(File file, long chunkGroupSizeThreshold
117117
"SHA-256 algorithm not found while using SHA-256 to generate data key", e);
118118
}
119119
md.update("IoTDB is the best".getBytes());
120-
md.update(config.getEncryptKey().getBytes());
120+
md.update(config.getEncryptKey());
121121
dataEncryptKey = Arrays.copyOfRange(md.digest(), 0, 16);
122122
encryptKey =
123-
IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey().getBytes())
123+
IEncryptor.getEncryptor(config.getEncryptType(), config.getEncryptKey())
124124
.encrypt(dataEncryptKey);
125125
} else {
126126
encryptLevel = "0";

0 commit comments

Comments
 (0)