AWS Networking Setup

# This script is provided as a sample network deployment for Poolside.
# Before use, it should be thoroughly reviewed by your security and networking teams
# to ensure compliance with your organization's policies and best practices.
# Modify as needed to meet your specific requirements and security standards.

variable "resource_prefix" {
  description = "Prefix to be used for all resource names and tags"
  type        = string
}

variable "region" {
  description = "AWS region to deploy resources"
  type        = string
  default     = "us-east-1"
  # Please ensure that you select a region with available GPU capacity for Poolside.
}

variable "vpc_cidr" {
  description = "CIDR block for the VPC"
  type        = string
  default     = "10.0.0.0/16"
}

variable "availability_zones" {
  description = "List of AZs to use for the subnets"
  type        = list(string)
  default     = ["us-east-1a", "us-east-1b"] # Specify your preferred AZs here, ensuring that your GPUs are in at least one of the AZs
}

variable "create_nat_gateway" {
  description = "Whether to create a NAT Gateway for private subnets"
  type        = bool
  default     = false
  # Note: A NAT Gateway is not strictly required by Poolside.
  # Its creation is at the discretion of the customer based on their network requirements.
}

provider "aws" {
  region = var.region
}

resource "aws_vpc" "vpc" {
  cidr_block           = var.vpc_cidr
  enable_dns_hostnames = true
  enable_dns_support   = true

  tags = {
    Name = "${var.resource_prefix}_vpc"
  }
}

resource "aws_internet_gateway" "igw" {
  vpc_id = aws_vpc.vpc.id

  tags = {
    Name = "${var.resource_prefix}_igw"
  }
}

resource "aws_subnet" "public" {
  count                   = 1
  vpc_id                  = aws_vpc.vpc.id
  cidr_block              = cidrsubnet(var.vpc_cidr, 8, 0)
  availability_zone       = var.availability_zones[count.index]
  map_public_ip_on_launch = true

  tags = {
    Name = "${var.resource_prefix}_public_subnet_${var.availability_zones[count.index]}"
    Type = "Public"
  }
}

resource "aws_subnet" "private_worker" {
  count             = 2
  vpc_id            = aws_vpc.vpc.id
  cidr_block        = cidrsubnet(var.vpc_cidr, 8, count.index + 1)
  availability_zone = var.availability_zones[count.index]

  tags = {
    Name = "${var.resource_prefix}_private_worker_subnet_${var.availability_zones[count.index]}"
    Type = "Private"
  }
}

resource "aws_subnet" "private_control_plane" {
  count             = 2
  vpc_id            = aws_vpc.vpc.id
  cidr_block        = cidrsubnet(var.vpc_cidr, 8, count.index + 3)
  availability_zone = var.availability_zones[count.index]

  tags = {
    Name = "${var.resource_prefix}_private_control_plane_subnet_${var.availability_zones[count.index]}"
    Type = "Private"
  }
}

resource "aws_eip" "nat" {
  count  = var.create_nat_gateway ? 1 : 0
  domain = "vpc"

  tags = {
    Name = "${var.resource_prefix}_nat_eip"
  }
}

resource "aws_nat_gateway" "nat" {
  count         = var.create_nat_gateway ? 1 : 0
  allocation_id = aws_eip.nat[0].id
  subnet_id     = aws_subnet.public[0].id

  tags = {
    Name = "${var.resource_prefix}_nat_gateway"
  }
}

resource "aws_route_table" "public" {
  vpc_id = aws_vpc.vpc.id

  route {
    cidr_block = "0.0.0.0/0"
    gateway_id = aws_internet_gateway.igw.id
  }

  tags = {
    Name = "${var.resource_prefix}_public_route_table"
  }
}

resource "aws_route_table" "private" {
  vpc_id = aws_vpc.vpc.id

  dynamic "route" {
    for_each = var.create_nat_gateway ? [1] : []
    content {
      cidr_block     = "0.0.0.0/0"
      nat_gateway_id = aws_nat_gateway.nat[0].id
    }
  }

  tags = {
    Name = "${var.resource_prefix}_private_route_table"
  }
}

resource "aws_route_table_association" "public" {
  count          = 1
  subnet_id      = aws_subnet.public[count.index].id
  route_table_id = aws_route_table.public.id
}

resource "aws_route_table_association" "private_worker" {
  count          = 2
  subnet_id      = aws_subnet.private_worker[count.index].id
  route_table_id = aws_route_table.private.id
}

resource "aws_route_table_association" "private_control_plane" {
  count          = 2
  subnet_id      = aws_subnet.private_control_plane[count.index].id
  route_table_id = aws_route_table.private.id
}

# Note: The following Network ACLs are set to allow all traffic by default.
# It is highly recommended to review and modify these rules to only allow
# necessary traffic based on your security requirements.
resource "aws_network_acl" "public" {
  vpc_id     = aws_vpc.vpc.id
  subnet_ids = aws_subnet.public[*].id

  egress {
    protocol   = "-1"
    rule_no    = 100
    action     = "allow"
    cidr_block = "0.0.0.0/0"
    from_port  = 0
    to_port    = 0
  }

  ingress {
    protocol   = "-1"
    rule_no    = 100
    action     = "allow"
    cidr_block = "0.0.0.0/0"
    from_port  = 0
    to_port    = 0
  }

  tags = {
    Name = "${var.resource_prefix}_public_network_acl"
  }
}

resource "aws_network_acl" "private" {
  vpc_id = aws_vpc.vpc.id
  subnet_ids = concat(
    aws_subnet.private_worker[*].id,
    aws_subnet.private_control_plane[*].id
  )

  egress {
    protocol   = "-1"
    rule_no    = 100
    action     = "allow"
    cidr_block = "0.0.0.0/0"
    from_port  = 0
    to_port    = 0
  }

  ingress {
    protocol   = "-1"
    rule_no    = 100
    action     = "allow"
    cidr_block = "0.0.0.0/0"
    from_port  = 0
    to_port    = 0
  }

  tags = {
    Name = "${var.resource_prefix}_private_network_acl"
  }
}

output "vpc_id" {
  value = aws_vpc.vpc.id
}

output "private_subnet_ids" {
  value = aws_subnet.private_worker[*].id
}

output "control_plane_subnet_ids" {
  value = aws_subnet.private_control_plane[*].id
}

output "public_subnet_ids" {
  value = aws_subnet.public[*].id
}