Merge pull request #8 from braintrustdata/mdeeks/kms

Add support for KMS keys

Author: mdeeks

main.tf

@@ -1,3 +1,14 @@
+module "kms" {
+  source = "./modules/kms"
+  count  = var.kms_key_arn == null ? 1 : 0
+
+  deployment_name = var.deployment_name
+}
+
+locals {
+  kms_key_arn = var.kms_key_arn != null ? var.kms_key_arn : module.kms[0].key_arn
+}
+
 module "main_vpc" {
   source = "./modules/vpc"
 
@@ -49,6 +60,8 @@ module "database" {
 
   postgres_storage_iops       = var.postgres_storage_iops
   postgres_storage_throughput = var.postgres_storage_throughput
+
+  kms_key_arn = local.kms_key_arn
 }
 
 module "redis" {
@@ -106,4 +119,5 @@ module "services" {
     module.quarantine_vpc[0].private_subnet_3_id
   ] : []
 
+  kms_key_arn = local.kms_key_arn
 }

modules/database/main.tf

@@ -36,6 +36,8 @@ resource "aws_db_instance" "main" {
 
   performance_insights_enabled          = true
   performance_insights_retention_period = 7
+
+  kms_key_id = var.kms_key_arn
 }
 
 resource "aws_db_parameter_group" "main" {
@@ -117,4 +119,5 @@ resource "aws_secretsmanager_secret_version" "database_secret" {
 resource "aws_secretsmanager_secret" "database_secret" {
   name_prefix = "${var.deployment_name}/DatabaseSecret-"
   description = "Username/password for the main Braintrust RDS database"
+  kms_key_id  = var.kms_key_arn
 }

modules/database/variables.tf

@@ -49,3 +49,9 @@ variable "database_security_group_ids" {
   description = "Security Group IDs for the RDS instance."
   type        = list(string)
 }
+
+variable "kms_key_arn" {
+  description = "KMS key ARN to use for encrypting resources. If not provided, the default AWS managed key is used. DO NOT change this after deployment. If you do, it will attempt to destroy your DB."
+  type        = string
+  default     = null
+}

modules/database/versions.tf

@@ -1,9 +1,9 @@
 terraform {
-  required_version = ">= 1.0.0"
+  required_version = ">= 1.9.0"
   required_providers {
     aws = {
       source  = "hashicorp/aws"
       version = "~> 5.0"
     }
   }
-}
+}

modules/elasticache/versions.tf

@@ -1,9 +1,9 @@
 terraform {
-  required_version = ">= 1.0.0"
+  required_version = ">= 1.9.0"
   required_providers {
     aws = {
       source  = "hashicorp/aws"
       version = "~> 5.0"
     }
   }
-}
+}

modules/kms/main.tf

@@ -0,0 +1,33 @@
+resource "aws_kms_key" "braintrust" {
+  description             = "KMS key for encrypting resources in the Braintrust data plane"
+  key_usage               = "ENCRYPT_DECRYPT"
+  deletion_window_in_days = 7
+  enable_key_rotation     = true
+
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Sid    = "Enable IAM User Permissions"
+        Effect = "Allow"
+        Principal = {
+          AWS = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:root"
+        }
+        Action   = "kms:*"
+        Resource = "*"
+      }
+    ]
+  })
+
+  tags = {
+    Name        = "${var.deployment_name}-main"
+    Environment = var.deployment_name
+  }
+}
+
+resource "aws_kms_alias" "braintrust" {
+  name          = "alias/braintrust/${var.deployment_name}"
+  target_key_id = aws_kms_key.braintrust.key_id
+}
+
+data "aws_caller_identity" "current" {}

modules/kms/outputs.tf

@@ -0,0 +1,9 @@
+output "key_arn" {
+  description = "The ARN of the KMS key"
+  value       = aws_kms_key.braintrust.arn
+}
+
+output "key_alias" {
+  description = "The alias of the KMS key"
+  value       = aws_kms_alias.braintrust.name
+}

modules/kms/variables.tf

@@ -0,0 +1,4 @@
+variable "deployment_name" {
+  description = "Name of the deployment, used for resource naming"
+  type        = string
+}

modules/kms/versions.tf

@@ -0,0 +1,9 @@
+terraform {
+  required_version = ">= 1.9.0"
+  required_providers {
+    aws = {
+      source  = "hashicorp/aws"
+      version = "~> 5.0"
+    }
+  }
+}

modules/services/lambda-aiproxy.tf

@@ -1,5 +1,8 @@
+locals {
+  ai_proxy_function_name = "${var.deployment_name}-AIProxy"
+}
 resource "aws_lambda_function" "ai_proxy" {
-  function_name = "${var.deployment_name}-AIProxy"
+  function_name = local.ai_proxy_function_name
   s3_bucket     = local.lambda_s3_bucket
   s3_key        = local.lambda_versions["AIProxy"]
   role          = aws_iam_role.api_handler_role.arn
@@ -8,6 +11,12 @@ resource "aws_lambda_function" "ai_proxy" {
   memory_size   = 1024
   timeout       = 900
   publish       = true
+  kms_key_arn   = var.kms_key_arn
+
+  logging_config {
+    log_format = "Text"
+    log_group  = "/braintrust/${var.deployment_name}/${local.ai_proxy_function_name}"
+  }
 
   ephemeral_storage {
     size = 1024