Skip to content

Commit 06c8300

Browse files
HastyshellHastyshell
committed
support rotating root keys in key manager
1 parent 1defff5 commit 06c8300

File tree

9 files changed

+314
-8
lines changed

9 files changed

+314
-8
lines changed

‎fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -844,4 +844,21 @@ public void testLambdaJoin() {
844844
.assertThrowsExactly(ParseException.class)
845845
.assertMessageContains("mismatched input '->' expecting {<EOF>, ';'}");
846846
}
847+
848+
@Test
849+
public void testAdminRotateTdeRootKey() {
850+
NereidsParser nereidsParser = new NereidsParser();
851+
String sql = "admin rotate tde root key";
852+
nereidsParser.parseSingle(sql);
853+
854+
sql = "admin rotate tde root key properties(\"k\" = \"v\")";
855+
nereidsParser.parseSingle(sql);
856+
857+
sql = "admin rotate tde root key properties(\"k0\" = \"v0\", \"k1\" = \"v1\")";
858+
nereidsParser.parseSingle(sql);
859+
860+
parsePlan("admin rotate tde root key properties()")
861+
.assertThrowsExactly(ParseException.class)
862+
.assertMessageContains("mismatched input ')' expecting");
863+
}
847864
}

‎fe/fe-enterprise/pom.xml‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,24 @@
3030
</dependency>
3131
</dependencies>
3232

33-
<properties>
34-
<maven.compiler.source>8</maven.compiler.source>
35-
<maven.compiler.target>8</maven.compiler.target>
36-
</properties>
33+
<build>
34+
<plugins>
35+
<plugin>
36+
<groupId>org.apache.maven.plugins</groupId>
37+
<artifactId>maven-surefire-plugin</artifactId>
38+
<configuration>
39+
<!-->set larger, eg, 3, to reduce the time or running FE unit tests<-->
40+
<forkCount>3</forkCount>
41+
<!-->not reuse forked jvm, so that each unit test will run in separate jvm. to avoid singleton conflict<-->
42+
<reuseForks>false</reuseForks>
43+
<useFile>false</useFile>
44+
<argLine>
45+
-Xmx512m -javaagent:${settings.localRepository}/org/jmockit/jmockit/${jmockit.version}/jmockit-${jmockit.version}.jar @{argLine}
46+
</argLine>
47+
</configuration>
48+
</plugin>
49+
</plugins>
50+
</build>
51+
3752

3853
</project>

‎fe/fe-enterprise/src/main/java/org/apache/doris/enterprise/AwsKmsRootKeyProvider.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@
3535
import software.amazon.awssdk.services.kms.model.DecryptResponse;
3636
import software.amazon.awssdk.services.kms.model.DescribeKeyRequest;
3737
import software.amazon.awssdk.services.kms.model.DescribeKeyResponse;
38+
import software.amazon.awssdk.services.kms.model.EncryptRequest;
39+
import software.amazon.awssdk.services.kms.model.EncryptResponse;
3840
import software.amazon.awssdk.services.kms.model.GenerateDataKeyRequest;
3941
import software.amazon.awssdk.services.kms.model.GenerateDataKeyResponse;
4042
import software.amazon.awssdk.services.kms.model.KeyMetadata;
@@ -49,6 +51,9 @@ public class AwsKmsRootKeyProvider extends KmsRootKeyProvider {
4951

5052
@Override
5153
public void init(RootKeyInfo info) {
54+
if (kms != null) {
55+
kms.close();
56+
}
5257
AwsCredentialsProvider credProvider;
5358
String ak = System.getenv("DORIS_TDE_AK");
5459
String sk = System.getenv("DORIS_TDE_SK");
@@ -75,6 +80,24 @@ public void init(RootKeyInfo info) {
7580
this.rootKeyInfo = info;
7681
}
7782

83+
public byte[] encrypt(byte[] plaintext) {
84+
try {
85+
EncryptRequest request = EncryptRequest.builder()
86+
.keyId(rootKeyInfo.cmkId)
87+
.plaintext(SdkBytes.fromByteArray(plaintext))
88+
.build();
89+
EncryptResponse response = kms.encrypt(request);
90+
SdkBytes sdkBytes = response.ciphertextBlob();
91+
return sdkBytes.asByteArray();
92+
} catch (KmsException kmsEx) {
93+
LOG.error("KMS exception in describeKey: {}", kmsEx.getMessage(), kmsEx);
94+
throw new RuntimeException("describeKey failed", kmsEx);
95+
} catch (Exception e) {
96+
LOG.error("Unexpected error in describeKey: {}", e.getMessage(), e);
97+
throw new RuntimeException("describeKey failed", e);
98+
}
99+
}
100+
78101
public void describeKey() {
79102
try {
80103
DescribeKeyRequest req = DescribeKeyRequest.builder()

‎fe/fe-enterprise/src/main/java/org/apache/doris/enterprise/KeyManager.java‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,15 @@
1818
package org.apache.doris.enterprise;
1919

2020
import org.apache.doris.catalog.Env;
21+
import org.apache.doris.common.AnalysisException;
2122
import org.apache.doris.common.Config;
2223
import org.apache.doris.encryption.DataKeyMaterial;
2324
import org.apache.doris.encryption.EncryptionKey;
2425
import org.apache.doris.encryption.KeyManagerInterface;
2526
import org.apache.doris.encryption.KeyManagerStore;
2627
import org.apache.doris.encryption.RootKeyInfo;
28+
import org.apache.doris.encryption.RootKeyInfo.RootKeyType;
29+
import org.apache.doris.nereids.trees.plans.commands.AdminRotateTdeRootKeyCommand;
2730
import org.apache.doris.persist.KeyOperationInfo;
2831

2932
import com.google.common.base.Preconditions;
@@ -32,7 +35,10 @@
3235

3336
import java.util.Arrays;
3437
import java.util.Base64;
38+
import java.util.HashMap;
3539
import java.util.List;
40+
import java.util.Map;
41+
import java.util.Objects;
3642
import java.util.zip.CRC32;
3743

3844
public class KeyManager implements KeyManagerInterface {
@@ -157,6 +163,7 @@ public void init() {
157163
public void replayKeyOperation(KeyOperationInfo keyOpInfo) {
158164
store = Env.getCurrentEnv().getKeyManagerStore();
159165
store.setRootKeyInfo(keyOpInfo.getRootKeyInfo());
166+
store.clearMasterKeys();
160167
for (EncryptionKey key : keyOpInfo.getMasterKeys()) {
161168
store.addMasterKey(key);
162169
}
@@ -219,4 +226,86 @@ private void decryptMasterKey() {
219226
LOG.info("Successfully decrypted the plaintext of {}, version {}", versionKey.id, versionKey.version);
220227
}
221228
}
229+
230+
@Override
231+
public void rotateRootKey(Map<String, String> properties) throws Exception {
232+
if (properties == null) {
233+
properties = new HashMap<>();
234+
}
235+
236+
store.writeLock();
237+
try {
238+
RootKeyInfo rootKeyInfo = store.getRootKeyInfo();
239+
RootKeyInfo newRootKeyInfo = new RootKeyInfo(rootKeyInfo);
240+
241+
String keyProviderType = properties.get(AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_PROVIDER);
242+
if (keyProviderType != null) {
243+
newRootKeyInfo.type = RootKeyType.tryFrom(keyProviderType);
244+
// TODO(tsy): remove this branch after local key provider is supported
245+
if (newRootKeyInfo.type == RootKeyType.LOCAL) {
246+
throw new IllegalArgumentException("Local key provider is currently unsupported");
247+
}
248+
}
249+
250+
String kmsMasterKeyId = properties.get(AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_ID);
251+
if (kmsMasterKeyId != null) {
252+
newRootKeyInfo.cmkId = kmsMasterKeyId;
253+
}
254+
255+
String kmsEndpoint = properties.get(AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_ENDPOINT);
256+
if (kmsEndpoint != null) {
257+
newRootKeyInfo.endpoint = kmsEndpoint;
258+
}
259+
260+
String kmsRegion = properties.get(AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_REGION);
261+
if (kmsRegion != null) {
262+
newRootKeyInfo.region = kmsRegion;
263+
}
264+
265+
String password = properties.get(AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_PASSWORD);
266+
if (newRootKeyInfo.type.equals(RootKeyType.LOCAL) && !Objects.equals(password, rootKeyInfo.password)) {
267+
if (password == null) {
268+
throw new IllegalArgumentException("Missing required `"
269+
+ AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_PASSWORD + "` for local key provider");
270+
}
271+
String originalPasswd = properties.get(AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_ORIGINAL_PASSWORD);
272+
if (rootKeyInfo.type.equals(RootKeyType.LOCAL)) {
273+
if (originalPasswd == null) {
274+
throw new IllegalArgumentException("Missing required ` + "
275+
+ AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_ORIGINAL_PASSWORD
276+
+ "` for local key provider");
277+
}
278+
if (!Objects.equals(originalPasswd, rootKeyInfo.password)) {
279+
throw new IllegalArgumentException("Password in `"
280+
+ AdminRotateTdeRootKeyCommand.DORIS_TDE_KEY_ORIGINAL_PASSWORD + "` is incorrect");
281+
}
282+
}
283+
}
284+
if (password != null) {
285+
if (!newRootKeyInfo.type.equals(RootKeyType.LOCAL)) {
286+
throw new IllegalArgumentException("Password is only required for local key provider");
287+
}
288+
newRootKeyInfo.password = password;
289+
}
290+
291+
store.setRootKeyInfo(newRootKeyInfo);
292+
rootKeyProvider.init(newRootKeyInfo);
293+
294+
KeyOperationInfo opInfo = new KeyOperationInfo();
295+
opInfo.setRootKeyInfo(newRootKeyInfo);
296+
List<EncryptionKey> masterKeys = store.getMasterKeys();
297+
for (EncryptionKey masterKey : masterKeys) {
298+
byte[] newCiphertext = rootKeyProvider.encrypt(masterKey.plaintext);
299+
masterKey.ciphertext = Base64.getEncoder().encodeToString(newCiphertext);
300+
long now = System.currentTimeMillis();
301+
masterKey.mtime = now;
302+
masterKey.ctime = now;
303+
opInfo.addMasterKey(masterKey);
304+
}
305+
opInfo.setOpType(KeyOperationInfo.KeyOPType.SET_ROOT_KEY);
306+
Env.getCurrentEnv().getEditLog().logOperateKey(opInfo);
307+
} finally {
308+
store.writeUnlock();
309+
}
310+
}
222311
}

‎fe/fe-enterprise/src/main/java/org/apache/doris/enterprise/KmsRootKeyProvider.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import org.apache.doris.encryption.DataKeyMaterial;
2121
import org.apache.doris.encryption.RootKeyInfo;
2222

23+
import java.util.Map;
24+
2325
class KmsRootKeyProvider implements RootKeyProvider {
2426
@Override
2527
public void init(RootKeyInfo info) {
@@ -30,6 +32,11 @@ public void describeKey() {
3032

3133
}
3234

35+
@Override
36+
public byte[] encrypt(byte[] plaintext) {
37+
return new byte[0];
38+
}
39+
3340
public byte[] decrypt(byte[] ciphertext) {
3441
return null;
3542
}

‎fe/fe-enterprise/src/main/java/org/apache/doris/enterprise/LocalRootKeyProvider.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import java.nio.charset.StandardCharsets;
2424
import java.security.spec.KeySpec;
25+
import java.util.Map;
2526
import javax.crypto.Cipher;
2627
import javax.crypto.SecretKey;
2728
import javax.crypto.SecretKeyFactory;
@@ -56,6 +57,11 @@ public void describeKey() {
5657
System.out.println("LocalRootKeyProvider using password-derived AES-256 root key (CTR mode).");
5758
}
5859

60+
@Override
61+
public byte[] encrypt(byte[] plaintext) {
62+
return new byte[0];
63+
}
64+
5965
@Override
6066
public byte[] decrypt(byte[] ciphertext) {
6167
try {

‎fe/fe-enterprise/src/main/java/org/apache/doris/enterprise/RootKeyProvider.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,17 @@
2020
import org.apache.doris.encryption.DataKeyMaterial;
2121
import org.apache.doris.encryption.RootKeyInfo;
2222

23+
import java.util.Map;
24+
2325
public interface RootKeyProvider {
2426

25-
public void init(RootKeyInfo info);
27+
void init(RootKeyInfo info);
28+
29+
void describeKey();
2630

27-
public void describeKey();
31+
byte[] encrypt(byte[] plaintext);
2832

29-
public byte[] decrypt(byte[] ciphertext);
33+
byte[] decrypt(byte[] ciphertext);
3034

31-
public DataKeyMaterial generateSymmetricDataKey(int length);
35+
DataKeyMaterial generateSymmetricDataKey(int length);
3236
}

0 commit comments

Comments
 (0)