ggershinsky commented on code in PR #7770:
URL: https://github.com/apache/iceberg/pull/7770#discussion_r1703115064
##########
core/src/main/java/org/apache/iceberg/encryption/StandardEncryptionManager.java:
##########
@@ -92,13 +112,98 @@ public ByteBuffer wrapKey(ByteBuffer secretKey) {
public ByteBuffer unwrapKey(ByteBuffer wrappedSecretKey) {
if (kmsClient == null) {
- throw new IllegalStateException(
- "Cannot wrap key after called after serialization (missing KMS
client)");
+ throw new IllegalStateException("Cannot unwrap key after serialization
(missing KMS client)");
}
return kmsClient.unwrapKey(wrappedSecretKey, tableKeyId);
}
+ public void addKekCache(Map<String, KeyEncryptionKey> wrappedKekCache) {
+ for (Map.Entry<String, KeyEncryptionKey> entry :
wrappedKekCache.entrySet()) {
+ KeyEncryptionKey wrappedKek = entry.getValue();
+ KeyEncryptionKey cachedKek = kekCache.get(entry.getKey());
+
+ if (cachedKek != null) {
+ Preconditions.checkState(
+ cachedKek.wrappedKey().equals(wrappedKek.wrappedKey()),
+ "Cached kek wrap differs from newly added for %s",
+ entry.getKey());
+ } else {
+ ByteBuffer encryptedKEK =
+
ByteBuffer.wrap(Base64.getDecoder().decode(wrappedKek.wrappedKey()));
+ // Unwrap the key in KMS
+ byte[] kekBytes = unwrapKey(encryptedKEK).array();
+
+ kekCache.put(
+ entry.getKey(),
+ new KeyEncryptionKey(
+ wrappedKek.id(), kekBytes, wrappedKek.wrappedKey(),
wrappedKek.timestamp()));
+ }
+ }
+ }
+
+ public KeyEncryptionKey currentKEK() {
+ if (kekCache.isEmpty()) {
+ KeyEncryptionKey keyEncryptionKey = generateNewKEK();
+ kekCache.put(keyEncryptionKey.id(), keyEncryptionKey);
+
+ return keyEncryptionKey;
+ }
+
+ KeyEncryptionKey lastKek = lastKek();
+ long timeNow = System.currentTimeMillis();
+ if (timeNow - lastKek.timestamp() > kekCacheTimeout) {
+ KeyEncryptionKey keyEncryptionKey = generateNewKEK();
+ kekCache.put(keyEncryptionKey.id(), keyEncryptionKey);
+
+ return keyEncryptionKey;
+ }
+
+ return lastKek;
+ }
+
+ private KeyEncryptionKey lastKek() {
+ int counter = 0;
+ KeyEncryptionKey result = null;
+ for (Map.Entry<String, KeyEncryptionKey> entry : kekCache.entrySet()) {
+ if (counter == 0) {
+ result = entry.getValue();
+ } else {
+ if (entry.getValue().timestamp() > result.timestamp()) {
+ result = entry.getValue();
+ }
+ }
+
+ counter++;
+ }
+
+ return result;
+ }
+
+ private KeyEncryptionKey generateNewKEK() {
+ byte[] newKek = new byte[dataKeyLength];
+ workerRNG().nextBytes(newKek);
+ String wrappedNewKek =
+
Base64.getEncoder().encodeToString(wrapKey(ByteBuffer.wrap(newKek)).array());
+
+ byte[] idBytes = new byte[KEK_ID_LENGTH];
+ workerRNG().nextBytes(idBytes);
+ String kekID = Base64.getEncoder().encodeToString(idBytes);
+
+ return new KeyEncryptionKey(kekID, newKek, wrappedNewKek,
System.currentTimeMillis());
+ }
+
+ public KeyEncryptionKey keyEncryptionKey(String kekID) {
+ KeyEncryptionKey result = kekCache.get(kekID);
+ Preconditions.checkState(result != null, "Key encryption key %s not
found", kekID);
+
+ return result;
+ }
+
+ public Map<String, KeyEncryptionKey> kekCache() {
Review Comment:
Ok, I'll move this to the EncryptionUtil .
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org
For additional commands, e-mail: issues-h...@iceberg.apache.org
