feat: Added private-only traffic feature

main.tf

@@ -165,27 +165,31 @@ module "app_lb" {
 
   fqdn                      = local.full_fqdn
   extra_fqdn                = local.extra_fqdn
-  allowed_inbound_cidr      = var.allowed_inbound_cidr
-  allowed_inbound_ipv6_cidr = var.allowed_inbound_ipv6_cidr
-  target_port               = local.internal_app_port
+  allowed_inbound_cidr        = var.allowed_inbound_cidr
+  allowed_inbound_ipv6_cidr   = var.allowed_inbound_ipv6_cidr
+  target_port                 = local.internal_app_port
+  network_id                  = local.network_id
+  network_private_subnets     = local.network_private_subnets
+  network_public_subnets      = local.network_public_subnets
+  enable_private_only_traffic = var.private_only_traffic
+  private_endpoint_cidr       = var.allowed_private_endpoint_cidr
 
-  network_id              = local.network_id
-  network_private_subnets = local.network_private_subnets
-  network_public_subnets  = local.network_public_subnets
 }
 
 module "private_link" {
   count  = length(var.private_link_allowed_account_ids) > 0 ? 1 : 0
   source = "./modules/private_link"
 
-  namespace               = var.namespace
-  allowed_account_ids     = var.private_link_allowed_account_ids
-  deletion_protection     = var.deletion_protection
-  network_private_subnets = local.network_private_subnets
-  alb_name                = local.lb_name_truncated
-  vpc_id                  = local.network_id
-
+  namespace                   = var.namespace
+  allowed_account_ids         = var.private_link_allowed_account_ids
+  deletion_protection         = var.deletion_protection
+  network_private_subnets     = local.network_private_subnets
+  alb_name                    = local.lb_name_truncated
+  vpc_id                      = local.network_id
+  enable_private_only_traffic = var.private_only_traffic
+  nlb_security_group          = module.app_lb.nlb_security_group
   depends_on = [
+    module.app_lb,
     module.wandb
   ]
 }

modules/app_lb/main.tf

@@ -3,6 +3,36 @@ locals {
   https_port = 443
 }
 
+resource "aws_security_group" "inbound-private" {
+  count       = var.enable_private_only_traffic ? 1 : 0
+  name        = "${var.namespace}-nlb-inbound-private-endpoint"
+  description = "Allow http(s) traffic to wandb"
+  vpc_id      = var.network_id
+
+  dynamic "ingress" {
+    for_each = var.private_endpoint_cidr
+    content {
+      from_port   = local.https_port
+      to_port     = local.https_port
+      protocol    = "tcp"
+      description = "Allow HTTPS (port ${local.https_port}) traffic inbound to W&B LB"
+      cidr_blocks = [ingress.value]
+    }
+  }
+
+  dynamic "ingress" {
+    for_each = var.private_endpoint_cidr
+    content {
+      from_port   = local.http_port
+      to_port     = local.http_port
+      protocol    = "tcp"
+      description = "Allow HTTP (port ${local.http_port}) traffic inbound to W&B LB"
+      cidr_blocks = [ingress.value]
+    }
+  }
+}
+
+
 resource "aws_security_group" "inbound" {
   name        = "${var.namespace}-alb-inbound"
   description = "Allow http(s) traffic to wandb"
@@ -27,6 +57,27 @@ resource "aws_security_group" "inbound" {
   }
 }
 
+resource "aws_security_group_rule" "alb_http_traffic" {
+  count                    = var.enable_private_only_traffic ? 1 : 0
+  type                     = "ingress"
+  from_port                = local.http_port
+  to_port                  = local.http_port
+  protocol                 = "tcp"
+  security_group_id        = aws_security_group.inbound.id
+  source_security_group_id = aws_security_group.inbound-private[0].id
+}
+
+resource "aws_security_group_rule" "alb_https_traffic" {
+  count                    = var.enable_private_only_traffic ? 1 : 0
+  type                     = "ingress"
+  from_port                = local.https_port
+  to_port                  = local.https_port
+  protocol                 = "tcp"
+  security_group_id        = aws_security_group.inbound.id
+  source_security_group_id = aws_security_group.inbound-private[0].id
+}
+
+
 resource "aws_security_group" "outbound" {
   name        = "${var.namespace}-alb-outbound"
   vpc_id      = var.network_id

modules/app_lb/outputs.tf

@@ -12,4 +12,12 @@ output "lb_arn" {
 
 output "tg_app_arn" {
   value = aws_lb_target_group.app.arn
+}
+
+output "alb_name" {
+value =  aws_lb.alb.arn
+}
+
+output "nlb_security_group" {
+  value = var.enable_private_only_traffic? aws_security_group.inbound-private[0].id : null
 }

modules/app_lb/variables.tf

@@ -73,4 +73,15 @@ variable "network_public_subnets" {
 variable "target_port" {
   type    = number
   default = 32543
+}
+
+
+variable "private_endpoint_cidr" {
+  description = "List of CIDR blocks allowed to access the wandb-server"
+  type        = list(string)
+}
+
+variable "enable_private_only_traffic" {
+  description = "Boolean flag to create sg"
+  type        = bool
 }

modules/private_link/main.tf

@@ -1,6 +1,6 @@
 locals {
   max_lb_name_length = 32 - length("-nlb")
-  lb_name_truncated  = "${substr(var.namespace, 0, local.max_lb_name_length)}-nlb"
+  lb_name_truncated  = var.enable_private_only_traffic ?  "${substr(var.namespace, 0, local.max_lb_name_length)}-private-link-nlb" : "${substr(var.namespace, 0, local.max_lb_name_length)}-nlb"
 }
 
 resource "aws_lb" "nlb" {
@@ -9,6 +9,10 @@ resource "aws_lb" "nlb" {
   load_balancer_type         = "network"
   subnets                    = var.network_private_subnets
   enable_deletion_protection = var.deletion_protection
+  security_groups = var.enable_private_only_traffic ? [var.nlb_security_group] : []
+lifecycle {
+  create_before_destroy = true
+}
 }
 
 resource "aws_lb_target_group" "nlb" {

modules/private_link/variables.tf

@@ -27,3 +27,10 @@ variable "vpc_id" {
   description = "ID of the VPC to create the VPC Endpoint Service in"
   type        = string
 }
+
+variable "enable_private_only_traffic" {
+  type        = bool
+}
+variable "nlb_security_group" {
+  type        = string
+}

variables.tf

@@ -267,6 +267,19 @@ variable "private_link_allowed_account_ids" {
   default     = []
 }
 
+variable "allowed_private_endpoint_cidr" {
+  description = "Private CIDRs allowed to access wandb-server."
+  nullable    = false
+  type        = list(string)
+  default = ["10.20.0.0/20"]
+}
+
+variable "private_only_traffic" {
+  description = "Enable private only traffic from customer private network"
+  type = bool
+  default = true
+}
+
 ##########################################
 # EKS Cluster                            #
 ##########################################