Use dynamo for encrypted cache (#186)

Lambda decryption performance completely wipes out all gains from using
redis

Author: devksingh4

src/api/functions/encryption.ts

@@ -30,8 +30,6 @@ import { UserProfileData } from "common/types/msGraphApi.js";
 import { SecretsManagerClient } from "@aws-sdk/client-secrets-manager";
 import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
 import { checkPaidMembershipFromTable } from "./membership.js";
-import { getKey, setKey } from "./redisCache.js";
-import RedisClient from "ioredis";
 import type pino from "pino";
 import { type FastifyBaseLogger } from "fastify";
 
@@ -41,16 +39,14 @@ function validateGroupId(groupId: string): boolean {
 }
 
 type GetEntraIdTokenInput = {
-  clients: { smClient: SecretsManagerClient; redisClient: RedisClient.default };
-  encryptionSecret: string;
+  clients: { smClient: SecretsManagerClient; dynamoClient: DynamoDBClient };
   clientId: string;
   scopes?: string[];
   secretName?: string;
   logger: pino.Logger | FastifyBaseLogger;
 };
 export async function getEntraIdToken({
   clients,
-  encryptionSecret,
   clientId,
   scopes = ["https://graph.microsoft.com/.default"],
   secretName,
@@ -72,14 +68,13 @@ export async function getEntraIdToken({
     "base64",
   ).toString("utf8");
   const cacheKey = `entra_id_access_token_${localSecretName}_${clientId}`;
-  const cachedTokenObject = await getKey<{ token: string }>({
-    redisClient: clients.redisClient,
-    key: cacheKey,
-    encryptionSecret,
-    logger,
-  });
-  if (cachedTokenObject) {
-    return cachedTokenObject.token;
+  const startTime = Date.now();
+  const cachedToken = await getItemFromCache(clients.dynamoClient, cacheKey);
+  logger.debug(
+    `Took ${Date.now() - startTime} ms to get cached entra ID token.`,
+  );
+  if (cachedToken) {
+    return cachedToken.token as string;
   }
   const config = {
     auth: {
@@ -103,17 +98,13 @@ export async function getEntraIdToken({
       });
     }
     date.setTime(date.getTime() - 30000);
-    if (result?.accessToken && result?.expiresOn) {
-      await setKey({
-        redisClient: clients.redisClient,
-        key: cacheKey,
-        data: JSON.stringify({ token: result.accessToken }),
-        expiresIn:
-          Math.floor(
-            (result.expiresOn.getTime() - new Date().getTime()) / 1000,
-          ) - 120, // get new token 2 min before expiry
-        encryptionSecret,
-      });
+    if (result?.accessToken) {
+      await insertItemIntoCache(
+        clients.dynamoClient,
+        cacheKey,
+        { token: result?.accessToken },
+        date,
+      );
     }
     return result?.accessToken ?? null;
   } catch (error) {

src/api/functions/entraId.ts

@@ -1,19 +1,10 @@
-import { DecryptionError } from "common/errors/index.js";
 import type RedisModule from "ioredis";
-import { z } from "zod";
-import {
-  CORRUPTED_DATA_MESSAGE,
-  decrypt,
-  encrypt,
-  INVALID_DECRYPTION_MESSAGE,
-} from "./encryption.js";
 import type pino from "pino";
 import { type FastifyBaseLogger } from "fastify";
 
 export type GetFromCacheInput = {
   redisClient: RedisModule.default;
   key: string;
-  encryptionSecret?: string;
   logger: pino.Logger | FastifyBaseLogger;
 };
 
@@ -22,70 +13,31 @@ export type SetInCacheInput = {
   key: string;
   data: string;
   expiresIn?: number;
-  encryptionSecret?: string;
+  logger: pino.Logger | FastifyBaseLogger;
 };
 
-const redisEntrySchema = z.object({
-  isEncrypted: z.boolean(),
-  data: z.string(),
-});
-
 export async function getKey<T extends object>({
   redisClient,
   key,
-  encryptionSecret,
   logger,
 }: GetFromCacheInput): Promise<T | null> {
+  logger.debug(`Getting redis key "${key}".`);
   const data = await redisClient.get(key);
   if (!data) {
     return null;
   }
-  const decoded = await redisEntrySchema.parseAsync(JSON.parse(data));
-  if (!decoded.isEncrypted) {
-    return JSON.parse(decoded.data) as T;
-  }
-  if (!encryptionSecret) {
-    throw new DecryptionError({
-      message: "Encrypted data found but no decryption key provided.",
-    });
-  }
-  try {
-    const decryptedData = decrypt({
-      cipherText: decoded.data,
-      encryptionSecret,
-    });
-    return JSON.parse(decryptedData) as T;
-  } catch (e) {
-    if (
-      e instanceof DecryptionError &&
-      (e.message === INVALID_DECRYPTION_MESSAGE ||
-        e.message === CORRUPTED_DATA_MESSAGE)
-    ) {
-      logger.info(
-        `Invalid decryption, deleting old Redis key and continuing...`,
-      );
-      await redisClient.del(key);
-      return null;
-    }
-    throw e;
-  }
+  return JSON.parse(data) as T;
 }
 
 export async function setKey({
   redisClient,
   key,
-  encryptionSecret,
   data,
   expiresIn,
+  logger,
 }: SetInCacheInput) {
-  const realData = encryptionSecret
-    ? encrypt({ plaintext: data, encryptionSecret })
-    : data;
-  const redisPayload: z.infer<typeof redisEntrySchema> = {
-    isEncrypted: !!encryptionSecret,
-    data: realData,
-  };
-  const strRedisPayload = JSON.stringify(redisPayload);
+  const strRedisPayload = data;
+  logger.debug(`Setting redis key "${key}".`);
   return expiresIn
     ? await redisClient.set(key, strRedisPayload, "EX", expiresIn)
     : await redisClient.set(key, strRedisPayload);

src/api/functions/redisCache.ts

@@ -159,7 +159,6 @@ const authPlugin: FastifyPluginAsync = async (fastify, _options) => {
       disableApiKeyAuth: boolean,
     ): Promise<Set<AppRoles>> => {
       const { redisClient } = fastify;
-      const encryptionSecret = fastify.secretConfig.encryption_key;
       const startTime = new Date().getTime();
       try {
         if (!disableApiKeyAuth) {
@@ -252,6 +251,7 @@ const authPlugin: FastifyPluginAsync = async (fastify, _options) => {
               key: `jwksKey:${header.kid}`,
               data: JSON.stringify({ key: signingKey }),
               expiresIn: JWKS_CACHE_SECONDS,
+              logger: request.log,
             });
             request.log.debug("Got JWKS signing key from server.");
           }
@@ -332,6 +332,7 @@ const authPlugin: FastifyPluginAsync = async (fastify, _options) => {
             data: JSON.stringify([...userRoles]),
             redisClient,
             expiresIn: GENERIC_CACHE_SECONDS,
+            logger: request.log,
           });
           request.log.debug("Retrieved user roles from database.");
         }

src/api/plugins/auth.ts

@@ -100,7 +100,6 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
         clients: await getAuthorizedClients(),
         clientId: fastify.environmentConfig.AadValidClientId,
         secretName: genericConfig.EntraSecretName,
-        encryptionSecret: fastify.secretConfig.encryption_key,
         logger: request.log,
       });
       await patchUserProfile(
@@ -166,7 +165,6 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
           clients: await getAuthorizedClients(),
           clientId: fastify.environmentConfig.AadValidClientId,
           secretName: genericConfig.EntraSecretName,
-          encryptionSecret: fastify.secretConfig.encryption_key,
           logger: request.log,
         });
         const groupMembers = listGroupMembers(entraIdToken, groupId);
@@ -234,7 +232,6 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
         clients: await getAuthorizedClients(),
         clientId: fastify.environmentConfig.AadValidClientId,
         secretName: genericConfig.EntraSecretName,
-        encryptionSecret: fastify.secretConfig.encryption_key,
         logger: request.log,
       });
       if (!entraIdToken) {
@@ -330,7 +327,6 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
         clients: await getAuthorizedClients(),
         clientId: fastify.environmentConfig.AadValidClientId,
         secretName: genericConfig.EntraSecretName,
-        encryptionSecret: fastify.secretConfig.encryption_key,
         logger: request.log,
       });
       const groupMetadataPromise = getGroupMetadata(entraIdToken, groupId);
@@ -585,7 +581,6 @@ No action is required from you at this time.
         clients: await getAuthorizedClients(),
         clientId: fastify.environmentConfig.AadValidClientId,
         secretName: genericConfig.EntraSecretName,
-        encryptionSecret: fastify.secretConfig.encryption_key,
         logger: request.log,
       });
       const response = await listGroupMembers(entraIdToken, groupId);
@@ -604,22 +599,24 @@ No action is required from you at this time.
       onRequest: fastify.authorizeFromSchema,
     },
     async (request, reply) => {
-      const entraIdToken = await getEntraIdToken({
-        clients: await getAuthorizedClients(),
-        clientId: fastify.environmentConfig.AadValidClientId,
-        secretName: genericConfig.EntraSecretName,
-        encryptionSecret: fastify.secretConfig.encryption_key,
-        logger: request.log,
-      });
       const { redisClient } = fastify;
       const key = `entra_manageable_groups_${fastify.environmentConfig.EntraServicePrincipalId}`;
       const redisResponse = await getKey<{ displayName: string; id: string }[]>(
         { redisClient, key, logger: request.log },
       );
       if (redisResponse) {
         request.log.debug("Got manageable groups from Redis cache.");
-        return reply.status(200).send(redisResponse);
+        return reply
+          .header("X-ACM-Data-Source", "redis")
+          .status(200)
+          .send(redisResponse);
       }
+      const entraIdToken = await getEntraIdToken({
+        clients: await getAuthorizedClients(),
+        clientId: fastify.environmentConfig.AadValidClientId,
+        secretName: genericConfig.EntraSecretName,
+        logger: request.log,
+      });
       // get groups, but don't show protected groups as manageable
       const freshData = (
         await getServicePrincipalOwnedGroups(
@@ -639,6 +636,7 @@ No action is required from you at this time.
         key,
         data: JSON.stringify(freshData),
         expiresIn: GENERIC_CACHE_SECONDS,
+        logger: request.log,
       });
       return reply.status(200).send(freshData);
     },

src/api/routes/iam.ts

@@ -112,7 +112,6 @@ const membershipPlugin: FastifyPluginAsync = async (fastify, _options) => {
           clients: await getAuthorizedClients(),
           clientId: fastify.environmentConfig.AadValidClientId,
           secretName: genericConfig.EntraSecretName,
-          encryptionSecret: fastify.secretConfig.encryption_key,
           logger: request.log,
         });
         const paidMemberGroup = fastify.environmentConfig.PaidMemberGroupId;
@@ -233,7 +232,6 @@ const membershipPlugin: FastifyPluginAsync = async (fastify, _options) => {
           clients: await getAuthorizedClients(),
           clientId: fastify.environmentConfig.AadValidClientId,
           secretName: genericConfig.EntraSecretName,
-          encryptionSecret: fastify.secretConfig.encryption_key,
           logger: request.log,
         });
         const paidMemberGroup = fastify.environmentConfig.PaidMemberGroupId;