build rate limiter logic

Author: devksingh4

src/api/index.ts

@@ -24,6 +24,7 @@ import { SecretsManagerClient } from "@aws-sdk/client-secrets-manager";
 import mobileWalletRoute from "./routes/mobileWallet.js";
 import stripeRoutes from "./routes/stripe.js";
 import membershipPlugin from "./routes/membership.js";
+import rateLimiterPlugin from "./plugins/rateLimiter.js";
 
 dotenv.config();
 
@@ -71,6 +72,10 @@ async function init() {
   await app.register(fastifyZodValidationPlugin);
   await app.register(FastifyAuthProvider);
   await app.register(errorHandlerPlugin);
+  await app.register(rateLimiterPlugin, {
+    limit: 50,
+    duration: 20,
+  });
   if (!process.env.RunEnvironment) {
     process.env.RunEnvironment = "dev";
   }

src/api/plugins/rateLimiter.ts

@@ -0,0 +1,118 @@
+import fp from "fastify-plugin";
+import {
+  ConditionalCheckFailedException,
+  UpdateItemCommand,
+} from "@aws-sdk/client-dynamodb";
+import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
+import { genericConfig } from "common/config.js";
+import { FastifyPluginAsync, FastifyRequest, FastifyReply } from "fastify";
+
+interface RateLimiterOptions {
+  limit?: number;
+  duration?: number;
+}
+
+interface RateLimitParams {
+  ddbClient: DynamoDBClient;
+  rateLimitIdentifier: string;
+  duration: number;
+  limit: number;
+  userIdentifier: string;
+}
+
+async function isAtLimit({
+  ddbClient,
+  rateLimitIdentifier,
+  duration,
+  limit,
+  userIdentifier,
+}: RateLimitParams): Promise<{
+  limited: boolean;
+  resetTime: number;
+  used: number;
+}> {
+  const nowInSeconds = Math.floor(Date.now() / 1000);
+  const timeWindow = Math.floor(nowInSeconds / duration) * duration;
+  const PK = `rate-limit:${rateLimitIdentifier}:${userIdentifier}:${timeWindow}`;
+
+  try {
+    const result = await ddbClient.send(
+      new UpdateItemCommand({
+        TableName: genericConfig.RateLimiterDynamoTableName,
+        Key: {
+          PK: { S: PK },
+          SK: { S: "counter" },
+        },
+        UpdateExpression: "ADD #rateLimitCount :inc SET #ttl = :ttl",
+        ConditionExpression:
+          "attribute_not_exists(#rateLimitCount) OR #rateLimitCount <= :limit",
+        ExpressionAttributeValues: {
+          ":inc": { N: "1" },
+          ":limit": { N: limit.toString() },
+          ":ttl": { N: (timeWindow + duration).toString() },
+        },
+        ExpressionAttributeNames: {
+          "#rateLimitCount": "rateLimitCount",
+          "#ttl": "ttl",
+        },
+        ReturnValues: "UPDATED_NEW",
+        ReturnValuesOnConditionCheckFailure: "ALL_OLD",
+      }),
+    );
+    return {
+      limited: false,
+      used: parseInt(result.Attributes?.rateLimitCount.N || "1", 10),
+      resetTime: timeWindow + duration,
+    };
+  } catch (error) {
+    if (error instanceof ConditionalCheckFailedException) {
+      return { limited: true, resetTime: timeWindow + duration, used: limit };
+    }
+    throw error;
+  }
+}
+
+const rateLimiterPlugin: FastifyPluginAsync<RateLimiterOptions> = async (
+  fastify,
+  options,
+) => {
+  const { limit = 10, duration = 60 } = options;
+
+  fastify.addHook(
+    "onRequest",
+    async (request: FastifyRequest, reply: FastifyReply) => {
+      const userIdentifier = request.ip;
+      const rateLimitIdentifier = "api-request";
+
+      const { limited, resetTime, used } = await isAtLimit({
+        ddbClient: fastify.dynamoClient,
+        rateLimitIdentifier,
+        duration,
+        limit,
+        userIdentifier,
+      });
+      reply.header("X-RateLimit-Limit", limit.toString());
+      reply.header("X-RateLimit-Reset", resetTime?.toString() || "0");
+      reply.header(
+        "X-RateLimit-Remaining",
+        limited ? 0 : used ? limit - used : limit - 1,
+      );
+      if (limited) {
+        const retryAfter = resetTime
+          ? resetTime - Math.floor(Date.now() / 1000)
+          : undefined;
+        reply.header("Retry-After", retryAfter?.toString() || "0");
+        return reply.status(429).send({
+          error: true,
+          name: "RateLimitExceededError",
+          id: 409,
+          message: "Rate limit exceeded.",
+        });
+      }
+    },
+  );
+};
+
+export default fp(rateLimiterPlugin, {
+  name: "fastify-rate-limiter",
+});

src/common/config.ts

@@ -21,6 +21,7 @@ export type ConfigType = {
 };
 
 export type GenericConfigType = {
+  RateLimiterDynamoTableName: string;
   EventsDynamoTableName: string;
   CacheDynamoTableName: string;
   StripeLinksDynamoTableName: string;
@@ -52,6 +53,7 @@ export const commChairsGroupId = "105e7d32-7289-435e-a67a-552c7f215507";
 export const miscTestingGroupId = "ff25ec56-6a33-420d-bdb0-51d8a3920e46";
 
 const genericConfig: GenericConfigType = {
+  RateLimiterDynamoTableName: "infra-core-api-rate-limiter",
   EventsDynamoTableName: "infra-core-api-events",
   StripeLinksDynamoTableName: "infra-core-api-stripe-links",
   CacheDynamoTableName: "infra-core-api-cache",