Use individual SGs instead of Default SG (#77)

When deploying the Braintrust data plane, we will no longer use the
default VPC security group, but instead create individual security
groups for Brainstore, elasticache, rds, and lambdas.

Existing deployments will have all resources deployed from Terraform
switch from using the default security groups to the individual security
groups. The quarantine lambdas that are created outside of Terraform
will still have the default VPC Security group of the Quarantine VPC,
however these can be changed to use the new Quarantine Security Group
instead via the AWS Console or via script.

---------

Co-authored-by: Jeff McCollum <[email protected]>
Co-authored-by: Mike Deeks <[email protected]>

Author: jeffmccollum

main.tf

@@ -11,6 +11,9 @@ locals {
   clickhouse_address = var.use_external_clickhouse_address != null ? var.use_external_clickhouse_address : (
     var.enable_clickhouse ? module.clickhouse[0].clickhouse_instance_private_ip : null
   )
+  bastion_security_group = var.enable_braintrust_support_shell_access ? {
+    "Remote Support Bastion" = module.remote_support[0].remote_support_security_group_id
+  } : {}
 }
 
 module "main_vpc" {
@@ -62,8 +65,14 @@ module "database" {
     module.main_vpc.private_subnet_2_id,
     module.main_vpc.private_subnet_3_id
   ]
-  database_security_group_ids = [module.main_vpc.default_security_group_id]
-
+  vpc_id = module.main_vpc.vpc_id
+  authorized_security_groups = merge(
+    {
+      "Lambda Services" = module.services.lambda_security_group_id
+      "Brainstore"      = var.enable_brainstore ? module.brainstore[0].brainstore_instance_security_group_id : null
+    },
+    local.bastion_security_group,
+  )
   postgres_storage_iops       = var.postgres_storage_iops
   postgres_storage_throughput = var.postgres_storage_throughput
   auto_minor_version_upgrade  = var.postgres_auto_minor_version_upgrade
@@ -80,7 +89,14 @@ module "redis" {
     module.main_vpc.private_subnet_2_id,
     module.main_vpc.private_subnet_3_id
   ]
-  security_group_ids  = [module.main_vpc.default_security_group_id]
+  vpc_id = module.main_vpc.vpc_id
+  authorized_security_groups = merge(
+    {
+      "Lambda Services" = module.services.lambda_security_group_id
+      "Brainstore"      = var.enable_brainstore ? module.brainstore[0].brainstore_instance_security_group_id : null
+    },
+    local.bastion_security_group,
+  )
   redis_instance_type = var.redis_instance_type
   redis_version       = var.redis_version
 }
@@ -127,17 +143,16 @@ module "services" {
   extra_env_vars                             = var.service_extra_env_vars
 
   # Networking
-  service_security_group_ids = [module.main_vpc.default_security_group_id]
+  vpc_id = module.main_vpc.vpc_id
   service_subnet_ids = [
     module.main_vpc.private_subnet_1_id,
     module.main_vpc.private_subnet_2_id,
     module.main_vpc.private_subnet_3_id
   ]
 
   # Quarantine VPC
-  use_quarantine_vpc                       = var.enable_quarantine_vpc
-  quarantine_vpc_id                        = var.enable_quarantine_vpc ? module.quarantine_vpc[0].vpc_id : null
-  quarantine_vpc_default_security_group_id = var.enable_quarantine_vpc ? module.quarantine_vpc[0].default_security_group_id : null
+  use_quarantine_vpc = var.enable_quarantine_vpc
+  quarantine_vpc_id  = var.enable_quarantine_vpc ? module.quarantine_vpc[0].vpc_id : null
   quarantine_vpc_private_subnets = var.enable_quarantine_vpc ? [
     module.quarantine_vpc[0].private_subnet_1_id,
     module.quarantine_vpc[0].private_subnet_2_id,
@@ -189,8 +204,15 @@ module "brainstore" {
   internal_observability_env_name = var.internal_observability_env_name
   internal_observability_region   = var.internal_observability_region
 
-  vpc_id            = module.main_vpc.vpc_id
-  security_group_id = module.main_vpc.default_security_group_id
+  vpc_id = module.main_vpc.vpc_id
+  authorized_security_groups = merge(
+    {
+      "Lambda Services" = module.services.lambda_security_group_id
+    },
+    local.bastion_security_group
+  )
+  authorized_security_groups_ssh = local.bastion_security_group
+
   private_subnet_ids = [
     module.main_vpc.private_subnet_1_id,
     module.main_vpc.private_subnet_2_id,

module-docs.md

@@ -1,3 +1,4 @@
+<!-- BEGIN_TF_DOCS -->
 ## Required Inputs
 
 The following input variables are required:
@@ -44,22 +45,6 @@ Type: `number`
 
 Default: `-1`
 
-### <a name="input_brainstore_backfill_disable_historical"></a> [brainstore\_backfill\_disable\_historical](#input\_brainstore\_backfill\_disable\_historical)
-
-Description: Disable historical backfill for Brainstore. Don't modify this unless instructed by Braintrust.
-
-Type: `bool`
-
-Default: `false`
-
-### <a name="input_brainstore_backfill_disable_nonhistorical"></a> [brainstore\_backfill\_disable\_nonhistorical](#input\_brainstore\_backfill\_disable\_nonhistorical)
-
-Description: Disable non-historical backfill for Brainstore. Don't modify this unless instructed by Braintrust.
-
-Type: `bool`
-
-Default: `false`
-
 ### <a name="input_brainstore_backfill_new_objects"></a> [brainstore\_backfill\_new\_objects](#input\_brainstore\_backfill\_new\_objects)
 
 Description: Enable backfill for new objects for Brainstore. Don't modify this unless instructed by Braintrust.
@@ -74,11 +59,11 @@ Description: Whether to set Brainstore as the default rather than requiring user
 
 Type: `string`
 
-Default: `"true"`
+Default: `"force"`
 
 ### <a name="input_brainstore_disable_optimization_worker"></a> [brainstore\_disable\_optimization\_worker](#input\_brainstore\_disable\_optimization\_worker)
 
-Description: Disable the optimization worker in Brainstore
+Description: Disable the optimization worker globally in Brainstore
 
 Type: `bool`
 
@@ -92,14 +77,6 @@ Type: `bool`
 
 Default: `true`
 
-### <a name="input_brainstore_enable_index_validation"></a> [brainstore\_enable\_index\_validation](#input\_brainstore\_enable\_index\_validation)
-
-Description: Enable index validation for Brainstore
-
-Type: `bool`
-
-Default: `false`
-
 ### <a name="input_brainstore_etl_batch_size"></a> [brainstore\_etl\_batch\_size](#input\_brainstore\_etl\_batch\_size)
 
 Description: The batch size for the ETL process
@@ -110,19 +87,19 @@ Default: `null`
 
 ### <a name="input_brainstore_extra_env_vars"></a> [brainstore\_extra\_env\_vars](#input\_brainstore\_extra\_env\_vars)
 
-Description: Extra environment variables to set for Brainstore
+Description: Extra environment variables to set for Brainstore reader or dual use nodes
 
 Type: `map(string)`
 
 Default: `{}`
 
-### <a name="input_brainstore_index_validation_only_deletes"></a> [brainstore\_index\_validation\_only\_deletes](#input\_brainstore\_index\_validation\_only\_deletes)
+### <a name="input_brainstore_extra_env_vars_writer"></a> [brainstore\_extra\_env\_vars\_writer](#input\_brainstore\_extra\_env\_vars\_writer)
 
-Description: Scope index validation to only deletes in Brainstore. Only applies if brainstore\_enable\_index\_validation is true
+Description: Extra environment variables to set for Brainstore writer nodes
 
-Type: `bool`
+Type: `map(string)`
 
-Default: `true`
+Default: `{}`
 
 ### <a name="input_brainstore_instance_count"></a> [brainstore\_instance\_count](#input\_brainstore\_instance\_count)
 
@@ -284,6 +261,30 @@ Type: `bool`
 
 Default: `true`
 
+### <a name="input_internal_observability_api_key"></a> [internal\_observability\_api\_key](#input\_internal\_observability\_api\_key)
+
+Description: Support for internal observability agent. Do not set this unless instructed by support.
+
+Type: `string`
+
+Default: `""`
+
+### <a name="input_internal_observability_env_name"></a> [internal\_observability\_env\_name](#input\_internal\_observability\_env\_name)
+
+Description: Support for internal observability agent. Do not set this unless instructed by support.
+
+Type: `string`
+
+Default: `""`
+
+### <a name="input_internal_observability_region"></a> [internal\_observability\_region](#input\_internal\_observability\_region)
+
+Description: Support for internal observability agent. Do not set this unless instructed by support.
+
+Type: `string`
+
+Default: `"us5"`
+
 ### <a name="input_kms_key_arn"></a> [kms\_key\_arn](#input\_kms\_key\_arn)
 
 Description: Existing KMS key ARN to use for encrypting resources. If not provided, a new key will be created. DO NOT change this after deployment. If you do, it will attempt to destroy your DB and prior S3 objects will no longer be readable.
@@ -476,6 +477,14 @@ Type: `string`
 
 Default: `"7.0"`
 
+### <a name="input_s3_additional_allowed_origins"></a> [s3\_additional\_allowed\_origins](#input\_s3\_additional\_allowed\_origins)
+
+Description: Additional origins to allow for S3 bucket CORS configuration. Supports a wildcard in the domain name.
+
+Type: `list(string)`
+
+Default: `[]`
+
 ### <a name="input_service_additional_policy_arns"></a> [service\_additional\_policy\_arns](#input\_service\_additional\_policy\_arns)
 
 Description: Additional policy ARNs to attach to the lambda functions that are the main braintrust service
@@ -548,6 +557,10 @@ Description: The primary endpoint for the dataplane API. This is the value that
 
 Description: Instance ID of the bastion host that Braintrust support staff can connect to using EC2 Instance Connect. Share this with the Braintrust team.
 
+### <a name="output_brainstore_security_group_id"></a> [brainstore\_security\_group\_id](#output\_brainstore\_security\_group\_id)
+
+Description: ID of the security group for the Brainstore instances
+
 ### <a name="output_braintrust_support_role_arn"></a> [braintrust\_support\_role\_arn](#output\_braintrust\_support\_role\_arn)
 
 Description: ARN of the Role that grants Braintrust team remote support. Share this with the Braintrust team.
@@ -564,13 +577,13 @@ Description: Name of the Clickhouse S3 bucket
 
 Description: ID of the Clickhouse secret. Note this is the Terraform ID attribute which is a pipe delimited combination of secret ID and version ID
 
-### <a name="output_main_vpc_cidr"></a> [main\_vpc\_cidr](#output\_main\_vpc\_cidr)
+### <a name="output_lambda_security_group_id"></a> [lambda\_security\_group\_id](#output\_lambda\_security\_group\_id)
 
-Description: CIDR block of the main VPC
+Description: ID of the security group for the Lambda functions
 
-### <a name="output_main_vpc_default_security_group_id"></a> [main\_vpc\_default\_security\_group\_id](#output\_main\_vpc\_default\_security\_group\_id)
+### <a name="output_main_vpc_cidr"></a> [main\_vpc\_cidr](#output\_main\_vpc\_cidr)
 
-Description: ID of the default security group in the main VPC
+Description: CIDR block of the main VPC
 
 ### <a name="output_main_vpc_id"></a> [main\_vpc\_id](#output\_main\_vpc\_id)
 
@@ -608,6 +621,19 @@ Description: ARN of the main Braintrust Postgres database
 
 Description: ID of the quarantine VPC that user functions run inside of.
 
+### <a name="output_rds_security_group_id"></a> [rds\_security\_group\_id](#output\_rds\_security\_group\_id)
+
+Description: ID of the security group for the RDS instance
+
 ### <a name="output_redis_arn"></a> [redis\_arn](#output\_redis\_arn)
 
 Description: ARN of the Redis instance
+
+### <a name="output_redis_security_group_id"></a> [redis\_security\_group\_id](#output\_redis\_security\_group\_id)
+
+Description: ID of the security group for the Elasticache instance
+
+### <a name="output_remote_support_security_group_id"></a> [remote\_support\_security\_group\_id](#output\_remote\_support\_security\_group\_id)
+
+Description: Security Group ID for the Remote Support bastion host.
+<!-- END_TF_DOCS -->

modules/brainstore/main-writer.tf

@@ -11,7 +11,7 @@ resource "aws_launch_template" "brainstore_writer" {
     arn = aws_iam_instance_profile.brainstore.arn
   }
 
-  vpc_security_group_ids = [var.security_group_id]
+  vpc_security_group_ids = [aws_security_group.brainstore_instance.id]
 
   block_device_mappings {
     device_name = "/dev/sda1"
@@ -92,7 +92,7 @@ resource "aws_lb" "brainstore_writer" {
   internal           = true
   load_balancer_type = "network"
   subnets            = var.private_subnet_ids
-  security_groups    = [var.security_group_id]
+  security_groups    = [aws_security_group.brainstore_elb.id]
 
   lifecycle {
     create_before_destroy = true

modules/brainstore/main.tf

@@ -18,7 +18,7 @@ resource "aws_launch_template" "brainstore" {
     arn = aws_iam_instance_profile.brainstore.arn
   }
 
-  vpc_security_group_ids = [var.security_group_id]
+  vpc_security_group_ids = [aws_security_group.brainstore_instance.id]
 
   block_device_mappings {
     device_name = "/dev/sda1"
@@ -99,7 +99,7 @@ resource "aws_lb" "brainstore" {
   internal           = true
   load_balancer_type = "network"
   subnets            = var.private_subnet_ids
-  security_groups    = [var.security_group_id]
+  security_groups    = [aws_security_group.brainstore_elb.id]
 
   lifecycle {
     # Changing security groups requires a new NLB.

modules/brainstore/outputs.tf

@@ -17,3 +17,13 @@ output "port" {
   description = "The port used by Brainstore"
   value       = var.port
 }
+
+output "brainstore_elb_security_group_id" {
+  description = "The ID of the security group for the Brainstore ELB"
+  value       = aws_security_group.brainstore_elb.id
+}
+
+output "brainstore_instance_security_group_id" {
+  description = "The ID of the security group for the Brainstore instances"
+  value       = aws_security_group.brainstore_instance.id
+}

modules/brainstore/security-groups.tf

@@ -0,0 +1,63 @@
+resource "aws_security_group" "brainstore_elb" {
+  name   = "${var.deployment_name}-brainstore-elb"
+  vpc_id = var.vpc_id
+  tags   = merge({ "Name" = "${var.deployment_name}-brainstore-elb" }, local.common_tags)
+}
+
+resource "aws_vpc_security_group_ingress_rule" "brainstore_elb_allow_ingress_from_authorized_security_groups" {
+  for_each = var.authorized_security_groups
+
+  from_port                    = var.port
+  to_port                      = var.port
+  ip_protocol                  = "tcp"
+  referenced_security_group_id = each.value
+  description                  = "Allow inbound to brainstore from ${each.key}."
+  security_group_id            = aws_security_group.brainstore_elb.id
+}
+
+resource "aws_vpc_security_group_egress_rule" "brainstore_elb_allow_egress_all" {
+  from_port         = -1
+  to_port           = -1
+  ip_protocol       = "-1"
+  cidr_ipv4         = "0.0.0.0/0"
+  description       = "Allow all outbound traffic from Brainstore ELB."
+  security_group_id = aws_security_group.brainstore_elb.id
+}
+
+resource "aws_security_group" "brainstore_instance" {
+  name   = "${var.deployment_name}-brainstore-instance"
+  vpc_id = var.vpc_id
+  tags   = merge({ "Name" = "${var.deployment_name}-brainstore-instance" }, local.common_tags)
+}
+
+resource "aws_vpc_security_group_ingress_rule" "brainstore_instance_allow_ingress_from_nlb" {
+
+  from_port                    = var.port
+  to_port                      = var.port
+  ip_protocol                  = "tcp"
+  referenced_security_group_id = aws_security_group.brainstore_elb.id
+  description                  = "Allow inbound to Brainstore instances from NLB."
+  security_group_id            = aws_security_group.brainstore_instance.id
+}
+
+resource "aws_vpc_security_group_ingress_rule" "brainstore_instance_allow_ingress_from_authorized_security_groups_ssh" {
+  for_each = var.authorized_security_groups_ssh
+
+  from_port                    = 22
+  to_port                      = 22
+  ip_protocol                  = "tcp"
+  referenced_security_group_id = each.value
+  description                  = "Allow inbound SSH to Brainstore instances from ${each.key}."
+
+  security_group_id = aws_security_group.brainstore_instance.id
+}
+
+resource "aws_vpc_security_group_egress_rule" "brainstore_instance_allow_egress_all" {
+
+  from_port         = -1
+  to_port           = -1
+  ip_protocol       = "-1"
+  cidr_ipv4         = "0.0.0.0/0"
+  description       = "Allow all outbound traffic from Brainstore instances."
+  security_group_id = aws_security_group.brainstore_instance.id
+}

modules/brainstore/variables.tf

@@ -52,9 +52,16 @@ variable "vpc_id" {
   description = "The ID of the VPC where Brainstore resources will be created"
 }
 
-variable "security_group_id" {
-  type        = string
-  description = "The ID of the security group to use for Brainstore resources"
+variable "authorized_security_groups" {
+  type        = map(string)
+  description = "Map of security group names to their IDs that are authorized to access the Brainstore ELB. Format: { name = <security_group_id> }"
+  default     = {}
+}
+
+variable "authorized_security_groups_ssh" {
+  type        = map(string)
+  description = "Map of security group names to their IDs that are authorized to access Brainstore instances via SSH. Format: { name = <security_group_id> }"
+  default     = {}
 }
 
 variable "private_subnet_ids" {