feat: Set up KMS key for clickhouse CMEK and endpoint for PL (#243)

set up KMS key for clickhouse CMEK and endpoint for PL
wandb · Jul 17, 2024 · 1d2fb92 · 1d2fb92
1 parent e03de16
commit 1d2fb92
Show file tree

Hide file tree

Showing 8 changed files with 139 additions and 11 deletions.
diff --git a/main.tf b/main.tf
@@ -1,15 +1,19 @@
 module "kms" {
   source = "./modules/kms"
 
-  key_alias           = var.kms_key_alias == null ? "${var.namespace}-kms-alias" : var.kms_key_alias
   key_deletion_window = var.kms_key_deletion_window
 
+  key_alias  = var.kms_key_alias == null ? "${var.namespace}-kms-alias" : var.kms_key_alias
   key_policy = var.kms_key_policy
+
+  clickhouse_key_alias  = var.kms_clickhouse_key_alias == null ? "${var.namespace}-kms-clickhouse-alias" : var.kms_clickhouse_key_alias
+  clickhouse_key_policy = var.kms_clickhouse_key_policy
 }
 
 locals {
 
   default_kms_key                           = module.kms.key.arn
+  clickhouse_kms_key                        = module.kms.clickhouse_key.arn
   s3_kms_key_arn                            = length(var.bucket_kms_key_arn) > 0 ? var.bucket_kms_key_arn : local.default_kms_key
   database_kms_key_arn                      = length(var.database_kms_key_arn) > 0 ? var.database_kms_key_arn : local.default_kms_key
   database_performance_insights_kms_key_arn = length(var.database_performance_insights_kms_key_arn) > 0 ? var.database_performance_insights_kms_key_arn : local.default_kms_key
@@ -37,12 +41,13 @@ module "networking" {
   namespace  = var.namespace
   create_vpc = var.create_vpc
 
-  cidr                      = var.network_cidr
-  private_subnet_cidrs      = var.network_private_subnet_cidrs
-  public_subnet_cidrs       = var.network_public_subnet_cidrs
-  database_subnet_cidrs     = var.network_database_subnet_cidrs
-  create_elasticache_subnet = var.create_elasticache
-  elasticache_subnet_cidrs  = var.network_elasticache_subnet_cidrs
+  cidr                           = var.network_cidr
+  private_subnet_cidrs           = var.network_private_subnet_cidrs
+  public_subnet_cidrs            = var.network_public_subnet_cidrs
+  database_subnet_cidrs          = var.network_database_subnet_cidrs
+  create_elasticache_subnet      = var.create_elasticache
+  elasticache_subnet_cidrs       = var.network_elasticache_subnet_cidrs
+  clickhouse_endpoint_service_id = var.clickhouse_endpoint_service_id
 }
 
 locals {

diff --git a/modules/kms/main.tf b/modules/kms/main.tf
@@ -89,3 +89,69 @@ resource "aws_kms_grant" "main" {
     "ReEncryptTo",
   ]
 }
+
+resource "aws_kms_key" "clickhouse_key" {
+  deletion_window_in_days = var.key_deletion_window
+  description             = "AWS KMS Customer-managed key to encrypt Weave resources in Clickhouse"
+  key_usage               = "ENCRYPT_DECRYPT"
+
+  policy = var.clickhouse_key_policy != "" ? var.clickhouse_key_policy : jsonencode({
+    "Version" : "2012-10-17",
+    "Statement" : [
+      {
+        "Sid" : "Allow administration of the key",
+        "Effect" : "Allow",
+        "Principal" : { "AWS" : "${data.aws_caller_identity.current.arn}" },
+        "Action" : "kms:*",
+        "Resource" : "*"
+      },
+      {
+        "Sid" : "Allow ClickHouse Access",
+        "Effect" : "Allow",
+        "Principal" : {
+          "AWS" : "arn:aws:iam::576599896960:role/prod-kms-request-role"
+        },
+        "Action" : [
+          "kms:GetPublicKey",
+          "kms:Decrypt",
+          "kms:GenerateDataKeyPair",
+          "kms:Encrypt",
+          "kms:GetKeyRotationStatus",
+          "kms:GenerateDataKey",
+          "kms:DescribeKey"
+        ],
+        "Resource" : "*"
+      },
+    ]
+  })
+
+  tags = {
+    Name = "wandb-kms-clickhouse-key"
+  }
+}
+
+
+
+resource "aws_kms_alias" "clickhouse_key" {
+  name          = "alias/${var.clickhouse_key_alias}"
+  target_key_id = aws_kms_key.clickhouse_key.key_id
+}
+
+
+resource "aws_kms_grant" "clickhouse" {
+  count = var.iam_principal_arn == "" ? 0 : 1
+
+  grantee_principal = var.iam_principal_arn
+  key_id            = aws_kms_key.clickhouse_key.key_id
+  operations = [
+    "Decrypt",
+    "DescribeKey",
+    "Encrypt",
+    "GenerateDataKey",
+    "GenerateDataKeyPair",
+    "GenerateDataKeyPairWithoutPlaintext",
+    "GenerateDataKeyPairWithoutPlaintext",
+    "ReEncryptFrom",
+    "ReEncryptTo",
+  ]
+}
diff --git a/modules/kms/outputs.tf b/modules/kms/outputs.tf
@@ -1,4 +1,10 @@
 output "key" {
   value       = aws_kms_key.key
-  description = "The KMS key used to encrypt data."
-}
+  description = "The KMS key used to encrypt Models data."
+}
+
+
+output "clickhouse_key" {
+  value       = aws_kms_key.clickhouse_key
+  description = "The KMS key used to encrypt Weave data in Clickhouse."
+}
diff --git a/modules/kms/variables.tf b/modules/kms/variables.tf
@@ -18,4 +18,15 @@ variable "key_policy" {
   description = "The policy that will define the permissions for the kms key."
   type        = string
   default     = ""
-}
+}
+
+variable "clickhouse_key_alias" {
+  description = "The key alias for AWS KMS Customer managed key."
+  type        = string
+}
+
+variable "clickhouse_key_policy" {
+  description = "The policy that will define the permissions for the kms clickhouse key."
+  type        = string
+  default     = ""
+}
diff --git a/modules/networking/main.tf b/modules/networking/main.tf
@@ -37,3 +37,13 @@ module "vpc" {
     "kubernetes.io/role/elb" = "1"
   }
 }
+
+resource "aws_vpc_endpoint" "clickhouse" {
+  count = var.create_vpc && var.clickhouse_endpoint_service_id != "" ? 1 : 0
+
+  vpc_id              = module.vpc.vpc_id
+  service_name        = var.clickhouse_endpoint_service_id
+  vpc_endpoint_type   = "Interface"
+  subnet_ids          = module.vpc.private_subnets
+  private_dns_enabled = true
+}
diff --git a/modules/networking/variables.tf b/modules/networking/variables.tf
@@ -61,4 +61,10 @@ variable "amazon_side_asn" {
   description = "The Autonomous System Number (ASN) for the Amazon side of the gateway. By default the virtual private gateway is created with the current default Amazon ASN."
   type        = string
   default     = "64512"
-}
+}
+
+variable "clickhouse_endpoint_service_id" {
+  description = "The ID of the Clickhouse service endpoint"
+  type        = string
+  default     = ""
+}
diff --git a/outputs.tf b/outputs.tf
@@ -52,6 +52,12 @@ output "kms_key_arn" {
   description = "The Amazon Resource Name of the KMS key used to encrypt data at rest."
 }
 
+output "kms_clickhouse_key_arn" {
+  value       = local.clickhouse_kms_key
+  description = "The Amazon Resource Name of the KMS key used to encrypt Weave data at rest in Clickhouse."
+
+}
+
 output "network_id" {
   value       = local.network_id
   description = "The identity of the VPC in which resources are deployed."

diff --git a/variables.tf b/variables.tf
@@ -199,6 +199,18 @@ variable "kms_key_policy" {
   default     = ""
 }
 
+variable "kms_clickhouse_key_alias" {
+  type        = string
+  description = "KMS key alias for AWS KMS Customer managed key used by Clickhouse CMEK."
+  default     = null
+}
+
+variable "kms_clickhouse_key_policy" {
+  type        = string
+  description = "The policy that will define the permissions for the clickhouse kms key."
+  default     = ""
+}
+
 ##########################################
 # Network                                #
 ##########################################
@@ -474,3 +486,9 @@ variable "yace_sa_name" {
   type    = string
   default = "wandb-yace"
 }
+
+variable "clickhouse_endpoint_service_id" {
+  type        = string
+  description = "The service ID of the VPC endpoint service for Clickhouse."
+  default     = ""
+}