WAF for all ALBs in an account using terraform

Question

This is my waf.tf terraform file:

resource "aws_wafv2_web_acl" "waf_acl-dev" {
  name        = "waf_log4j_Protections-dev"
  description = "WAFv2 for dev"
  scope       = "REGIONAL"
  default_action {
    allow {}
  }
  rule {
    name     = "AWSManagedRulesKnownBadInputsRule"
    priority = 1

    override_action {
      count {}
    }
    statement {
      managed_rule_group_statement {
        name        = "AWSManagedRulesKnownBadInputsRuleSet"
        vendor_name = "AWS"
        # Excluding all these leaves only Log4JRCE
        excluded_rule {
          name = "Host_localhost_HEADER"
        }
        excluded_rule {
          name = "PROPFIND_METHOD"
        }
        excluded_rule {
          name = "ExploitablePaths_URIPATH"
        }
      }
    }
    visibility_config {
      cloudwatch_metrics_enabled = true
      metric_name                = "AWSManagedRulesKnownBadInputsRule"
      sampled_requests_enabled   = true
    }
  }
  visibility_config {
            cloudwatch_metrics_enabled = true
            metric_name = "AWSManagedKnownBadInputsRule"
            sampled_requests_enabled = true
    }
}

variable "lb_arn" {
  type = string
  default = ""
}

data "aws_lb" "all_alb" {
    

    tags = {
        Environment = "Dev"
    }
    # arn = var.lb_arn
    # for_each = data.aws_lb.all_alb
    # id = each.value
}

output "all_alb" {
    value = data.aws_lb.all_alb
}

resource "aws_wafv2_web_acl_association" "waf_acl-association-dev" {
    for_each = data.aws_lb.all_alb.arn
    resource_arn = each.value
    web_acl_arn = aws_wafv2_web_acl.waf_acl-dev.arn

}

My objective is to create a WAF with rules (that portion works fine), and attach multiple ALBs to the WAF using tags on each ALB. Im running it inside a teamcity loop that loops into multiple AWS accounts (That is out of scope). Each account can have one or multiple ALBs with the tag provided. When I run this code I get the following error:

Error: Search returned 2 results, please revise so only one is returned
17:39:18     │ 
17:39:18     │   with data.aws_lb.all_alb,
17:39:18     │   on xa-waf-inow.tf line 49, in data "aws_lb" "all_alb":
17:39:18     │   49: data "aws_lb" "all_alb" {

I also tried a few other alternatives like putting [0] at the end of line 49 but Im still getting hit with some syntax error or the other. Can someone please help?Thanks

Edit:

I also tried for_each in data block:

data "aws_lb" "all_alb" {
  for_each = {
    tags = {
        Environment = "Dev"
    }
  }
  arn = each.key
 
}
 
output "all_alb" {
    value = data.aws_lb.all_alb
}
 
resource "aws_wafv2_web_acl_association" "waf_acl-association-dev" {
    # for_each = data.aws_lb.all_alb.arn
    resource_arn = data.aws_lb.all_alb[each.key]
    web_acl_arn = aws_wafv2_web_acl.waf_acl-dev.arn

}

But got this error:

Error: Reference to "each" in context without for_each
18:24:01     │ 
18:24:01     │   on xa-waf-inow.tf line 65, in resource "aws_wafv2_web_acl_association" "waf_acl-association-dev":
18:24:01     │   65:     resource_arn = data.aws_lb.all_alb[each.key]
18:24:01     │ 
18:24:01     │ The "each" object can be used only in "module" or "resource" blocks, and
18:24:01     │ only when the "for_each" argument is set.

Answer 1

The aws_lb data source must return only one result. You can't change that, as this is how it was destined.

If you want to return multiple ALBs you have two choices:

1. Use for_each with the data source. This way your data source will run for each id of your alb. This means that you have to provide the alb ids as an input variable.
2. Or, create your own custom data source. Since this is fully custom code that you have to write, it can overcome any limitations of TF's build in data sources.

Answer 2

I was able to get this to work for myself. My issue was a mixture of lists and sets. I believe I can change my variable to a set, and not have to deal with any sets here, but I know the below works as is.

VARIABLES.TF

variable "list_of_alb" {
  type = list(string)
}

MAIN.TF

list_of_alb = [
    "${terraform.workspace}-unique-1",
    "${terraform.workspace}-unique-2"
  ]

DATA.TF - the problem child

data "aws_lb" "main" {
  for_each = toset( var.list_of_alb )
  name = each.value
}

WAF.TF

resource "aws_wafv2_web_acl_association" "main" {
  for_each = toset(var.list_of_alb)
  resource_arn = data.aws_lb.main[each.value].arn
  web_acl_arn  = aws_wafv2_web_acl.main.arn
}