added security group to load-balancer

Author: dnks0

terraform/aws/modules/load-balancer/local.tf

@@ -4,17 +4,16 @@ locals {
   nlb_dns_name = var.bootstrap_load_balancer ? aws_lb.this[0].dns_name : data.aws_lb.this[0].dns_name
   nlb_zone_id  = var.bootstrap_load_balancer ? aws_lb.this[0].zone_id : data.aws_lb.this[0].zone_id
 
-  nlb_has_security_groups = var.bootstrap_load_balancer ? false : length(data.aws_lb.this[0].security_groups) > 0
+  nlb_security_group_ids = var.bootstrap_load_balancer ? aws_lb.this[0].security_groups : data.aws_lb.this[0].security_groups
 
-  nlb_security_group_ids = local.nlb_has_security_groups ? data.aws_lb.this[0].security_groups : []
 
-  nlb_listener_for_egress_rules = concat(
+  nlb_ports_for_egress_rules = concat(
     [for l in var.dbx_proxy_listener : { port = l.port, description = "Databricks to NLB to dbx-proxy listener ${l.name}" }],
-    contains([for l in var.dbx_proxy_listener : l.port], var.dbx_proxy_health_port) ? [] : [{ port = var.dbx_proxy_health_port, description = "Databricks to NLB to dbx-proxy health check" }],
+    [{ port = var.dbx_proxy_health_port, description = "Databricks to NLB to dbx-proxy health check" }],
   )
 
-  nlb_sg_egress_rules = local.nlb_has_security_groups ? [
-    for pair in setproduct(local.nlb_security_group_ids, local.nlb_listener_for_egress_rules, var.subnet_cidrs) : {
+  nlb_sg_egress_rules = length(local.nlb_security_group_ids) > 0 ? [
+    for pair in setproduct(local.nlb_security_group_ids, local.nlb_ports_for_egress_rules, var.subnet_cidrs) : {
       security_group_id = pair[0]
       description       = pair[1].description
       port              = pair[1].port

terraform/aws/modules/load-balancer/nlb.tf

@@ -6,15 +6,52 @@ resource "aws_lb" "this" {
   load_balancer_type = "network"
   internal           = true
   subnets            = var.subnet_ids
+  security_groups    = [aws_security_group.this[0].id]
 
   enable_cross_zone_load_balancing = true
   enable_deletion_protection = false
 
   tags = var.tags
 }
 
+resource "aws_security_group" "this" {
+  count = var.bootstrap_load_balancer ? 1 : 0
+
+  name        = "${var.prefix}-nlb-sg"
+  description = "Security group for dbx-proxy NLB"
+  vpc_id      = var.vpc_id
+
+  # Outbound from NLB on any listener port
+  dynamic "egress" {
+    for_each = { for l in var.dbx_proxy_listener : l.name => l }
+    content {
+      description = "Databricks to NLB to dbx-proxy listener ${egress.key}"
+      from_port   = egress.value.port
+      to_port     = egress.value.port
+      protocol    = "tcp"
+      cidr_blocks = var.subnet_cidrs
+    }
+  }
+
+  # Health check port
+  egress {
+    description = "Databricks to NLB to dbx-proxy health check"
+    from_port   = var.dbx_proxy_health_port
+    to_port     = var.dbx_proxy_health_port
+    protocol    = "tcp"
+    cidr_blocks = var.subnet_cidrs
+  }
+
+  tags = merge(
+    var.tags,
+    {
+      Name = "${var.prefix}-nlb-sg"
+    },
+  )
+}
+
 resource "aws_vpc_security_group_egress_rule" "this" {
-  count = local.nlb_has_security_groups ? length(local.nlb_sg_egress_rules) : 0
+  count = var.bootstrap_load_balancer ? 0 : length(local.nlb_sg_egress_rules)
 
   security_group_id = local.nlb_sg_egress_rules[count.index].security_group_id
   from_port         = local.nlb_sg_egress_rules[count.index].port

terraform/aws/modules/load-balancer/outputs.tf

@@ -16,6 +16,11 @@ output "nlb_target_group_arns" {
   )
 }
 
+output "nlb_security_group_ids" {
+  description = "Security group IDs attached to the NLB (if any)."
+  value       = tolist(local.nlb_security_group_ids)
+}
+
 output "vpc_endpoint_service_arn" {
   description = "ARN of the VPC endpoint service if created; otherwise null."
   value       = length(aws_vpc_endpoint_service.this) > 0 ? aws_vpc_endpoint_service.this[0].arn : null

terraform/aws/outputs.tf

@@ -17,6 +17,7 @@ output "load_balancer" {
     nlb_arn                   = module.load_balancer.nlb_arn
     nlb_dns_name              = module.load_balancer.nlb_dns_name
     nlb_target_group_arns     = module.load_balancer.nlb_target_group_arns
+    nlb_security_group_ids    = module.load_balancer.nlb_security_group_ids
     vpc_endpoint_service_arn  = module.load_balancer.vpc_endpoint_service_arn
     vpc_endpoint_service_name = module.load_balancer.vpc_endpoint_service_name
   }