use redis cache for critical paths

Author: devksingh4

cloudformation/iam.yml

@@ -99,14 +99,6 @@ Resources:
             Resource:
               - Fn::Sub: arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/infra-core-api-cache
 
-          - Sid: DynamoDBRateLimitTableAccess
-            Effect: Allow
-            Action:
-              - dynamodb:DescribeTable
-              - dynamodb:UpdateItem
-            Resource:
-              - Fn::Sub: arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/infra-core-api-rate-limiter
-
           - Sid: DynamoDBAuditLogTableAccess
             Effect: Allow
             Action:

cloudformation/main.yml

@@ -408,30 +408,6 @@ Resources:
         - AttributeName: userEmail
           KeyType: HASH
 
-  RateLimiterTable:
-    Type: "AWS::DynamoDB::Table"
-    DeletionPolicy: "Delete"
-    UpdateReplacePolicy: "Delete"
-    Properties:
-      BillingMode: "PAY_PER_REQUEST"
-      TableName: infra-core-api-rate-limiter
-      DeletionProtectionEnabled: true
-      PointInTimeRecoverySpecification:
-        PointInTimeRecoveryEnabled: false
-      AttributeDefinitions:
-        - AttributeName: PK
-          AttributeType: S
-        - AttributeName: SK
-          AttributeType: S
-      KeySchema:
-        - AttributeName: PK
-          KeyType: HASH
-        - AttributeName: SK
-          KeyType: RANGE
-      TimeToLiveSpecification:
-        AttributeName: ttl
-        Enabled: true
-
   EventRecordsTable:
     Type: "AWS::DynamoDB::Table"
     DeletionPolicy: "Retain"

src/api/functions/authorization.ts

@@ -9,14 +9,8 @@ export const AUTH_DECISION_CACHE_SECONDS = 180;
 
 export async function getUserRoles(
   dynamoClient: DynamoDBClient,
-  fastifyApp: FastifyInstance,
   userId: string,
 ): Promise<AppRoles[]> {
-  const cachedValue = fastifyApp.nodeCache.get(`userroles-${userId}`);
-  if (cachedValue) {
-    fastifyApp.log.info(`Returning cached auth decision for user ${userId}`);
-    return cachedValue as AppRoles[];
-  }
   const tableName = `${genericConfig.IAMTablePrefix}-userroles`;
   const command = new GetItemCommand({
     TableName: tableName,
@@ -37,32 +31,13 @@ export async function getUserRoles(
   if (!("roles" in items)) {
     return [];
   }
-  if (items.roles[0] === "all") {
-    fastifyApp.nodeCache.set(
-      `userroles-${userId}`,
-      allAppRoles,
-      AUTH_DECISION_CACHE_SECONDS,
-    );
-    return allAppRoles;
-  }
-  fastifyApp.nodeCache.set(
-    `userroles-${userId}`,
-    items.roles,
-    AUTH_DECISION_CACHE_SECONDS,
-  );
   return items.roles as AppRoles[];
 }
 
 export async function getGroupRoles(
   dynamoClient: DynamoDBClient,
-  fastifyApp: FastifyInstance,
   groupId: string,
 ) {
-  const cachedValue = fastifyApp.nodeCache.get(`grouproles-${groupId}`);
-  if (cachedValue) {
-    fastifyApp.log.info(`Returning cached auth decision for group ${groupId}`);
-    return cachedValue as AppRoles[];
-  }
   const tableName = `${genericConfig.IAMTablePrefix}-grouproles`;
   const command = new GetItemCommand({
     TableName: tableName,
@@ -77,34 +52,14 @@ export async function getGroupRoles(
     });
   }
   if (!response.Item) {
-    fastifyApp.nodeCache.set(
-      `grouproles-${groupId}`,
-      [],
-      AUTH_DECISION_CACHE_SECONDS,
-    );
     return [];
   }
   const items = unmarshall(response.Item) as { roles: AppRoles[] | ["all"] };
   if (!("roles" in items)) {
-    fastifyApp.nodeCache.set(
-      `grouproles-${groupId}`,
-      [],
-      AUTH_DECISION_CACHE_SECONDS,
-    );
     return [];
   }
   if (items.roles[0] === "all") {
-    fastifyApp.nodeCache.set(
-      `grouproles-${groupId}`,
-      allAppRoles,
-      AUTH_DECISION_CACHE_SECONDS,
-    );
     return allAppRoles;
   }
-  fastifyApp.nodeCache.set(
-    `grouproles-${groupId}`,
-    items.roles,
-    AUTH_DECISION_CACHE_SECONDS,
-  );
   return items.roles as AppRoles[];
 }

src/api/functions/discord.ts

@@ -14,7 +14,7 @@ import moment from "moment-timezone";
 import { FastifyBaseLogger } from "fastify";
 import { DiscordEventError } from "../../common/errors/index.js";
 import { getSecretValue } from "../plugins/auth.js";
-import { genericConfig } from "../../common/config.js";
+import { genericConfig, SecretConfig } from "../../common/config.js";
 import { SecretsManagerClient } from "@aws-sdk/client-secrets-manager";
 
 // https://stackoverflow.com/a/3809435/5684541
@@ -26,19 +26,16 @@ export type IUpdateDiscord = EventPostRequest & { id: string };
 const urlRegex = /https:\/\/[a-z0-9.-]+\/calendar\?id=([a-f0-9-]+)/;
 
 export const updateDiscord = async (
-  smClient: SecretsManagerClient,
+  secretApiConfig: SecretConfig,
   event: IUpdateDiscord,
   actor: string,
   isDelete: boolean = false,
   logger: FastifyBaseLogger,
 ): Promise<null | GuildScheduledEventCreateOptions> => {
-  const secretApiConfig =
-    (await getSecretValue(smClient, genericConfig.ConfigSecretName)) || {};
   const client = new Client({ intents: [GatewayIntentBits.Guilds] });
   let payload: GuildScheduledEventCreateOptions | null = null;
-
   client.once(Events.ClientReady, async (readyClient: Client<true>) => {
-    logger.info(`Logged in as ${readyClient.user.tag}`);
+    logger.debug(`Logged in as ${readyClient.user.tag}`);
     const guildID = secretApiConfig.discord_guild_id;
     const guild = await client.guilds.fetch(guildID?.toString() || "");
     const discordEvents = await guild.scheduledEvents.fetch();
@@ -69,6 +66,7 @@ export const updateDiscord = async (
         logger.warn(`Event with id ${id} not found in Discord`);
       }
       await client.destroy();
+      logger.debug("Logged out of Discord.");
       return null;
     }
 
@@ -108,6 +106,7 @@ export const updateDiscord = async (
     }
 
     await client.destroy();
+    logger.debug("Logged out of Discord.");
     return payload;
   });
 

src/api/functions/rateLimit.ts

@@ -1,66 +1,60 @@
-import {
-  ConditionalCheckFailedException,
-  UpdateItemCommand,
-  DynamoDBClient,
-} from "@aws-sdk/client-dynamodb";
-import { genericConfig } from "common/config.js";
+import { Redis } from "ioredis"; // Make sure you have ioredis installed (npm install ioredis)
 
 interface RateLimitParams {
-  ddbClient: DynamoDBClient;
+  redisClient: Redis;
   rateLimitIdentifier: string;
   duration: number;
   limit: number;
   userIdentifier: string;
 }
 
+interface RateLimitResult {
+  limited: boolean;
+  resetTime: number;
+  used: number;
+}
+
+const LUA_SCRIPT_INCREMENT_AND_EXPIRE = `
+  local count = redis.call("INCR", KEYS[1])
+  -- If the count is 1, this means the key was just created by INCR (first request in this window).
+  -- So, we set its expiration time to the end of the current window.
+  if tonumber(count) == 1 then
+    redis.call("EXPIREAT", KEYS[1], ARGV[1])
+  end
+  return count
+`;
+
 export async function isAtLimit({
-  ddbClient,
+  redisClient,
   rateLimitIdentifier,
   duration,
   limit,
   userIdentifier,
-}: RateLimitParams): Promise<{
-  limited: boolean;
-  resetTime: number;
-  used: number;
-}> {
+}: RateLimitParams): Promise<RateLimitResult> {
+  if (duration <= 0) {
+    throw new Error("Rate limit duration must be a positive number.");
+  }
+  if (limit < 0) {
+    throw new Error("Rate limit must be a non-negative number.");
+  }
+
   const nowInSeconds = Math.floor(Date.now() / 1000);
-  const timeWindow = Math.floor(nowInSeconds / duration) * duration;
-  const PK = `rate-limit:${rateLimitIdentifier}:${userIdentifier}:${timeWindow}`;
+  const timeWindowStart = Math.floor(nowInSeconds / duration) * duration;
+  const key = `rate-limit:${rateLimitIdentifier}:${userIdentifier}:${timeWindowStart}`;
+  const expiryTimestamp = timeWindowStart + duration;
 
-  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 currentUsedCount = (await redisClient.eval(
+    LUA_SCRIPT_INCREMENT_AND_EXPIRE,
+    1, // Number of keys
+    key, // KEYS[1]
+    expiryTimestamp.toString(), // ARGV[1]
+  )) as number; // The script returns the count, which is a number.
+  const isLimited = currentUsedCount > limit;
+  const resetTime = expiryTimestamp;
+
+  return {
+    limited: isLimited,
+    resetTime,
+    used: isLimited ? limit : currentUsedCount,
+  };
 }

src/api/index.ts

@@ -4,14 +4,18 @@ import { randomUUID } from "crypto";
 import fastify, { FastifyInstance } from "fastify";
 import FastifyAuthProvider from "@fastify/auth";
 import fastifyStatic from "@fastify/static";
-import fastifyAuthPlugin from "./plugins/auth.js";
+import fastifyAuthPlugin, { getSecretValue } from "./plugins/auth.js";
 import protectedRoute from "./routes/protected.js";
 import errorHandlerPlugin from "./plugins/errorHandler.js";
 import { RunEnvironment, runEnvironments } from "../common/roles.js";
 import { InternalServerError } from "../common/errors/index.js";
 import eventsPlugin from "./routes/events.js";
 import cors from "@fastify/cors";
-import { environmentConfig, genericConfig } from "../common/config.js";
+import {
+  environmentConfig,
+  genericConfig,
+  SecretConfig,
+} from "../common/config.js";
 import organizationsPlugin from "./routes/organizations.js";
 import authorizeFromSchemaPlugin from "./plugins/authorizeFromSchema.js";
 import evaluatePoliciesPlugin from "./plugins/evaluatePolicies.js";
@@ -43,6 +47,8 @@ import {
 import { ZodOpenApiVersion } from "zod-openapi";
 import { withTags } from "./components/index.js";
 import apiKeyRoute from "./routes/apiKey.js";
+import RedisModule from "ioredis";
+import fastifyCron from "fastify-cron";
 
 dotenv.config();
 
@@ -56,6 +62,12 @@ async function init(prettyPrint: boolean = false) {
   const secretsManagerClient = new SecretsManagerClient({
     region: genericConfig.AwsRegion,
   });
+  const secret = (await getSecretValue(
+    secretsManagerClient,
+    genericConfig.ConfigSecretName,
+  )) as SecretConfig;
+  const redisClient = new RedisModule.default(secret.redis_url);
+
   const transport = prettyPrint
     ? {
         target: "pino-pretty",
@@ -224,6 +236,26 @@ async function init(prettyPrint: boolean = false) {
   app.nodeCache = new NodeCache({ checkperiod: 30 });
   app.dynamoClient = dynamoClient;
   app.secretsManagerClient = secretsManagerClient;
+  app.redisClient = redisClient;
+  app.secretConfig = secret;
+  app.refreshSecretConfig = async () => {
+    app.secretConfig = (await getSecretValue(
+      app.secretsManagerClient,
+      genericConfig.ConfigSecretName,
+    )) as SecretConfig;
+  };
+  app.register(fastifyCron.default, {
+    // refresh secrets config
+    jobs: [
+      {
+        cronTime: "*/15 * * * *",
+        onTick: async (server) => {
+          server.log.info("Refreshing secrets manager config.");
+          await server.refreshSecretConfig();
+        },
+      },
+    ],
+  });
   app.addHook("onRequest", (req, _, done) => {
     req.startTime = now();
     const hostname = req.hostname;
@@ -250,7 +282,13 @@ async function init(prettyPrint: boolean = false) {
         summary: "Verify that the API server is healthy.",
       }),
     },
-    (_, reply) => reply.send({ message: "UP" }),
+    async (_, reply) => {
+      const startTime = new Date().getTime();
+      await app.redisClient.ping();
+      const redisTime = new Date().getTime();
+      app.log.debug(`Redis latency: ${redisTime - startTime} ms.`);
+      return reply.send({ message: "UP" });
+    },
   );
   await app.register(
     async (api, _options) => {
@@ -295,7 +333,7 @@ if (import.meta.url === `file://${process.argv[1]}`) {
     process.exit(1);
   }
   const app = await init(true);
-  app.listen({ port: 8080 }, async (err) => {
+  app.listen({ port: 8080 }, (err) => {
     /* eslint no-console: ["error", {"allow": ["log", "error"]}] */
     if (err) {
       console.error(err);

src/api/package.json

@@ -41,10 +41,12 @@
     "dotenv": "^16.5.0",
     "esbuild": "^0.25.3",
     "fastify": "^5.3.2",
+    "fastify-cron": "^1.4.0",
     "fastify-plugin": "^5.0.1",
     "fastify-raw-body": "^5.0.0",
     "fastify-zod-openapi": "^4.1.1",
     "ical-generator": "^8.1.1",
+    "ioredis": "^5.6.1",
     "jsonwebtoken": "^9.0.2",
     "jwks-rsa": "^3.2.0",
     "moment": "^2.30.1",