Merge pull request #12 from dnks0/feature/dbx-proxy

added support for existing Security groups of an NLB in proxy-only mode

Author: dnks0

terraform/aws/main.tf

@@ -36,6 +36,7 @@ module "load_balancer" {
 
   vpc_id                  = local.vpc_id
   subnet_ids              = local.subnet_ids
+  subnet_cidrs            = local.subnet_cidrs
 
   dbx_proxy_health_port   = var.dbx_proxy_health_port
   dbx_proxy_listener      = var.dbx_proxy_listener

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

@@ -4,6 +4,24 @@ 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 = local.nlb_has_security_groups ? data.aws_lb.this[0].security_groups : []
+
+  nlb_listener_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 = "NLB to dbx-proxy health checks" }],
+  )
+
+  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) : {
+      security_group_id = pair[0]
+      description       = pair[1].description
+      port              = pair[1].port
+      cidr              = pair[2]
+    }
+  ] : []
+
   allowed_principals = [
     "arn:aws:iam::565502421330:role/private-connectivity-role-${var.region}"
   ]

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

@@ -13,6 +13,17 @@ resource "aws_lb" "this" {
   tags = var.tags
 }
 
+resource "aws_vpc_security_group_egress_rule" "this" {
+  count = local.nlb_has_security_groups ? length(local.nlb_sg_egress_rules) : 0
+
+  security_group_id = local.nlb_sg_egress_rules[count.index].security_group_id
+  from_port         = local.nlb_sg_egress_rules[count.index].port
+  to_port           = local.nlb_sg_egress_rules[count.index].port
+  ip_protocol       = "tcp"
+  cidr_ipv4         = local.nlb_sg_egress_rules[count.index].cidr
+  description       = local.nlb_sg_egress_rules[count.index].description
+}
+
 # Optional: expose the dbx-proxy health port via the NLB so callers can reach it directly
 # (e.g. through the PrivateLink endpoint). If the health port is already used as a regular
 # listener port, we skip creating this additional listener/TG to avoid a conflict.

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

@@ -33,6 +33,11 @@ variable "subnet_ids" {
   type        = list(string)
 }
 
+variable "subnet_cidrs" {
+  description = "Subnet CIDR blocks used for NLB security group egress rules."
+  type        = list(string)
+}
+
 variable "dbx_proxy_health_port" {
   description = "Port on which the dbx-proxy instances expose a TCP health check."
   type        = number