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

* added validation to avoid conflicts where the proxy health port is also used as a listener port
* streamlined bootstrap deployments to also include a security group on the NLB
   * ingress is bypassed on the NLB SG due to limitations
   * ingress to the NLB is controlled via the endpoint-service
* included NLB security groups into module outputs

Author: dnks0

README.md

@@ -86,4 +86,4 @@ curl -sS -w '\nHTTP %{http_code}\n' http://<ncc-endpoint-rule-domain>:8080/statu
 ```
 
 ### Limitations / Trade-Offs
-Before going to production, please review the following [limitations & trade-offs](terraform/README.md#limitations--tradeoffs-of-the-current-implementation).
+Before going to production, please review the following [limitations & trade-offs](terraform/README.md#limitations--tradeoffs-of-the-current-implementation).

terraform/README.md

@@ -283,3 +283,9 @@ This module is intentionally minimal right now. The following limitations are im
   - Databricks enforces limits around NCCs, private endpoints, and private endpoint rules (including limits on the number of domain names per rule).
   - Treat these as **external constraints** that influence how you model `dbx_proxy_listener`.
   - Reference: [Configure private connectivity to resources in your VPC](https://docs.databricks.com/aws/en/security/network/serverless-network-security/pl-to-internal-network).
+
+- **PrivateLink NLB SG ingress enforcement**
+  - Currently, if bootstrapped, the NLB gets created with a dedicated Security Group attached, which would allow to control ingress to the NLB, enforced with `enforce_security_group_inbound_rules_on_private_link_traffic = on`.
+  - The above setting would also allow to see individual client-IPs as source in network packets (which would be favorable in general)
+  - However, we are not able to lock-down the inbound rules on the NLB Security Group due to missing information like Security Group ID or VPC endpoint ID of the consumer side (Databricks Serverless).
+  - Therefore, we are using `enforce_security_group_inbound_rules_on_private_link_traffic = on`, which bypasses the Security Group inbound rules, and solely rely on the endpoint service `allowed_principals` and manual acceptance to control inbound traffic to the NLB.

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,63 @@ 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
 
+  # PrivateLink traffic bypasses NLB SG ingress when this is off. We use off
+  # because the service owner cannot restrict ingress by endpoint SG/CIDR without
+  # consumer-provided details; access is instead controlled via endpoint service
+  # allowed_principals and manual acceptance.
+  enforce_security_group_inbound_rules_on_private_link_traffic = "off"
+
   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
+  }
+
+  # We can not lock down ingress to individual sources since we don't know the
+  # source IP addresses or Security Group IDs the traffic is originating from.
+  # Therefore, we allow all ingress traffic on the SG level. Ingress is controlled
+  # by the endpoint service allowed_principals and manual acceptance.
+
+  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
@@ -24,12 +72,7 @@ resource "aws_vpc_security_group_egress_rule" "this" {
   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.
 resource "aws_lb_target_group" "health" {
-  count = contains([for l in var.dbx_proxy_listener : l.port], var.dbx_proxy_health_port) ? 0 : 1
-
   name        = "dbx-proxy-tg-health"
   port        = var.dbx_proxy_health_port
   protocol    = "TCP"
@@ -49,16 +92,21 @@ resource "aws_lb_target_group" "health" {
 }
 
 resource "aws_lb_listener" "health" {
-  count = length(aws_lb_target_group.health)
-
   load_balancer_arn = local.nlb_arn
   port              = var.dbx_proxy_health_port
   protocol          = "TCP"
 
   default_action {
     type             = "forward"
-    target_group_arn = aws_lb_target_group.health[0].arn
+    target_group_arn = aws_lb_target_group.health.arn
   }
+
+  tags = merge(
+    var.tags,
+    {
+      Name = "${var.prefix}-l-health"
+    },
+  )
 }
 
 # One target group per listener port for simple configuration.
@@ -94,4 +142,11 @@ resource "aws_lb_listener" "this" {
     type             = "forward"
     target_group_arn = each.value.arn
   }
+
+  tags = merge(
+    var.tags,
+    {
+      Name = "${var.prefix}-l-${each.key}"
+    },
+  )
 }

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

@@ -12,10 +12,15 @@ output "nlb_target_group_arns" {
   description = "ARNs of the NLB target groups, keyed by listener name (plus health when created)."
   value = merge(
     { for name, tg in aws_lb_target_group.this : name => tg.arn },
-    length(aws_lb_target_group.health) > 0 ? { health = aws_lb_target_group.health[0].arn } : {},
+    { health = aws_lb_target_group.health.arn },
   )
 }
 
+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
   }

terraform/aws/variables.tf

@@ -145,4 +145,8 @@ EOT
     }))
   }))
   default = []
+  validation {
+    condition     = alltrue([for listener in var.dbx_proxy_listener : listener.port != var.dbx_proxy_health_port])
+    error_message = "dbx_proxy_health_port must not overlap with any dbx_proxy_listener port."
+  }
 }