bug: Exclude internal calls from WAF rate limiter (#591)

aws-samples · Oct 18, 2024 · b6a5d5a · b6a5d5a
1 parent 77ec531
commit b6a5d5a
Show file tree

Hide file tree

Showing 3 changed files with 133 additions and 3 deletions.
diff --git a/lib/chatbot-api/index.ts b/lib/chatbot-api/index.ts
@@ -1,5 +1,6 @@
 import * as cognito from "aws-cdk-lib/aws-cognito";
 import * as dynamodb from "aws-cdk-lib/aws-dynamodb";
+import * as ec2 from "aws-cdk-lib/aws-ec2";
 import * as s3 from "aws-cdk-lib/aws-s3";
 import * as sqs from "aws-cdk-lib/aws-sqs";
 import * as sns from "aws-cdk-lib/aws-sns";
@@ -110,7 +111,10 @@ export class ChatBotApi extends Construct {
           name: "WafAppsync",
           rules: [
             ...props.shared.webACLRules,
-            ...this.createWafRules(props.config.llms.rateLimitPerIP ?? 100),
+            ...this.createWafRules(
+              props.config.llms.rateLimitPerIP ?? 100,
+              props.shared.vpc
+            ),
           ],
         }).attrArn,
         resourceArn: api.arn,
@@ -175,7 +179,10 @@ export class ChatBotApi extends Construct {
     ]);
   }
 
-  private createWafRules(llmRatePerIP: number): wafv2.CfnWebACL.RuleProperty[] {
+  private createWafRules(
+    llmRatePerIP: number,
+    vpc: ec2.Vpc
+  ): wafv2.CfnWebACL.RuleProperty[] {
     /**
      * The rate limit is the maximum number of requests from a
      * single IP address that are allowed in a ten-minute period.
@@ -242,6 +249,39 @@ export class ChatBotApi extends Construct {
         metricName: "LimitRequestsPerIP",
       },
     };
-    return [ruleLimitRequests];
+
+    // The following rule is disabling throttling for calls coming from the VPC.
+    const eips: string[] = [];
+    vpc.node.findAll().forEach((resource) => {
+      if (resource instanceof ec2.CfnEIP) {
+        // NAT Gateways IP
+        eips.push(resource.attrPublicIp + "/32");
+      }
+    });
+
+    const vpcnIpSet = new wafv2.CfnIPSet(this, "VPCPublicIPs", {
+      addresses: eips,
+      ipAddressVersion: "IPV4",
+      scope: "REGIONAL",
+    });
+
+    const allowInternalCalls: wafv2.CfnWebACL.RuleProperty = {
+      name: "AllowInternalCalls",
+      priority: 2,
+      action: {
+        allow: {},
+      },
+      statement: {
+        ipSetReferenceStatement: {
+          arn: vpcnIpSet.attrArn,
+        },
+      },
+      visibilityConfig: {
+        sampledRequestsEnabled: false,
+        cloudWatchMetricsEnabled: false,
+        metricName: "AllowInternalCalls",
+      },
+    };
+    return [ruleLimitRequests, allowInternalCalls];
   }
 }
diff --git a/tests/__snapshots__/cdk-app.test.ts.snap b/tests/__snapshots__/cdk-app.test.ts.snap
@@ -4057,6 +4057,29 @@ schema {
       },
       "Type": "AWS::IAM::Policy",
     },
+    "ChatBotApiVPCPublicIPsAE6206D3": {
+      "Properties": {
+        "Addresses": [
+          {
+            "Fn::Join": [
+              "",
+              [
+                {
+                  "Fn::GetAtt": [
+                    "SharedVPCpublicSubnet1EIPA9E2FC1C",
+                    "PublicIp",
+                  ],
+                },
+                "/32",
+              ],
+            ],
+          },
+        ],
+        "IPAddressVersion": "IPV4",
+        "Scope": "REGIONAL",
+      },
+      "Type": "AWS::WAFv2::IPSet",
+    },
     "ChatBotApiWafAppsync9FEB4E22": {
       "Properties": {
         "DefaultAction": {
@@ -4145,6 +4168,28 @@ schema {
               "SampledRequestsEnabled": true,
             },
           },
+          {
+            "Action": {
+              "Allow": {},
+            },
+            "Name": "AllowInternalCalls",
+            "Priority": 2,
+            "Statement": {
+              "IPSetReferenceStatement": {
+                "Arn": {
+                  "Fn::GetAtt": [
+                    "ChatBotApiVPCPublicIPsAE6206D3",
+                    "Arn",
+                  ],
+                },
+              },
+            },
+            "VisibilityConfig": {
+              "CloudWatchMetricsEnabled": false,
+              "MetricName": "AllowInternalCalls",
+              "SampledRequestsEnabled": false,
+            },
+          },
         ],
         "Scope": "REGIONAL",
         "VisibilityConfig": {

diff --git a/tests/chatbot-api/__snapshots__/chatbot-api-construct.test.ts.snap b/tests/chatbot-api/__snapshots__/chatbot-api-construct.test.ts.snap
@@ -4119,6 +4119,29 @@ schema {
       },
       "Type": "AWS::IAM::Policy",
     },
+    "ChatBotApiConstructVPCPublicIPs5C63CEBB": {
+      "Properties": {
+        "Addresses": [
+          {
+            "Fn::Join": [
+              "",
+              [
+                {
+                  "Fn::GetAtt": [
+                    "SharedVPCpublicSubnet1EIPA9E2FC1C",
+                    "PublicIp",
+                  ],
+                },
+                "/32",
+              ],
+            ],
+          },
+        ],
+        "IPAddressVersion": "IPV4",
+        "Scope": "REGIONAL",
+      },
+      "Type": "AWS::WAFv2::IPSet",
+    },
     "ChatBotApiConstructWafAppsyncE6AACFFE": {
       "Properties": {
         "DefaultAction": {
@@ -4207,6 +4230,28 @@ schema {
               "SampledRequestsEnabled": true,
             },
           },
+          {
+            "Action": {
+              "Allow": {},
+            },
+            "Name": "AllowInternalCalls",
+            "Priority": 2,
+            "Statement": {
+              "IPSetReferenceStatement": {
+                "Arn": {
+                  "Fn::GetAtt": [
+                    "ChatBotApiConstructVPCPublicIPs5C63CEBB",
+                    "Arn",
+                  ],
+                },
+              },
+            },
+            "VisibilityConfig": {
+              "CloudWatchMetricsEnabled": false,
+              "MetricName": "AllowInternalCalls",
+              "SampledRequestsEnabled": false,
+            },
+          },
         ],
         "Scope": "REGIONAL",
         "VisibilityConfig": {