feat: Added rate limit + security headers.

cli/magic-config.ts

@@ -206,6 +206,8 @@ const embeddingModels = [
       options.advancedMonitoring = config.advancedMonitoring;
       options.createVpcEndpoints = config.vpc?.createVpcEndpoints;
       options.logRetention = config.logRetention;
+      options.rateLimitPerAIP = config.rateLimitPerIP;
+      options.llmRateLimitPerIP = config.llms.rateLimitPerIP;
       options.privateWebsite = config.privateWebsite;
       options.certificate = config.certificate;
       options.domain = config.domain;
@@ -294,15 +296,15 @@ async function processCreateOptions(options: any): Promise<void> {
       name: "createCMKs",
       message:
         "Do you want to create KMS Customer Managed Keys (CMKs)? (It will be used to encrypt the data at rest.)",
-      initial: true,
+      initial: options.createCMKs ?? true,
       hint: "It is recommended but enabling it on an existing environment will cause the re-creation of some of the resources (for example Aurora cluster, Open Search collection). To prevent data loss, it is recommended to use it on a new environment or at least enable retain on cleanup (needs to be deployed before enabling the use of CMK). For more information on Aurora migration, please refer to the documentation.",
     },
     {
       type: "confirm",
       name: "retainOnDelete",
       message:
         "Do you want to retain data stores on cleanup of the project (Logs, S3, Tables, Indexes, Cognito User pools)?",
-      initial: true,
+      initial: options.retainOnDelete ?? true,
       hint: "It reduces the risk of deleting data. It will however not delete all the resources on cleanup (would require manual removal if relevant)",
     },
     {
@@ -828,6 +830,38 @@ async function processCreateOptions(options: any): Promise<void> {
   const models: any = await enquirer.prompt(modelsPrompts);
 
   const advancedSettingsPrompts = [
+    {
+      type: "input",
+      name: "llmRateLimitPerIP",
+      message:
+        "What is the allowed rate per IP for Gen AI calls (over 10 minutes)? This is used by the SendQuery mutation only",
+      initial: options.llmRateLimitPerIP
+        ? String(options.llmRateLimitPerIP)
+        : "100",
+      validate(value: string) {
+        if (Number(value) >= 10) {
+          return true;
+        } else {
+          return "Should be more than 10";
+        }
+      },
+    },
+    {
+      type: "input",
+      name: "rateLimitPerIP",
+      message:
+        "What the allowed per IP for all calls (over 10 minutes)? This is used by the all the AppSync APIs and CloudFront",
+      initial: options.rateLimitPerAIP
+        ? String(options.rateLimitPerAIP)
+        : "400",
+      validate(value: string) {
+        if (Number(value) >= 10) {
+          return true;
+        } else {
+          return "Should be more than 10";
+        }
+      },
+    },
     {
       type: "input",
       name: "logRetention",
@@ -874,7 +908,7 @@ async function processCreateOptions(options: any): Promise<void> {
       name: "customPublicDomain",
       message:
         "Do you want to provide a custom domain name and corresponding certificate arn for the public website ?",
-      initial: options.customPublicDomain || false,
+      initial: options.domain ? true : false,
       skip(): boolean {
         return (this as any).state.answers.privateWebsite;
       },
@@ -1137,6 +1171,9 @@ async function processCreateOptions(options: any): Promise<void> {
     logRetention: advancedSettings.logRetention
       ? Number(advancedSettings.logRetention)
       : undefined,
+    rateLimitPerAIP: advancedSettings?.rateLimitPerIP
+      ? Number(advancedSettings?.rateLimitPerIP)
+      : undefined,
     certificate: advancedSettings.certificate,
     domain: advancedSettings.domain,
     cognitoFederation: advancedSettings.cognitoFederationEnabled
@@ -1182,6 +1219,9 @@ async function processCreateOptions(options: any): Promise<void> {
         }
       : undefined,
     llms: {
+      rateLimitPerAIP: advancedSettings?.llmRateLimitPerIP
+        ? Number(advancedSettings?.llmRateLimitPerIP)
+        : undefined,
       sagemaker: answers.sagemakerModels,
       huggingfaceApiSecretArn: answers.huggingfaceApiSecretArn,
       sagemakerSchedule: answers.enableSagemakerModelsSchedule

lib/aws-genai-llm-chatbot-stack.ts

@@ -158,6 +158,7 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
       userPoolClientId: authentication.userPoolClient.userPoolClientId,
       api: chatBotApi,
       chatbotFilesBucket: chatBotApi.filesBucket,
+      uploadBucket: ragEngines?.uploadBucket,
       crossEncodersEnabled:
         typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
       sagemakerEmbeddingsEnabled:

lib/chatbot-api/index.ts

@@ -5,6 +5,7 @@ import * as sqs from "aws-cdk-lib/aws-sqs";
 import * as sns from "aws-cdk-lib/aws-sns";
 import * as ssm from "aws-cdk-lib/aws-ssm";
 import * as iam from "aws-cdk-lib/aws-iam";
+import * as wafv2 from "aws-cdk-lib/aws-wafv2";
 import * as cdk from "aws-cdk-lib";
 import * as path from "path";
 import { Construct } from "constructs";
@@ -95,6 +96,27 @@ export class ChatBotApi extends Construct {
         : appsync.Visibility.GLOBAL,
     });
 
+    if (props.shared.webACLRules.length > 0) {
+      new wafv2.CfnWebACLAssociation(this, "WebACLAssociation", {
+        webAclArn: new wafv2.CfnWebACL(this, "WafAppsync", {
+          defaultAction: { allow: {} },
+          scope: "REGIONAL",
+          visibilityConfig: {
+            cloudWatchMetricsEnabled: true,
+            metricName: "WafAppsync",
+            sampledRequestsEnabled: true,
+          },
+          description: "WAFv2 ACL for APPSync",
+          name: "WafAppsync",
+          rules: [
+            ...props.shared.webACLRules,
+            ...this.createWafRules(props.config.llms.rateLimitPerIP ?? 100),
+          ],
+        }).attrArn,
+        resourceArn: api.arn,
+      });
+    }
+
     const apiResolvers = new ApiResolvers(this, "RestApi", {
       ...props,
       sessionsTable: chatTables.sessionsTable,
@@ -152,4 +174,74 @@ export class ChatBotApi extends Construct {
       },
     ]);
   }
+
+  private createWafRules(llmRatePerIP: number): 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.
+     * The IP address is automatically unblocked after it falls below the limit.
+     */
+    const ruleLimitRequests: wafv2.CfnWebACL.RuleProperty = {
+      name: "LimitLLMRequestsPerIP",
+      priority: 1,
+      action: {
+        block: {
+          customResponse: {
+            responseCode: 429,
+          },
+        },
+      },
+      statement: {
+        rateBasedStatement: {
+          limit: llmRatePerIP,
+          evaluationWindowSec: 60 * 10,
+          aggregateKeyType: "IP",
+          scopeDownStatement: {
+            andStatement: {
+              statements: [
+                {
+                  byteMatchStatement: {
+                    searchString: "/graphql",
+                    fieldToMatch: {
+                      uriPath: {},
+                    },
+                    textTransformations: [
+                      {
+                        priority: 0,
+                        type: "NONE",
+                      },
+                    ],
+                    positionalConstraint: "EXACTLY",
+                  },
+                },
+                {
+                  byteMatchStatement: {
+                    searchString: "mutation SendQuery(",
+                    fieldToMatch: {
+                      body: {
+                        oversizeHandling: "MATCH",
+                      },
+                    },
+                    textTransformations: [
+                      {
+                        priority: 0,
+                        type: "NONE",
+                      },
+                    ],
+                    positionalConstraint: "CONTAINS",
+                  },
+                },
+              ],
+            },
+          },
+        },
+      },
+      visibilityConfig: {
+        sampledRequestsEnabled: true,
+        cloudWatchMetricsEnabled: true,
+        metricName: "LimitRequestsPerIP",
+      },
+    };
+    return [ruleLimitRequests];
+  }
 }

lib/shared/index.ts

@@ -5,6 +5,7 @@ import * as lambda from "aws-cdk-lib/aws-lambda";
 import * as secretsmanager from "aws-cdk-lib/aws-secretsmanager";
 import * as ssm from "aws-cdk-lib/aws-ssm";
 import * as logs from "aws-cdk-lib/aws-logs";
+import * as wafv2 from "aws-cdk-lib/aws-wafv2";
 import { Construct } from "constructs";
 import * as path from "path";
 import { Layer } from "../layer";
@@ -36,6 +37,7 @@ export class Shared extends Construct {
   readonly powerToolsLayer: lambda.ILayerVersion;
   readonly sharedCode: SharedAssetBundler;
   readonly s3vpcEndpoint: ec2.InterfaceVpcEndpoint;
+  readonly webACLRules: wafv2.CfnWebACL.RuleProperty[] = [];
 
   constructor(scope: Construct, id: string, props: SharedProps) {
     super(scope, id);
@@ -250,6 +252,8 @@ export class Shared extends Construct {
       }
     }
 
+    this.webACLRules = this.createWafRules(props.config.rateLimitPerIP ?? 400);
+
     const configParameter = new ssm.StringParameter(this, "Config", {
       stringValue: JSON.stringify(props.config),
     });
@@ -316,4 +320,32 @@ export class Shared extends Construct {
       { id: "AwsSolutions-SMG4", reason: "Secret value is blank." },
     ]);
   }
+
+  private createWafRules(ratePerIP: number): 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.
+     * The IP address is automatically unblocked after it falls below the limit.
+     */
+    const ruleLimitRequests: wafv2.CfnWebACL.RuleProperty = {
+      name: "LimitRequestsPerIP",
+      priority: 10,
+      action: {
+        block: {},
+      },
+      statement: {
+        rateBasedStatement: {
+          limit: ratePerIP,
+          evaluationWindowSec: 60 * 10,
+          aggregateKeyType: "IP",
+        },
+      },
+      visibilityConfig: {
+        sampledRequestsEnabled: true,
+        cloudWatchMetricsEnabled: true,
+        metricName: "LimitRequestsPerIP",
+      },
+    };
+    return [ruleLimitRequests];
+  }
 }

lib/shared/types.ts

@@ -84,6 +84,7 @@ export interface SystemConfig {
   certificate?: string;
   domain?: string;
   privateWebsite?: boolean;
+  rateLimitPerIP?: number;
   cognitoFederation?: {
     enabled?: boolean;
     autoRedirect?: boolean;
@@ -113,6 +114,7 @@ export interface SystemConfig {
     };
   };
   llms: {
+    rateLimitPerIP?: number;
     sagemaker: SupportedSageMakerModels[];
     huggingfaceApiSecretArn?: string;
     sagemakerSchedule?: {

lib/user-interface/index.ts

@@ -25,6 +25,7 @@ export interface UserInterfaceProps {
   readonly userPoolClient: cognito.UserPoolClient;
   readonly api: ChatBotApi;
   readonly chatbotFilesBucket: s3.Bucket;
+  readonly uploadBucket?: s3.Bucket;
   readonly crossEncodersEnabled: boolean;
   readonly sagemakerEmbeddingsEnabled: boolean;
 }
@@ -56,8 +57,12 @@ export class UserInterface extends Construct {
       blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
       autoDeleteObjects: true,
       bucketName: props.config.privateWebsite ? props.config.domain : undefined,
-      websiteIndexDocument: "index.html",
-      websiteErrorDocument: "index.html",
+      websiteIndexDocument: props.config.privateWebsite
+        ? "index.html"
+        : undefined,
+      websiteErrorDocument: props.config.privateWebsite
+        ? "index.html"
+        : undefined,
       enforceSSL: true,
       serverAccessLogsBucket: uploadLogsBucket,
       // Cloudfront with OAI only supports S3 Managed Key (would need to migrate to OAC)
@@ -80,6 +85,8 @@ export class UserInterface extends Construct {
       const publicWebsite = new PublicWebsite(this, "PublicWebsite", {
         ...props,
         websiteBucket: websiteBucket,
+        chatbotFilesBucket: props.chatbotFilesBucket,
+        uploadBucket: props.uploadBucket,
       });
       this.cloudFrontDistribution = publicWebsite.distribution;
       this.publishedDomain = props.config.domain