feat: token bucket rate limiting

package.json

@@ -30,6 +30,7 @@
     "jose": "^5.9.3",
     "pg-promise": "^11.9.1",
     "pino": "^9.4.0",
+    "rate-limiter-flexible": "^5.0.4",
     "redis": "^4.7.0"
   },
   "devDependencies": {

src/@types/express/index.d.ts

@@ -5,6 +5,7 @@ declare global {
   // biome-ignore lint/style/noNamespace: <It is the only way to do this.  If we use module Biome also complains>
   namespace Express {
     interface Request {
+      tokenConsumedOnFormId?: string;
       serviceAccountId: string;
       serviceUserId: string;
     }

src/config.ts

@@ -4,6 +4,11 @@ export enum EnvironmentMode {
   Production = "production",
 }
 
+export type TokenBucketConfiguration = {
+  capacity: number;
+  numberOfSecondsBeforeRefill: number;
+};
+
 // AWS SDK
 
 export const AWS_REGION: string = "ca-central-1";
@@ -30,6 +35,18 @@ export const LOCALSTACK_ENDPOINT: string | undefined = loadOptionalEnvVar(
   "LOCALSTACK_ENDPOINT",
 );
 
+// Rate limiting
+
+export const lowRateLimiterConfiguration: TokenBucketConfiguration = {
+  capacity: 500,
+  numberOfSecondsBeforeRefill: 60,
+};
+
+export const highRateLimiterConfiguration: TokenBucketConfiguration = {
+  capacity: 1000,
+  numberOfSecondsBeforeRefill: 60,
+};
+
 // Redis
 
 export const REDIS_URL: string = loadRequiredEnvVar("REDIS_URL");

src/lib/formsClient/getPublicKey.ts

@@ -1,5 +1,8 @@
 import { DatabaseConnectorClient } from "@lib/integration/databaseConnector.js";
-import { getValueFromCache, cacheValue } from "@lib/utils/cache.js";
+import {
+  getValueFromRedis,
+  setValueInRedis,
+} from "@lib/integration/redis/redisClientAdapter.js";
 import { logMessage } from "@lib/logging/logger.js";
 
 const REDIS_PUBLIC_KEY_PREFIX: string = "api:publicKey";
@@ -52,14 +55,14 @@ function retrievePublicKeyFromDatabase(
 function getPublicKeyFromCache(
   serviceAccountId: string,
 ): Promise<string | undefined> {
-  return getValueFromCache(`${REDIS_PUBLIC_KEY_PREFIX}:${serviceAccountId}`);
+  return getValueFromRedis(`${REDIS_PUBLIC_KEY_PREFIX}:${serviceAccountId}`);
 }
 
 function cachePublicKey(
   publicKey: string,
   serviceAccountId: string,
 ): Promise<void> {
-  return cacheValue(
+  return setValueInRedis(
     `${REDIS_PUBLIC_KEY_PREFIX}:${serviceAccountId}`,
     publicKey,
     CACHE_EXPIRY_DELAY_IN_SECONDS,

src/lib/idp/verifyAccessToken.ts

@@ -1,5 +1,8 @@
 import { introspectAccessToken } from "@lib/integration/zitadelConnector.js";
-import { cacheValue, getValueFromCache } from "@lib/utils/cache.js";
+import {
+  setValueInRedis,
+  getValueFromRedis,
+} from "@lib/integration/redis/redisClientAdapter.js";
 import { createHash } from "node:crypto";
 import { logMessage } from "@lib/logging/logger.js";
 
@@ -60,7 +63,7 @@
 ): Promise<VerifiedAccessToken | undefined> {
   const key = getVerifiedAccessTokenCacheKey(accessToken);
 
-  return getValueFromCache(key).then((value) => {
+  return getValueFromRedis(key).then((value) => {
     return value ? (JSON.parse(value) as VerifiedAccessToken) : undefined;
   });
 }
@@ -71,7 +74,7 @@
 ): Promise<void> {
   const key = getVerifiedAccessTokenCacheKey(accessToken);
 
-  return cacheValue(
+  return setValueInRedis(
     key,
     JSON.stringify(introspectionResult),
     CACHE_EXPIRY_DELAY_IN_SECONDS,

src/lib/utils/cache.ts → src/lib/integration/redis/redisClientAdapter.ts

@@ -1,14 +1,14 @@
-import { RedisConnector } from "@lib/integration/redisConnector.js";
+import { RedisConnector } from "@lib/integration/redis/redisConnector.js";
 
-export function getValueFromCache(key: string): Promise<string | undefined> {
+export function getValueFromRedis(key: string): Promise<string | undefined> {
   return RedisConnector.getInstance()
     .then((redisConnector) => redisConnector.client.get(key))
     .then((value) => {
       return value ?? undefined;
     });
 }
 
-export function cacheValue(
+export function setValueInRedis(
   key: string,
   value: string,
   expiryDelayInSeconds?: number,

src/lib/rateLimiting/tokenBucketLimiter.ts

@@ -0,0 +1,53 @@
+import { getTokenBucketRateLimiterAssociatedToForm } from "@lib/rateLimiting/tokenBucketProvider.js";
+import type { RateLimiterRes } from "rate-limiter-flexible";
+
+export type BucketStatus = {
+  bucketCapacity: number;
+  remainingTokens: number;
+  numberOfMillisecondsBeforeRefill: number;
+};
+
+export type ConsumeTokenResult = {
+  wasAbleToConsumeToken: boolean;
+  bucketStatus: BucketStatus;
+};
+
+export async function consumeTokenIfAvailable(
+  formId: string,
+): Promise<ConsumeTokenResult> {
+  const tokenBucket = await getTokenBucketRateLimiterAssociatedToForm(formId);
+
+  try {
+    const consumptionResult = await tokenBucket.consume(formId);
+
+    return {
+      wasAbleToConsumeToken: true,
+      bucketStatus: buildBucketStatus(tokenBucket.points, consumptionResult),
+    };
+  } catch (rateLimiterRes) {
+    // Since our token buckets have an `insuranceLimiter` set up, the consume function promise can only be rejected with a `RateLimiterRes` object
+    return {
+      wasAbleToConsumeToken: false,
+      bucketStatus: buildBucketStatus(
+        tokenBucket.points,
+        rateLimiterRes as RateLimiterRes,
+      ),
+    };
+  }
+}
+
+export async function refundConsumedToken(formId: string): Promise<void> {
+  const tokenBucket = await getTokenBucketRateLimiterAssociatedToForm(formId);
+  await tokenBucket.reward(formId);
+}
+
+function buildBucketStatus(
+  bucketCapacity: number,
+  rateLimiterResponse: RateLimiterRes,
+): BucketStatus {
+  return {
+    bucketCapacity: bucketCapacity,
+    remainingTokens: rateLimiterResponse.remainingPoints,
+    numberOfMillisecondsBeforeRefill: rateLimiterResponse.msBeforeNext,
+  };
+}

src/lib/rateLimiting/tokenBucketProvider.ts

@@ -0,0 +1,77 @@
+import {
+  type RateLimiterAbstract,
+  RateLimiterMemory,
+  RateLimiterRedis,
+} from "rate-limiter-flexible";
+import { RedisConnector } from "@lib/integration/redis/redisConnector.js";
+import {
+  highRateLimiterConfiguration,
+  lowRateLimiterConfiguration,
+} from "@config";
+import { getValueFromRedis } from "@lib/integration/redis/redisClientAdapter.js";
+import { logMessage } from "@lib/logging/logger.js";
+
+const REDIS_RATE_LIMIT_KEY_PREFIX: string = "rate-limit";
+
+const redisClient = await RedisConnector.getInstance().then(
+  (instance) => instance.client,
+);
+
+/**
+ * From the official documentation:
+ * `inMemoryBlockOnConsumed`: Can be used against DDoS attacks. In-memory blocking works in current process memory and for consume method only.
+ *                            It blocks a key in memory for msBeforeNext milliseconds from the last consume result, if inMemoryBlockDuration is
+ *                            not set. This helps to avoid extra requests. inMemoryBlockOnConsumed value is supposed to be equal or more than
+ *                            points option. Sometimes it is not necessary to increment counter on store, if all points are consumed already.
+ * `insuranceLimiter`: Instance of RateLimiterAbstract extended object to store limits, when database comes up with any error.
+ */
+
+const lowCapacityTokenBucket = new RateLimiterRedis({
+  keyPrefix: "low-capacity-token-bucket",
+  storeClient: redisClient,
+  useRedisPackage: true,
+  points: lowRateLimiterConfiguration.capacity,
+  duration: lowRateLimiterConfiguration.numberOfSecondsBeforeRefill,
+  inMemoryBlockOnConsumed: lowRateLimiterConfiguration.capacity,
+  insuranceLimiter: new RateLimiterMemory({
+    keyPrefix: "backup-low-capacity-token-bucket",
+    points: lowRateLimiterConfiguration.capacity,
+    duration: lowRateLimiterConfiguration.numberOfSecondsBeforeRefill,
+  }),
+});
+
+const highCapacityTokenBucket = new RateLimiterRedis({
+  keyPrefix: "high-capacity-token-bucket",
+  storeClient: redisClient,
+  useRedisPackage: true,
+  points: highRateLimiterConfiguration.capacity,
+  duration: highRateLimiterConfiguration.numberOfSecondsBeforeRefill,
+  inMemoryBlockOnConsumed: highRateLimiterConfiguration.capacity,
+  insuranceLimiter: new RateLimiterMemory({
+    keyPrefix: "backup-high-capacity-token-bucket",
+    points: highRateLimiterConfiguration.capacity,
+    duration: highRateLimiterConfiguration.numberOfSecondsBeforeRefill,
+  }),
+});
+
+export function getTokenBucketRateLimiterAssociatedToForm(
+  formId: string,
+): Promise<RateLimiterAbstract> {
+  return getValueFromRedis(`${REDIS_RATE_LIMIT_KEY_PREFIX}:${formId}`)
+    .then((value) => {
+      switch (value) {
+        case "high":
+          return highCapacityTokenBucket;
+        default:
+          return lowCapacityTokenBucket;
+      }
+    })
+    .catch((error) => {
+      logMessage.warn(
+        error,
+        `[token-bucket-provider] Failed to retrieve token bucket capacity for form ${formId}. Will use low capacity bucket by default`,
+      );
+
+      return lowCapacityTokenBucket;
+    });
+}

src/middleware/globalErrorHandler.ts

@@ -1,12 +1,18 @@
 import type { Request, Response, NextFunction } from "express";
 import { logMessage } from "@lib/logging/logger.js";
+import { refundConsumedToken } from "@lib/rateLimiting/tokenBucketLimiter.js";
 
-export function globalErrorHandlerMiddleware(
+export async function globalErrorHandlerMiddleware(
   error: Error,
-  _request: Request,
+  request: Request,
   response: Response,
   _next: NextFunction,
-): void {
+): Promise<void> {
   logMessage.error(error, "Global unhandled error");
+
+  if (request.tokenConsumedOnFormId !== undefined) {
+    await refundConsumedToken(request.tokenConsumedOnFormId);
+  }
+
   response.sendStatus(500);
 }

src/middleware/rateLimiter.ts

@@ -1,9 +1,49 @@
 import type { NextFunction, Request, Response } from "express";
+import { consumeTokenIfAvailable } from "@lib/rateLimiting/tokenBucketLimiter.js";
+import { logMessage } from "@lib/logging/logger.js";
 
-export function rateLimiterMiddleware(
-  _request: Request,
-  _response: Response,
+export async function rateLimiterMiddleware(
+  request: Request,
+  response: Response,
   next: NextFunction,
-): void {
-  next();
+): Promise<void> {
+  try {
+    const formId = request.params.formId;
+
+    const consumeTokenResult = await consumeTokenIfAvailable(formId);
+
+    response.header({
+      "X-RateLimit-Limit": consumeTokenResult.bucketStatus.bucketCapacity,
+      "X-RateLimit-Remaining": consumeTokenResult.bucketStatus.remainingTokens,
+      "X-RateLimit-Reset": new Date(
+        Date.now() +
+          consumeTokenResult.bucketStatus.numberOfMillisecondsBeforeRefill,
+      ),
+    });
+
+    if (consumeTokenResult.wasAbleToConsumeToken === false) {
+      response.header({
+        "Retry-After":
+          consumeTokenResult.bucketStatus.numberOfMillisecondsBeforeRefill /
+          1000,
+      });
+
+      logMessage.info(
+        `[rate-limiter] Form ${formId} consumed all ${consumeTokenResult.bucketStatus.bucketCapacity} tokens. Bucket will be refilled in ${consumeTokenResult.bucketStatus.numberOfMillisecondsBeforeRefill / 1000} seconds`,
+      );
+
+      response.sendStatus(429);
+      return;
+    }
+
+    request.tokenConsumedOnFormId = formId;
+
+    next();
+  } catch (error) {
+    next(
+      new Error("[middleware] Internal error with rate limiter", {
+        cause: error,
+      }),
+    );
+  }
 }

src/router.ts

@@ -54,9 +54,12 @@ export function buildRouter(): Router {
     .use("/template", templateRoute)
     .use("/submission", submissionRoute);
 
-  const formsRoute = Router()
-    .use(rateLimiterMiddleware)
-    .use("/:formId(c[a-z0-9]{24})", authenticationMiddleware, formIdRoute);
+  const formsRoute = Router().use(
+    "/:formId(c[a-z0-9]{24})",
+    authenticationMiddleware,
+    rateLimiterMiddleware,
+    formIdRoute,
+  );
 
   const statusRoute = Router().get(
     "/",

src/server.ts

@@ -3,7 +3,7 @@ import express, { type Express } from "express";
 import { SERVER_PORT } from "@config";
 import { buildRouter } from "./router.js";
 import { getApiAuditLogSqsQueueUrl } from "@lib/integration/awsSqsQueueLoader.js";
-import { RedisConnector } from "@lib/integration/redisConnector.js";
+import { RedisConnector } from "@lib/integration/redis/redisConnector.js";
 import { logMessage } from "@lib/logging/logger.js";
 
 const server: Express = express();