feat: update congnito auth

Author: CommanderK5

main.tf

@@ -63,6 +63,8 @@ module "metaflow-ui" {
   ui_static_container_image       = local.ui_static_container_image
   alb_internal                    = var.ui_alb_internal
   ui_allow_list                   = var.ui_allow_list
+  authenticate_with_cognito       = var.authenticate_with_cognito
+  cognito                         = var.cognito
 
   METAFLOW_DATASTORE_SYSROOT_S3      = module.metaflow-datastore.METAFLOW_DATASTORE_SYSROOT_S3
   certificate_arn                    = var.ui_certificate_arn

modules/ui/ec2.tf

@@ -121,21 +121,26 @@ resource "aws_lb_listener" "this" {
   certificate_arn = var.certificate_arn
 
   dynamic "default_action" {
-    for_each = local.authenticate_cognito
+    for_each = local.default_actions
+
     content {
-      type = "authenticate-cognito"
-      authenticate_cognito {
-        user_pool_arn       = lookup(authenticate_cognito.value, "user_pool_arn", null)
-        user_pool_client_id = lookup(authenticate_cognito.value, "user_pool_client_id", null)
-        user_pool_domain    = lookup(authenticate_cognito.value, "user_pool_domain", null)
+      type = lookup(default_action.value, "type", null)
+
+      dynamic "authenticate_cognito" {
+        for_each = length(keys(lookup(default_action.value, "authenticate_cognito", {}))) > 0 ? [lookup(default_action.value, "authenticate_cognito", {})] : []
+
+        content {
+          user_pool_arn       = lookup(authenticate_cognito.value, "user_pool_arn", null)
+          user_pool_client_id = lookup(authenticate_cognito.value, "user_pool_client_id", null)
+          user_pool_domain    = lookup(authenticate_cognito.value, "user_pool_domain", null)
+        }
       }
     }
   }
 
   default_action {
     type             = "forward"
     target_group_arn = aws_lb_target_group.ui_static.id
-    order            = 100
   }
 }
 
@@ -144,13 +149,19 @@ resource "aws_lb_listener_rule" "ui_backend" {
   priority     = 1
 
   dynamic "action" {
-    for_each = local.authenticate_cognito
+    for_each = local.default_actions
+
     content {
-      type = "authenticate-cognito"
-      authenticate_cognito {
-        user_pool_arn       = lookup(authenticate_cognito.value, "user_pool_arn", null)
-        user_pool_client_id = lookup(authenticate_cognito.value, "user_pool_client_id", null)
-        user_pool_domain    = lookup(authenticate_cognito.value, "user_pool_domain", null)
+      type = lookup(action.value, "type", null)
+
+      dynamic "authenticate_cognito" {
+        for_each = length(keys(lookup(action.value, "authenticate_cognito", {}))) > 0 ? [lookup(action.value, "authenticate_cognito", {})] : []
+
+        content {
+          user_pool_arn       = lookup(authenticate_cognito.value, "user_pool_arn", null)
+          user_pool_client_id = lookup(authenticate_cognito.value, "user_pool_client_id", null)
+          user_pool_domain    = lookup(authenticate_cognito.value, "user_pool_domain", null)
+        }
       }
     }
   }

modules/ui/locals.tf

@@ -38,7 +38,7 @@ locals {
     var.ui_backend_container_image
   )
 
-  authenticate_cognito = var.authenticate_with_cognito ? [
+  default_actions = var.authenticate_with_cognito ? [
     {
       type             = "authenticate-cognito"
       target_group_arn = null