From a572105595df6d8242f0393829c09e80feca8340 Mon Sep 17 00:00:00 2001
From: Alberto Manfrinati <alberto.manfrinati@gmail.com>
Date: Fri, 14 Feb 2025 13:11:06 +0100
Subject: [PATCH] Add block-requests-to-queries, ip-rate-limit and
 block-bot-control-requests rules

---
 .../src/modules/chatbot/variables.tf          |  24 +++++
 .../infrastructure/src/modules/chatbot/waf.tf | 102 +++++++++++++++++-
 2 files changed, 124 insertions(+), 2 deletions(-)

diff --git a/apps/infrastructure/src/modules/chatbot/variables.tf b/apps/infrastructure/src/modules/chatbot/variables.tf
index 1267412b0..7177650ad 100644
--- a/apps/infrastructure/src/modules/chatbot/variables.tf
+++ b/apps/infrastructure/src/modules/chatbot/variables.tf
@@ -127,3 +127,27 @@ variable "ecs_monitoring" {
   })
   description = "Langfuse configuration for the AI chatbot"
 }
+
+################################################################################
+# WAF
+################################################################################
+
+variable "waf_ip_rate_limit_limit" {
+  type    = number
+  default = 100
+}
+
+variable "waf_ip_rate_limit_evaluation_window_sec" {
+  type    = number
+  default = 60
+}
+
+variable "waf_block_requests_to_queries_limit" {
+  type    = number
+  default = 100
+}
+
+variable "waf_block_requests_to_queries_evaluation_window_sec" {
+  type    = number
+  default = 60
+}
diff --git a/apps/infrastructure/src/modules/chatbot/waf.tf b/apps/infrastructure/src/modules/chatbot/waf.tf
index 7a273f517..71b9aa99e 100644
--- a/apps/infrastructure/src/modules/chatbot/waf.tf
+++ b/apps/infrastructure/src/modules/chatbot/waf.tf
@@ -1,5 +1,6 @@
 resource "aws_wafv2_web_acl" "chatbot" {
   name        = "${local.prefix}-waf"
+  description = "Web ACL Rule for Chatbot in ${var.environment}"
   scope       = "REGIONAL"
 
   default_action {
@@ -7,9 +8,9 @@ resource "aws_wafv2_web_acl" "chatbot" {
   }
 
   visibility_config {
-    cloudwatch_metrics_enabled = true
+    cloudwatch_metrics_enabled = false
     metric_name                = "${local.prefix}-waf"
-    sampled_requests_enabled    = true
+    sampled_requests_enabled   = false
   }
 
   association_config {
@@ -19,4 +20,101 @@ resource "aws_wafv2_web_acl" "chatbot" {
       }
     }
   }
+
+  # The following rule blocks requests that arrive in POST on the API Gateway
+  # with route "/queries" when their number is greater than 10 per minute
+  rule {
+    name     = "block-requests-to-queries"
+    priority = 0
+
+    action {
+      count {} # To test with use `count {}` instead of `block {}`
+    }
+
+    statement {
+      rate_based_statement {
+        aggregate_key_type    = "CONSTANT"
+        limit                 = var.waf_block_requests_to_queries_limit
+        evaluation_window_sec = var.waf_block_requests_to_queries_evaluation_window_sec
+
+        scope_down_statement {
+          byte_match_statement {
+            field_to_match {
+              uri_path {}
+            }
+
+            positional_constraint = "STARTS_WITH"
+            search_string         = "/queries"
+
+            text_transformation {
+              priority = 1
+              type     = "NONE"
+            }
+          }
+        }
+      }
+    }
+
+    visibility_config {
+      sampled_requests_enabled = true
+      cloudwatch_metrics_enabled = true
+      metric_name = "${local.prefix}-waf-metric-block-requests-to-queries"
+    }
+  }
+
+  # The following rule apply a rate limit of N requests from the same IP
+  rule {
+    name     = "ip-rate-limit"
+    priority = 0
+
+    action {
+      count {} # To test with use `count {}` instead of `block {}`
+    }
+
+    statement {
+      rate_based_statement {
+        limit                 = var.waf_ip_rate_limit_limit
+        aggregate_key_type    = "IP"
+        evaluation_window_sec = var.waf_ip_rate_limit_evaluation_window_sec
+      }
+    }
+
+    visibility_config {
+      cloudwatch_metrics_enabled = true
+      metric_name                = "${local.prefix}-waf-metric-ip-rate-limit"
+      sampled_requests_enabled   = false
+    }
+  }
+
+  # The following rule blocks requests that are identified as bot control
+  rule {
+    name     = "block-bot-control-requests"
+    priority = 1
+
+    override_action {
+      count {} # To test with use `count {}` instead of `block {}`
+    }
+
+    statement {
+      managed_rule_group_statement {
+        name        = "AWSManagedRulesBotControlRuleSet"
+        vendor_name = "AWS"
+      }
+    }
+    visibility_config {
+      cloudwatch_metrics_enabled = true
+      metric_name                = "${local.prefix}-waf-metric-bot-control"
+      sampled_requests_enabled   = false
+    }
+  }
+}
+
+resource "aws_wafv2_web_acl_association" "chatbot" {
+  resource_arn = aws_api_gateway_deployment.stage.execution_arn
+  web_acl_arn  = aws_wafv2_web_acl.chatbot.arn
+
+  depends_on = [
+    aws_api_gateway_deployment.stage,
+    aws_wafv2_web_acl.chatbot
+  ]
 }