feat: add CMEK Support (#656)

* feat: Implement hand written layer on top of CMEK protos

* fix tests

* fix tests

* fix: fix config file

* fix: fix test

* fix: fix and cleanup tests

* fix: lint fix

Co-authored-by: Kristen O'Leary <[email protected]>

Author: ad548

.kokoro/nightly/integration.cfg

@@ -8,7 +8,7 @@ env_vars: {
 
 env_vars: {
     key: "INTEGRATION_TEST_ARGS"
-    value: "-P bigtable-emulator-it,bigtable-prod-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests"
+    value: "-P bigtable-emulator-it,bigtable-prod-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests -Dbigtable.kms_key_name=projects/gcloud-devel/locations/us-central1/keyRings/cmek-test-key-ring/cryptoKeys/cmek-test-key -Dbigtable.wait-for-cmek-key-status=true"
 }
 
 env_vars: {

.kokoro/presubmit/integration.cfg

@@ -8,7 +8,7 @@ env_vars: {
 
 env_vars: {
     key: "INTEGRATION_TEST_ARGS"
-    value: "-P bigtable-emulator-it,bigtable-prod-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests"
+    value: "-P bigtable-emulator-it,bigtable-prod-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests -Dbigtable.kms_key_name=projects/gcloud-devel/locations/us-central1/keyRings/cmek-test-key-ring/cryptoKeys/cmek-test-key -Dbigtable.wait-for-cmek-key-status=true"
 }
 
 env_vars: {

google-cloud-bigtable/pom.xml

@@ -214,6 +214,11 @@
       <artifactId>truth</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.truth.extensions</groupId>
+      <artifactId>truth-proto-extension</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-testing</artifactId>

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java

@@ -30,6 +30,8 @@
 import com.google.bigtable.admin.v2.ListBackupsRequest;
 import com.google.bigtable.admin.v2.ListTablesRequest;
 import com.google.bigtable.admin.v2.RestoreTableMetadata;
+import com.google.bigtable.admin.v2.Table.ClusterState;
+import com.google.bigtable.admin.v2.Table.View;
 import com.google.cloud.Policy;
 import com.google.cloud.Policy.DefaultMarshaller;
 import com.google.cloud.bigtable.admin.v2.BaseBigtableTableAdminClient.ListBackupsPage;
@@ -40,6 +42,7 @@
 import com.google.cloud.bigtable.admin.v2.models.Backup;
 import com.google.cloud.bigtable.admin.v2.models.CreateBackupRequest;
 import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
+import com.google.cloud.bigtable.admin.v2.models.EncryptionInfo;
 import com.google.cloud.bigtable.admin.v2.models.GCRules;
 import com.google.cloud.bigtable.admin.v2.models.ModifyColumnFamiliesRequest;
 import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
@@ -49,6 +52,8 @@
 import com.google.cloud.bigtable.admin.v2.models.UpdateBackupRequest;
 import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.iam.v1.GetIamPolicyRequest;
@@ -60,6 +65,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nonnull;
 
@@ -519,6 +525,52 @@ private ApiFuture<Table> getTableAsync(
     return transformToTableResponse(this.stub.getTableCallable().futureCall(request));
   }
 
+  /**
+   * Gets the current encryption info for the table across all of the clusters.
+   *
+   * <p>The returned Map will be keyed by cluster id and contain a status for all of the keys in
+   * use.
+   */
+  public Map<String, List<EncryptionInfo>> getEncryptionInfo(String tableId) {
+    return ApiExceptions.callAndTranslateApiException(getEncryptionInfoAsync(tableId));
+  }
+
+  /**
+   * Asynchronously gets the current encryption info for the table across all of the clusters.
+   *
+   * <p>The returned Map will be keyed by cluster id and contain a status for all of the keys in
+   * use.
+   */
+  public ApiFuture<Map<String, List<EncryptionInfo>>> getEncryptionInfoAsync(String tableId) {
+    GetTableRequest request =
+        GetTableRequest.newBuilder()
+            .setName(getTableName(tableId))
+            .setView(View.ENCRYPTION_VIEW)
+            .build();
+    return ApiFutures.transform(
+        this.stub.getTableCallable().futureCall(request),
+        new ApiFunction<com.google.bigtable.admin.v2.Table, Map<String, List<EncryptionInfo>>>() {
+          @Override
+          public Map<String, List<EncryptionInfo>> apply(com.google.bigtable.admin.v2.Table table) {
+            ImmutableMap.Builder<String, List<EncryptionInfo>> result = ImmutableMap.builder();
+
+            for (Map.Entry<String, ClusterState> entry : table.getClusterStatesMap().entrySet()) {
+              ImmutableList.Builder<EncryptionInfo> infos = ImmutableList.builder();
+
+              for (com.google.bigtable.admin.v2.EncryptionInfo infoProto :
+                  entry.getValue().getEncryptionInfoList()) {
+                infos.add(EncryptionInfo.fromProto(infoProto));
+              }
+
+              result.put(entry.getKey(), infos.build());
+            }
+
+            return result.build();
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
   /**
    * Lists all table IDs in the instance.
    *

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/Backup.java

@@ -142,6 +142,18 @@ public State getState() {
     return State.fromProto(proto.getState());
   }
 
+  /**
+   * Get the encryption information for the backup.
+   *
+   * <p>If encryption_type is CUSTOMER_MANAGED_ENCRYPTION, kms_key_version will be filled in with
+   * status UNKNOWN.
+   *
+   * <p>If encryption_type is GOOGLE_DEFAULT_ENCRYPTION, all other fields will have default value.
+   */
+  public EncryptionInfo getEncryptionInfo() {
+    return EncryptionInfo.fromProto(proto.getEncryptionInfo());
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/Cluster.java

@@ -30,6 +30,7 @@
  * in the instance.
  */
 public class Cluster {
+
   public enum State {
     /** The state of the cluster could not be determined. */
     NOT_KNOWN(com.google.bigtable.admin.v2.Cluster.State.STATE_NOT_KNOWN),
@@ -156,6 +157,18 @@ public StorageType getStorageType() {
     return StorageType.fromProto(stateProto.getDefaultStorageType());
   }
 
+  /**
+   * Google Cloud Key Management Service (KMS) settings for a CMEK-protected Bigtable cluster. This
+   * returns the full resource name of the Cloud KMS key in the format
+   * `projects/{key_project_id}/locations/{location}/keyRings/{ring_name}/cryptoKeys/{key_name}`
+   */
+  public String getKmsKeyName() {
+    if (stateProto.hasEncryptionConfig()) {
+      return stateProto.getEncryptionConfig().getKmsKeyName();
+    }
+    return null;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateClusterRequest.java

@@ -49,6 +49,7 @@
  *     details</a>
  */
 public final class CreateClusterRequest {
+
   private final com.google.bigtable.admin.v2.CreateClusterRequest.Builder proto =
       com.google.bigtable.admin.v2.CreateClusterRequest.newBuilder();
   // instanceId and zone are short ids, which will be expanded to full names when the project name
@@ -104,6 +105,17 @@ public CreateClusterRequest setStorageType(@Nonnull StorageType storageType) {
     return this;
   }
 
+  /**
+   * Sets the Google Cloud Key Management Service (KMS) key for a CMEK-protected Bigtable. This
+   * requires the full resource name of the Cloud KMS key, in the format
+   * `projects/{key_project_id}/locations/{location}/keyRings/{ring_name}/cryptoKeys/{key_name}`
+   */
+  public CreateClusterRequest setKmsKeyName(@Nonnull String kmsKeyName) {
+    Preconditions.checkNotNull(kmsKeyName);
+    proto.getClusterBuilder().getEncryptionConfigBuilder().setKmsKeyName(kmsKeyName);
+    return this;
+  }
+
   /**
    * Creates the request protobuf. This method is considered an internal implementation detail and
    * not meant to be used by applications.

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateInstanceRequest.java

@@ -147,6 +147,35 @@ public CreateInstanceRequest addCluster(
     return this;
   }
 
+  /**
+   * Adds a CMEK protected cluster using the specified KMS key name.
+   *
+   * @param clusterId the name of the cluster.
+   * @param zone the zone where the cluster will be created.
+   * @param serveNodes the number of nodes that cluster will contain. DEVELOPMENT instance clusters
+   *     must have exactly one node.
+   * @param storageType the type of storage used by this cluster to serve its parent instance's
+   *     tables.
+   * @param kmsKeyName the full name of the KMS key name to use in the format
+   *     `projects/{key_project_id}/locations/{location}/keyRings/{ring_name}/cryptoKeys/{key_name}`
+   */
+  public CreateInstanceRequest addCmekCluster(
+      @Nonnull String clusterId,
+      @Nonnull String zone,
+      int serveNodes,
+      @Nonnull StorageType storageType,
+      @Nonnull String kmsKeyName) {
+    CreateClusterRequest clusterRequest =
+        CreateClusterRequest.of("ignored-instance-id", clusterId)
+            .setZone(zone)
+            .setServeNodes(serveNodes)
+            .setStorageType(storageType)
+            .setKmsKeyName(kmsKeyName);
+    clusterRequests.add(clusterRequest);
+
+    return this;
+  }
+
   /**
    * Adds a DEVELOPMENT cluster to the instance request.
    *

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/EncryptionInfo.java

@@ -0,0 +1,108 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed 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
+ *
+ *     https://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 com.google.cloud.bigtable.admin.v2.models;
+
+import com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType;
+import com.google.cloud.bigtable.common.Status;
+import com.google.common.base.Objects;
+
+/**
+ * Encryption information for a given resource.
+ *
+ * <p>If this resource is protected with customer managed encryption, the in-use Google Cloud Key
+ * Management Service (KMS) key versions will be specified along with their status.
+ */
+public final class EncryptionInfo {
+  public enum Type {
+    /** Encryption type was not specified, though data at rest remains encrypted. */
+    ENCRYPTION_TYPE_UNSPECIFIED(
+        com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType.ENCRYPTION_TYPE_UNSPECIFIED),
+    /**
+     * The data backing this resource is encrypted at rest with a key that is fully managed by
+     * Google. No key version or status will be populated.
+     */
+    GOOGLE_DEFAULT_ENCRYPTION(
+        com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION),
+    /**
+     * The data backing this resource is encrypted at rest with a key that is fully managed by
+     * Google. No key version or status will be populated. This is the default state.
+     */
+    CUSTOMER_MANAGED_ENCRYPTION(
+        com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType.CUSTOMER_MANAGED_ENCRYPTION),
+    /** Type not known by the client, please upgrade your client */
+    UNRECOGNIZED(com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType.UNRECOGNIZED);
+
+    private final com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType proto;
+
+    Type(EncryptionType proto) {
+      this.proto = proto;
+    }
+
+    /** Wraps the EncryptionInfo protobuf. */
+    public static Type fromProto(com.google.bigtable.admin.v2.EncryptionInfo.EncryptionType proto) {
+      for (Type type : values()) {
+        if (Objects.equal(type.proto, proto)) {
+          return type;
+        }
+      }
+      return UNRECOGNIZED;
+    }
+  }
+
+  private com.google.bigtable.admin.v2.EncryptionInfo proto;
+
+  public static EncryptionInfo fromProto(com.google.bigtable.admin.v2.EncryptionInfo proto) {
+    return new EncryptionInfo(proto);
+  }
+
+  private EncryptionInfo(com.google.bigtable.admin.v2.EncryptionInfo proto) {
+    this.proto = proto;
+  }
+
+  public Type getType() {
+    return EncryptionInfo.Type.fromProto(proto.getEncryptionType());
+  }
+
+  public String getKmsKeyVersion() {
+    return proto.getKmsKeyVersion();
+  }
+
+  public Status getStatus() {
+    return Status.fromProto(proto.getEncryptionStatus());
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    EncryptionInfo that = (EncryptionInfo) o;
+    return Objects.equal(proto, that.proto);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(proto);
+  }
+
+  @Override
+  public String toString() {
+    return proto.toString();
+  }
+}