Encryption for REST catalog

Author: Sreesh Maheshwar

core/src/main/java/org/apache/iceberg/encryption/EncryptionUtil.java

@@ -87,7 +87,7 @@ public static EncryptionManager createEncryptionManager(
     return createEncryptionManager(List.of(), tableProperties, kmsClient);
   }
 
-  static EncryptionManager createEncryptionManager(
+  public static EncryptionManager createEncryptionManager(
       List<EncryptedKey> keys, Map<String, String> tableProperties, KeyManagementClient kmsClient) {
     Preconditions.checkArgument(kmsClient != null, "Invalid KMS client: null");
     String tableKeyId = tableProperties.get(TableProperties.ENCRYPTION_TABLE_KEY);
@@ -108,7 +108,7 @@ static EncryptionManager createEncryptionManager(
         "Invalid data key length: %s (must be 16, 24, or 32)",
         dataKeyLength);
 
-    return new StandardEncryptionManager(tableKeyId, dataKeyLength, kmsClient);
+    return new StandardEncryptionManager(keys, tableKeyId, dataKeyLength, kmsClient);
   }
 
   public static EncryptedOutputFile plainAsEncryptedOutput(OutputFile encryptingOutputFile) {
@@ -156,4 +156,11 @@ static ByteBuffer encryptManifestListKeyMetadata(
         encryptor.encrypt(keyMetadataBytes, keyId.getBytes(StandardCharsets.UTF_8));
     return ByteBuffer.wrap(encryptedKeyMetadata);
   }
+
+  public static Map<String, EncryptedKey> encryptionKeys(EncryptionManager em) {
+    Preconditions.checkState(
+        em instanceof StandardEncryptionManager,
+        "Encryption keys are only available for StandardEncryptionManager");
+    return ((StandardEncryptionManager) em).encryptionKeys();
+  }
 }

core/src/main/java/org/apache/iceberg/encryption/KeyManagementClient.java

@@ -24,7 +24,7 @@
 import java.util.Map;
 
 /** A minimum client interface to connect to a key management service (KMS). */
-interface KeyManagementClient extends Serializable, Closeable {
+public interface KeyManagementClient extends Serializable, Closeable {
 
   /**
    * Wrap a secret key, using a wrapping/master key which is stored in KMS and referenced by an ID.

core/src/main/java/org/apache/iceberg/encryption/StandardEncryptionManager.java

@@ -23,6 +23,7 @@
 import java.nio.ByteBuffer;
 import java.security.SecureRandom;
 import java.util.Base64;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.apache.iceberg.TableProperties;
@@ -46,9 +47,16 @@ private class TransientEncryptionState {
     private final Map<String, EncryptedKey> encryptionKeys;
     private final LoadingCache<String, ByteBuffer> unwrappedKeyCache;
 
-    private TransientEncryptionState(KeyManagementClient kmsClient) {
+    private TransientEncryptionState(List<EncryptedKey> keys, KeyManagementClient kmsClient) {
       this.kmsClient = kmsClient;
       this.encryptionKeys = Maps.newLinkedHashMap();
+
+      for (EncryptedKey key : keys) {
+        Preconditions.checkArgument(
+            key.keyId() != null, "Key id cannot be null"); // Required by spec.
+        encryptionKeys.put(key.keyId(), key);
+      }
+
       this.unwrappedKeyCache =
           Caffeine.newBuilder()
               .expireAfterWrite(1, TimeUnit.HOURS)
@@ -64,20 +72,33 @@ private TransientEncryptionState(KeyManagementClient kmsClient) {
   private transient volatile SecureRandom lazyRNG = null;
 
   /**
+   * @deprecated will be removed in 1.11.0; use {@link #StandardEncryptionManager(List, String, int,
+   *     KeyManagementClient)} instead.
+   */
+  public StandardEncryptionManager(
+      String tableKeyId, int dataKeyLength, KeyManagementClient kmsClient) {
+    this(List.of(), tableKeyId, dataKeyLength, kmsClient);
+  }
+
+  /**
+   * @param keys a list of existing {@link EncryptedKey}s for this {@link EncryptionManager} to use
    * @param tableKeyId table encryption key id
    * @param dataKeyLength length of data encryption key (16/24/32 bytes)
    * @param kmsClient Client of KMS used to wrap/unwrap keys in envelope encryption
    */
   public StandardEncryptionManager(
-      String tableKeyId, int dataKeyLength, KeyManagementClient kmsClient) {
+      List<EncryptedKey> keys,
+      String tableKeyId,
+      int dataKeyLength,
+      KeyManagementClient kmsClient) {
     Preconditions.checkNotNull(tableKeyId, "Invalid encryption key ID: null");
     Preconditions.checkArgument(
         dataKeyLength == 16 || dataKeyLength == 24 || dataKeyLength == 32,
         "Invalid data key length: %s (must be 16, 24, or 32)",
         dataKeyLength);
     Preconditions.checkNotNull(kmsClient, "Invalid KMS client: null");
     this.tableKeyId = tableKeyId;
-    this.transientState = new TransientEncryptionState(kmsClient);
+    this.transientState = new TransientEncryptionState(keys, kmsClient);
     this.dataKeyLength = dataKeyLength;
   }
 
@@ -199,6 +220,14 @@ public String addManifestListKeyMetadata(NativeEncryptionKeyMetadata keyMetadata
     return manifestListKeyID;
   }
 
+  public Map<String, EncryptedKey> encryptionKeys() {
+    if (transientState == null) {
+      throw new IllegalStateException("Cannot return encryption keys after serialization");
+    }
+
+    return transientState.encryptionKeys;
+  }
+
   private String generateKeyId() {
     byte[] idBytes = new byte[16];
     workerRNG().nextBytes(idBytes);

core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java

@@ -47,6 +47,8 @@
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableCommit;
 import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.encryption.KeyManagementClient;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
 import org.apache.iceberg.exceptions.NoSuchNamespaceException;
 import org.apache.iceberg.exceptions.NoSuchTableException;
@@ -158,6 +160,7 @@ public class RESTSessionCatalog extends BaseViewSessionCatalog
   private Integer pageSize = null;
   private CloseableGroup closeables = null;
   private Set<Endpoint> endpoints;
+  private KeyManagementClient kmsClient = null;
 
   enum SnapshotMode {
     ALL,
@@ -251,6 +254,9 @@ public void initialize(String name, Map<String, String> unresolved) {
 
     this.reportingViaRestEnabled =
         PropertyUtil.propertyAsBoolean(mergedProps, REST_METRICS_REPORTING_ENABLED, true);
+    if (mergedProps.containsKey(CatalogProperties.ENCRYPTION_KMS_IMPL)) {
+      this.kmsClient = EncryptionUtil.createKmsClient(mergedProps);
+    }
     super.initialize(name, mergedProps);
   }
 
@@ -446,6 +452,7 @@ public Table loadTable(SessionContext context, TableIdentifier identifier) {
             paths.table(finalIdentifier),
             Map::of,
             tableFileIO(context, tableConf, response.credentials()),
+            kmsClient,
             tableMetadata,
             endpoints);
 
@@ -525,6 +532,7 @@ public Table registerTable(
             paths.table(ident),
             Map::of,
             tableFileIO(context, tableConf, response.credentials()),
+            kmsClient,
             response.tableMetadata(),
             endpoints);
 
@@ -784,6 +792,7 @@ public Table create() {
               paths.table(ident),
               Map::of,
               tableFileIO(context, tableConf, response.credentials()),
+              kmsClient,
               response.tableMetadata(),
               endpoints);
 
@@ -811,6 +820,7 @@ public Transaction createTransaction() {
               paths.table(ident),
               Map::of,
               tableFileIO(context, tableConf, response.credentials()),
+              kmsClient,
               RESTTableOperations.UpdateType.CREATE,
               createChanges(meta),
               meta,
@@ -874,6 +884,7 @@ public Transaction replaceTransaction() {
               paths.table(ident),
               Map::of,
               tableFileIO(context, tableConf, response.credentials()),
+              kmsClient,
               RESTTableOperations.UpdateType.REPLACE,
               changes.build(),
               base,

core/src/main/java/org/apache/iceberg/rest/RESTTableOperations.java

@@ -31,12 +31,19 @@
 import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.UpdateRequirement;
 import org.apache.iceberg.UpdateRequirements;
+import org.apache.iceberg.encryption.EncryptedKey;
+import org.apache.iceberg.encryption.EncryptingFileIO;
 import org.apache.iceberg.encryption.EncryptionManager;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.encryption.KeyManagementClient;
+import org.apache.iceberg.encryption.PlaintextEncryptionManager;
+import org.apache.iceberg.encryption.StandardEncryptionManager;
 import org.apache.iceberg.io.FileIO;
 import org.apache.iceberg.io.LocationProvider;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 import org.apache.iceberg.rest.requests.UpdateTableRequest;
 import org.apache.iceberg.rest.responses.ErrorResponse;
 import org.apache.iceberg.rest.responses.LoadTableResponse;
@@ -55,27 +62,45 @@ enum UpdateType {
   private final String path;
   private final Supplier<Map<String, String>> headers;
   private final FileIO io;
+  private final KeyManagementClient kmsClient;
   private final List<MetadataUpdate> createChanges;
   private final TableMetadata replaceBase;
   private final Set<Endpoint> endpoints;
   private UpdateType updateType;
   private TableMetadata current;
 
+  private EncryptionManager encryptionManager;
+  private EncryptingFileIO encryptingFileIO;
+  private String encryptionKeyId;
+  private int encryptionDekLength;
+  private List<EncryptedKey> encryptedKeysFromMetadata;
+
   RESTTableOperations(
       RESTClient client,
       String path,
       Supplier<Map<String, String>> headers,
       FileIO io,
+      KeyManagementClient kmsClient,
       TableMetadata current,
       Set<Endpoint> endpoints) {
-    this(client, path, headers, io, UpdateType.SIMPLE, Lists.newArrayList(), current, endpoints);
+    this(
+        client,
+        path,
+        headers,
+        io,
+        kmsClient,
+        UpdateType.SIMPLE,
+        Lists.newArrayList(),
+        current,
+        endpoints);
   }
 
   RESTTableOperations(
       RESTClient client,
       String path,
       Supplier<Map<String, String>> headers,
       FileIO io,
+      KeyManagementClient kmsClient,
       UpdateType updateType,
       List<MetadataUpdate> createChanges,
       TableMetadata current,
@@ -84,6 +109,7 @@ enum UpdateType {
     this.path = path;
     this.headers = headers;
     this.io = io;
+    this.kmsClient = kmsClient;
     this.updateType = updateType;
     this.createChanges = createChanges;
     this.replaceBase = current;
@@ -93,6 +119,10 @@ enum UpdateType {
       this.current = current;
     }
     this.endpoints = endpoints;
+
+    // N.B. We don't use this.current because for tables-to-be-created, because it would be null,
+    // and ee still want encrypted properties in this case for its TableOperations.
+    encryptionPropsFromMetadata(current);
   }
 
   @Override
@@ -113,6 +143,17 @@ public void commit(TableMetadata base, TableMetadata metadata) {
     Consumer<ErrorResponse> errorHandler;
     List<UpdateRequirement> requirements;
     List<MetadataUpdate> updates;
+
+    if (encryption() instanceof StandardEncryptionManager) {
+      TableMetadata.Builder builder = TableMetadata.buildFrom(metadata);
+      for (Map.Entry<String, EncryptedKey> entry :
+          EncryptionUtil.encryptionKeys(encryption()).entrySet()) {
+        builder.addEncryptionKey(entry.getValue());
+      }
+      metadata = builder.build();
+      // TODO(smaheshwar): Think about requirements.
+    }
+
     switch (updateType) {
       case CREATE:
         Preconditions.checkState(
@@ -166,7 +207,67 @@ public void commit(TableMetadata base, TableMetadata metadata) {
 
   @Override
   public FileIO io() {
-    return io;
+    if (encryptionKeyId == null) {
+      return io;
+    }
+
+    if (encryptingFileIO == null) {
+      encryptingFileIO = EncryptingFileIO.combine(io, encryption());
+    }
+
+    return encryptingFileIO;
+  }
+
+  @Override
+  public EncryptionManager encryption() {
+    if (encryptionManager != null) {
+      return encryptionManager;
+    }
+
+    if (encryptionKeyId != null) {
+      if (kmsClient == null) {
+        throw new RuntimeException(
+            "Cant create encryption manager, because key management client is not set");
+      }
+
+      Map<String, String> tableProperties = Maps.newHashMap();
+      tableProperties.put(TableProperties.ENCRYPTION_TABLE_KEY, encryptionKeyId);
+      tableProperties.put(
+          TableProperties.ENCRYPTION_DEK_LENGTH, String.valueOf(encryptionDekLength));
+      encryptionManager =
+          EncryptionUtil.createEncryptionManager(
+              encryptedKeysFromMetadata, tableProperties, kmsClient);
+    } else {
+      return PlaintextEncryptionManager.instance();
+    }
+
+    return encryptionManager;
+  }
+
+  private void encryptionPropsFromMetadata(TableMetadata metadata) {
+    // TODO(smaheshwar): Check generally for changed encryption-related properties!
+    if (metadata == null || metadata.properties() == null) {
+      return;
+    }
+
+    encryptedKeysFromMetadata = metadata.encryptionKeys();
+
+    Map<String, String> tableProperties = metadata.properties();
+    if (encryptionKeyId == null) {
+      encryptionKeyId = tableProperties.get(TableProperties.ENCRYPTION_TABLE_KEY);
+    }
+
+    if (encryptionKeyId != null && encryptionDekLength <= 0) {
+      String dekLength = tableProperties.get(TableProperties.ENCRYPTION_DEK_LENGTH);
+      encryptionDekLength =
+          (dekLength == null)
+              ? TableProperties.ENCRYPTION_DEK_LENGTH_DEFAULT
+              : Integer.parseInt(dekLength);
+    }
+
+    // Force re-creation of encryptingFileIO and encryptionManager
+    encryptingFileIO = null;
+    encryptionManager = null;
   }
 
   private TableMetadata updateCurrentMetadata(LoadTableResponse response) {
@@ -175,6 +276,7 @@ private TableMetadata updateCurrentMetadata(LoadTableResponse response) {
     // safely ignored. there is no requirement to update config on refresh or commit.
     if (current == null
         || !Objects.equals(current.metadataFileLocation(), response.metadataLocation())) {
+      encryptionPropsFromMetadata(response.tableMetadata());
       this.current = response.tableMetadata();
     }