encrypt thread safety (#278)

* encrypt thread safety * simplify duplicate codes --------- Co-authored-by: zhujt <[email protected]>
apache · Oct 31, 2024 · 64a9d8a · 64a9d8a
1 parent 53f86c2
commit 64a9d8a
Show file tree

Hide file tree

Showing 23 changed files with 246 additions and 139 deletions.
diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptParameter.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptParameter.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.tsfile.encrypt;
+
+public class EncryptParameter {
+
+  private final String type;
+  private final byte[] key;
+
+  public EncryptParameter(String type, byte[] key) {
+    this.type = type;
+    this.key = key;
+  }
+
+  public byte[] getKey() {
+    return key;
+  }
+
+  public String getType() {
+    return type;
+  }
+}
diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/EncryptUtils.java
@@ -39,11 +39,7 @@ public class EncryptUtils {
 
   public static String normalKeyStr = getNormalKeyStr();
 
-  public static IEncrypt encrypt = getEncrypt();
-
-  public static IEncryptor encryptor = encrypt.getEncryptor();
-
-  public static IDecryptor decryptor = encrypt.getDecryptor();
+  public static EncryptParameter encryptParam = getEncryptParameter();
 
   public static String getEncryptKeyFromPath(String path) {
     if (path == null) {
@@ -148,15 +144,19 @@ public static String getNormalKeyStr(TSFileConfig conf) {
     }
   }
 
-  public static IEncrypt getEncrypt() {
+  public static EncryptParameter getEncryptParameter() {
+    return getEncryptParameter(TSFileDescriptor.getInstance().getConfig());
+  }
+
+  public static EncryptParameter getEncryptParameter(TSFileConfig conf) {
     String encryptType;
     byte[] dataEncryptKey;
-    if (TSFileDescriptor.getInstance().getConfig().getEncryptFlag()) {
-      encryptType = TSFileDescriptor.getInstance().getConfig().getEncryptType();
+    if (conf.getEncryptFlag()) {
+      encryptType = conf.getEncryptType();
       try {
         MessageDigest md = MessageDigest.getInstance("MD5");
         md.update("IoTDB is the best".getBytes());
-        md.update(TSFileDescriptor.getInstance().getConfig().getEncryptKey().getBytes());
+        md.update(conf.getEncryptKey().getBytes());
         dataEncryptKey = md.digest();
       } catch (Exception e) {
         throw new EncryptException(
@@ -166,10 +166,22 @@ public static IEncrypt getEncrypt() {
       encryptType = "org.apache.tsfile.encrypt.UNENCRYPTED";
       dataEncryptKey = null;
     }
+    return new EncryptParameter(encryptType, dataEncryptKey);
+  }
+
+  public static IEncrypt getEncrypt() {
+    return getEncrypt(TSFileDescriptor.getInstance().getConfig());
+  }
+
+  public static IEncrypt getEncrypt(String encryptType, byte[] dataEncryptKey) {
     try {
+      if (IEncrypt.encryptMap.containsKey(encryptType)) {
+        return ((IEncrypt) IEncrypt.encryptMap.get(encryptType).newInstance(dataEncryptKey));
+      }
       Class<?> encryptTypeClass = Class.forName(encryptType);
       java.lang.reflect.Constructor<?> constructor =
           encryptTypeClass.getDeclaredConstructor(byte[].class);
+      IEncrypt.encryptMap.put(encryptType, constructor);
       return ((IEncrypt) constructor.newInstance(dataEncryptKey));
     } catch (ClassNotFoundException e) {
       throw new EncryptException("Get encryptor class failed: " + encryptType, e);

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IDecryptor.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IDecryptor.java
@@ -51,6 +51,27 @@ static IDecryptor getDecryptor(String type, byte[] key) {
     }
   }
 
+  static IDecryptor getDecryptor(EncryptParameter encryptParam) {
+    String type = encryptParam.getType();
+    byte[] key = encryptParam.getKey();
+    try {
+      if (IEncrypt.encryptMap.containsKey(type)) {
+        return ((IEncrypt) IEncrypt.encryptMap.get(type).newInstance(key)).getDecryptor();
+      }
+      Class<?> encryptClass = Class.forName(type);
+      java.lang.reflect.Constructor<?> constructor =
+          encryptClass.getDeclaredConstructor(byte[].class);
+      IEncrypt.encryptMap.put(type, constructor);
+      return ((IEncrypt) constructor.newInstance(key)).getDecryptor();
+    } catch (ClassNotFoundException e) {
+      throw new EncryptException("Get decryptor class failed: " + type, e);
+    } catch (NoSuchMethodException e) {
+      throw new EncryptException("Get constructor for decryptor failed: " + type, e);
+    } catch (InvocationTargetException | InstantiationException | IllegalAccessException e) {
+      throw new EncryptException("New decryptor instance failed: " + type, e);
+    }
+  }
+
   byte[] decrypt(byte[] data);
 
   byte[] decrypt(byte[] data, int offset, int size);

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptor.java b/java/tsfile/src/main/java/org/apache/tsfile/encrypt/IEncryptor.java
@@ -51,6 +51,27 @@ static IEncryptor getEncryptor(String type, byte[] key) {
     }
   }
 
+  static IEncryptor getEncryptor(EncryptParameter encryptParam) {
+    String type = encryptParam.getType();
+    byte[] key = encryptParam.getKey();
+    try {
+      if (IEncrypt.encryptMap.containsKey(type)) {
+        return ((IEncrypt) IEncrypt.encryptMap.get(type).newInstance(key)).getEncryptor();
+      }
+      Class<?> encryptClass = Class.forName(type);
+      java.lang.reflect.Constructor<?> constructor =
+          encryptClass.getDeclaredConstructor(byte[].class);
+      IEncrypt.encryptMap.put(type, constructor);
+      return ((IEncrypt) constructor.newInstance(key)).getEncryptor();
+    } catch (ClassNotFoundException e) {
+      throw new EncryptException("Get encryptor class failed: " + type, e);
+    } catch (NoSuchMethodException e) {
+      throw new EncryptException("Get constructor for encryptor failed: " + type, e);
+    } catch (InvocationTargetException | InstantiationException | IllegalAccessException e) {
+      throw new EncryptException("New encryptor instance failed: " + type, e);
+    }
+  }
+
   byte[] encrypt(byte[] data);
 
   byte[] encrypt(byte[] data, int offset, int size);

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/file/metadata/TsFileMetadata.java b/java/tsfile/src/main/java/org/apache/tsfile/file/metadata/TsFileMetadata.java
@@ -21,9 +21,9 @@
 
 import org.apache.tsfile.common.conf.TSFileDescriptor;
 import org.apache.tsfile.compatibility.DeserializeConfig;
+import org.apache.tsfile.encrypt.EncryptParameter;
 import org.apache.tsfile.encrypt.EncryptUtils;
 import org.apache.tsfile.encrypt.IDecryptor;
-import org.apache.tsfile.encrypt.IEncryptor;
 import org.apache.tsfile.exception.encrypt.EncryptException;
 import org.apache.tsfile.utils.BloomFilter;
 import org.apache.tsfile.utils.ReadWriteForEncodingUtils;
@@ -162,18 +162,11 @@ public static TsFileMetadata deserializeFrom(ByteBuffer buffer, DeserializeConfi
     return fileMetaData;
   }
 
-  public IEncryptor getIEncryptor() {
+  public EncryptParameter getEncryptParam() {
     if (dataEncryptKey == null) {
-      return IEncryptor.getEncryptor("org.apache.tsfile.encrypt.UNENCRYPTED", null);
+      return new EncryptParameter("org.apache.tsfile.encrypt.UNENCRYPTED", null);
     }
-    return IEncryptor.getEncryptor(encryptType, dataEncryptKey);
-  }
-
-  public IDecryptor getIDecryptor() {
-    if (dataEncryptKey == null) {
-      return IDecryptor.getDecryptor("org.apache.tsfile.encrypt.UNENCRYPTED", null);
-    }
-    return IDecryptor.getDecryptor(encryptType, dataEncryptKey);
+    return new EncryptParameter(encryptType, dataEncryptKey);
   }
 
   public void addProperty(String key, String value) {

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/TsFileSequenceReader.java
@@ -26,6 +26,7 @@
 import org.apache.tsfile.compatibility.DeserializeConfig;
 import org.apache.tsfile.compress.IUnCompressor;
 import org.apache.tsfile.encoding.decoder.Decoder;
+import org.apache.tsfile.encrypt.EncryptParameter;
 import org.apache.tsfile.encrypt.EncryptUtils;
 import org.apache.tsfile.encrypt.IDecryptor;
 import org.apache.tsfile.enums.TSDataType;
@@ -376,14 +377,14 @@ public BloomFilter readBloomFilter() throws IOException {
    * @return the decryptor for the TsFile
    * @throws IOException if an I/O error occurs while reading the file metadata
    */
-  public IDecryptor getDecryptor() throws IOException {
+  public EncryptParameter getEncryptParam() throws IOException {
     try {
       readFileMetadata();
     } catch (Exception e) {
       logger.error("Something error happened while reading file metadata of file {}", file, e);
-      return EncryptUtils.encrypt.getDecryptor();
+      return EncryptUtils.encryptParam;
     }
-    return tsFileMetaData.getIDecryptor();
+    return tsFileMetaData.getEncryptParam();
   }
 
   /**
@@ -1583,7 +1584,7 @@ public Chunk readMemChunk(long offset) throws IOException {
     try {
       ChunkHeader header = readChunkHeader(offset);
       ByteBuffer buffer = readChunk(offset + header.getSerializedSize(), header.getDataSize());
-      return new Chunk(header, buffer, getDecryptor());
+      return new Chunk(header, buffer, getEncryptParam());
     } catch (StopReadTsFileByInterruptException e) {
       throw e;
     } catch (Throwable t) {
@@ -1609,7 +1610,7 @@ public Chunk readMemChunk(ChunkMetadata metaData) throws IOException {
           buffer,
           metaData.getDeleteIntervalList(),
           metaData.getStatistics(),
-          getDecryptor());
+          getEncryptParam());
     } catch (StopReadTsFileByInterruptException e) {
       throw e;
     } catch (Throwable t) {
@@ -1635,7 +1636,7 @@ public Chunk readMemChunk(CachedChunkLoaderImpl.ChunkCacheKey chunkCacheKey) thr
         buffer,
         chunkCacheKey.getDeleteIntervalList(),
         chunkCacheKey.getStatistics(),
-        getDecryptor());
+        getEncryptParam());
   }
 
   /**
@@ -1707,7 +1708,7 @@ public ByteBuffer readCompressedPage(PageHeader header) throws IOException {
 
   public ByteBuffer readPage(PageHeader header, CompressionType type) throws IOException {
     ByteBuffer buffer = readData(-1, header.getCompressedSize());
-    IDecryptor decryptor = getDecryptor();
+    IDecryptor decryptor = IDecryptor.getDecryptor(getEncryptParam());
     if (header.getUncompressedSize() == 0) {
       return buffer;
     }

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/UnClosedTsFileReader.java b/java/tsfile/src/main/java/org/apache/tsfile/read/UnClosedTsFileReader.java
@@ -19,8 +19,8 @@
 
 package org.apache.tsfile.read;
 
+import org.apache.tsfile.encrypt.EncryptParameter;
 import org.apache.tsfile.encrypt.EncryptUtils;
-import org.apache.tsfile.encrypt.IDecryptor;
 import org.apache.tsfile.exception.NotImplementedException;
 import org.apache.tsfile.file.metadata.TsFileMetadata;
 
@@ -29,16 +29,16 @@
 /** A class for reading unclosed tsfile. */
 public class UnClosedTsFileReader extends TsFileSequenceReader {
 
-  private IDecryptor decryptor;
+  private EncryptParameter encryptParam;
 
   public UnClosedTsFileReader(String file) throws IOException {
     super(file, false);
-    decryptor = EncryptUtils.decryptor;
+    encryptParam = EncryptUtils.encryptParam;
   }
 
-  public UnClosedTsFileReader(String file, IDecryptor decryptor) throws IOException {
+  public UnClosedTsFileReader(String file, EncryptParameter decryptParam) throws IOException {
     super(file, false);
-    this.decryptor = decryptor;
+    this.encryptParam = encryptParam;
   }
 
   /** unclosed file has no tail magic data. */
@@ -54,7 +54,7 @@ public TsFileMetadata readFileMetadata() {
   }
 
   @Override
-  public IDecryptor getDecryptor() {
-    return decryptor;
+  public EncryptParameter getEncryptParam() {
+    return encryptParam;
   }
 }
diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/common/Chunk.java b/java/tsfile/src/main/java/org/apache/tsfile/read/common/Chunk.java
@@ -19,8 +19,8 @@
 
 package org.apache.tsfile.read.common;
 
+import org.apache.tsfile.encrypt.EncryptParameter;
 import org.apache.tsfile.encrypt.EncryptUtils;
-import org.apache.tsfile.encrypt.IDecryptor;
 import org.apache.tsfile.file.MetaMarker;
 import org.apache.tsfile.file.header.ChunkHeader;
 import org.apache.tsfile.file.metadata.statistics.Statistics;
@@ -46,7 +46,7 @@ public class Chunk {
   private ByteBuffer chunkData;
   private Statistics chunkStatistic;
 
-  private IDecryptor decryptor;
+  private EncryptParameter encryptParam;
 
   /** A list of deleted intervals. */
   private List<TimeRange> deleteIntervalList;
@@ -60,36 +60,36 @@ public Chunk(
     this.chunkData = buffer;
     this.deleteIntervalList = deleteIntervalList;
     this.chunkStatistic = chunkStatistic;
-    this.decryptor = EncryptUtils.encrypt.getDecryptor();
+    this.encryptParam = EncryptUtils.encryptParam;
   }
 
   public Chunk(
       ChunkHeader header,
       ByteBuffer buffer,
       List<TimeRange> deleteIntervalList,
       Statistics chunkStatistic,
-      IDecryptor decryptor) {
+      EncryptParameter encryptParam) {
     this.chunkHeader = header;
     this.chunkData = buffer;
     this.deleteIntervalList = deleteIntervalList;
     this.chunkStatistic = chunkStatistic;
-    this.decryptor = decryptor;
+    this.encryptParam = encryptParam;
   }
 
   public Chunk(ChunkHeader header, ByteBuffer buffer) {
     this.chunkHeader = header;
     this.chunkData = buffer;
-    this.decryptor = EncryptUtils.encrypt.getDecryptor();
+    this.encryptParam = EncryptUtils.encryptParam;
   }
 
-  public Chunk(ChunkHeader header, ByteBuffer buffer, IDecryptor decryptor) {
+  public Chunk(ChunkHeader header, ByteBuffer buffer, EncryptParameter encryptParam) {
     this.chunkHeader = header;
     this.chunkData = buffer;
-    this.decryptor = decryptor;
+    this.encryptParam = encryptParam;
   }
 
-  public IDecryptor getDecryptor() {
-    return decryptor;
+  public EncryptParameter getEncryptParam() {
+    return encryptParam;
   }
 
   public ChunkHeader getHeader() {

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/read/controller/CachedChunkLoaderImpl.java b/java/tsfile/src/main/java/org/apache/tsfile/read/controller/CachedChunkLoaderImpl.java
@@ -74,7 +74,7 @@ public Chunk loadChunk(ChunkMetadata chunkMetaData) throws IOException {
         chunk.getData().duplicate(),
         chunkMetaData.getDeleteIntervalList(),
         chunkMetaData.getStatistics(),
-        reader.getDecryptor());
+        reader.getEncryptParam());
   }
 
   @Override
@@ -92,7 +92,7 @@ public IChunkReader getChunkReader(IChunkMetadata chunkMetaData, Filter globalTi
             chunk.getData().duplicate(),
             chunkMetaData.getDeleteIntervalList(),
             chunkMetaData.getStatistics(),
-            reader.getDecryptor()),
+            reader.getEncryptParam()),
         globalTimeFilter);
   }