Add IP based rate limit middleware, powered by upstash

The minimum rate limitting limits in AWS WAF are too high for our
purposes of trying to avoid cost spikes caused by to many openai API
requests.

apps/aws-app/cdk/app.ts

@@ -1,4 +1,5 @@
 import * as cdk from 'aws-cdk-lib';
+import './env.js';
 import {MainStack} from './main-stack.js';
 import {WafStack} from './waf-stack.js';
 

apps/aws-app/cdk/env.ts

@@ -0,0 +1,20 @@
+import {z} from 'zod';
+
+declare global {
+  namespace NodeJS {
+    interface ProcessEnv extends z.infer<typeof envVariables> {}
+  }
+}
+
+const envVariables = z.object({
+  AWS_ACCESS_KEY_ID: z.string(),
+  AWS_HANDLER_VERIFY_HEADER: z.string(),
+  AWS_REGION: z.string(),
+  AWS_SECRET_ACCESS_KEY: z.string(),
+  GOOGLE_SEARCH_API_KEY: z.string(),
+  GOOGLE_SEARCH_SEARCH_ENGINE_ID: z.string(),
+  UPSTASH_REDIS_REST_TOKEN: z.string(),
+  UPSTASH_REDIS_REST_URL: z.string(),
+});
+
+envVariables.parse(process.env);

apps/aws-app/cdk/main-stack.ts

@@ -2,7 +2,6 @@ import path from 'path';
 import * as cdk from 'aws-cdk-lib';
 import type {Construct} from 'constructs';
 
-const verifyHeader = process.env.AWS_HANDLER_VERIFY_HEADER;
 const distDirname = path.join(import.meta.dirname, `../dist/`);
 
 export interface MainStackProps extends cdk.StackProps {
@@ -26,9 +25,14 @@ export class MainStack extends cdk.Stack {
         runtime: cdk.aws_lambda.Runtime.NODEJS_20_X,
         bundling: {format: cdk.aws_lambda_nodejs.OutputFormat.ESM},
         timeout: cdk.Duration.minutes(1),
-        environment: verifyHeader
-          ? {AWS_HANDLER_VERIFY_HEADER: verifyHeader}
-          : undefined,
+        environment: {
+          AWS_HANDLER_VERIFY_HEADER: process.env.AWS_HANDLER_VERIFY_HEADER,
+          GOOGLE_SEARCH_API_KEY: process.env.GOOGLE_SEARCH_API_KEY,
+          GOOGLE_SEARCH_SEARCH_ENGINE_ID:
+            process.env.GOOGLE_SEARCH_SEARCH_ENGINE_ID,
+          UPSTASH_REDIS_REST_TOKEN: process.env.UPSTASH_REDIS_REST_TOKEN,
+          UPSTASH_REDIS_REST_URL: process.env.UPSTASH_REDIS_REST_URL,
+        },
       },
     );
 
@@ -47,9 +51,9 @@ export class MainStack extends cdk.Stack {
     const distribution = new cdk.aws_cloudfront.Distribution(this, `cdn`, {
       defaultBehavior: {
         origin: new cdk.aws_cloudfront_origins.FunctionUrlOrigin(functionUrl, {
-          customHeaders: verifyHeader
-            ? {'X-Origin-Verify': verifyHeader}
-            : undefined,
+          customHeaders: {
+            'X-Origin-Verify': process.env.AWS_HANDLER_VERIFY_HEADER,
+          },
         }),
         allowedMethods: cdk.aws_cloudfront.AllowedMethods.ALLOW_ALL,
         cachePolicy: new cdk.aws_cloudfront.CachePolicy(this, `cache-policy`, {

apps/aws-app/package.json

@@ -18,6 +18,8 @@
   "dependencies": {
     "@mfng/core": "*",
     "@mfng/shared-app": "*",
+    "@upstash/ratelimit": "^1.0.1",
+    "@upstash/redis": "^1.28.4",
     "ai": "^3.0.11",
     "clsx": "^1.2.1",
     "openai": "^4.29.0",

apps/aws-app/src/handler/index.tsx

@@ -20,10 +20,12 @@ import {
   reactServerManifest,
   reactSsrManifest,
 } from './manifests.js';
+import {ratelimitMiddleware} from './ratelimit-middleware.js';
 
 export const app = new Hono();
 
 app.use(authMiddleware);
+app.use(ratelimitMiddleware);
 app.use(loggerMiddleware);
 app.get(`/*`, async (context) => handleGet(context.req.raw));
 app.post(`/*`, async (context) => handlePost(context.req.raw));

apps/aws-app/src/handler/ratelimit-middleware.ts

@@ -0,0 +1,38 @@
+import {Ratelimit} from '@upstash/ratelimit';
+import {Redis} from '@upstash/redis';
+import type {MiddlewareHandler} from 'hono';
+
+let ratelimit: Ratelimit | undefined;
+
+try {
+  ratelimit = new Ratelimit({
+    redis: Redis.fromEnv(),
+    limiter: Ratelimit.slidingWindow(10, `10 m`),
+    analytics: true,
+    ephemeralCache: new Map(),
+  });
+} catch {
+  console.warn(`Unable to create Redis instance, no rate limits are applied.`);
+}
+
+export const ratelimitMiddleware: MiddlewareHandler = async (context, next) => {
+  if (ratelimit && context.req.method === `POST`) {
+    const ip = context.req.header(`cloudfront-viewer-address`);
+
+    if (ip) {
+      const {success, reset} = await ratelimit.limit(ip);
+
+      if (!success) {
+        console.debug(`Rate limit reached for ${ip}`);
+
+        return context.text(`Too Many Requests`, 429, {
+          'Retry-After': ((reset - Date.now()) / 1000).toFixed(),
+        });
+      }
+    } else {
+      console.warn(`cloudfront-viewer-address not provided`);
+    }
+  }
+
+  return next();
+};