replace all getEntraToken with redis version

devksingh4 · devksingh4 · commit c1b80e2fd72d · 2025-07-03T13:43:41.000-05:00
diff --git a/src/api/functions/encryption.ts b/src/api/functions/encryption.ts
@@ -7,6 +7,9 @@ const KEY_LEN = 32;
 const ALGORITHM = "aes-256-gcm";
 const HASH_FUNCTION = "sha512";
 
+export const INVALID_DECRYPTION_MESSAGE =
+  "Could not decrypt data (check that the encryption secret is correct).";
+
 export function encrypt({
   plaintext,
   encryptionSecret,
@@ -62,8 +65,7 @@ export function decrypt({
     ]);
   } catch (e) {
     throw new DecryptionError({
-      message:
-        "Could not decrypt data (check that the encryption secret is correct).",
+      message: INVALID_DECRYPTION_MESSAGE,
     });
   }
 
diff --git a/src/api/functions/entraId.ts b/src/api/functions/entraId.ts
@@ -9,6 +9,7 @@ import {
 } from "../../common/config.js";
 import {
   BaseError,
+  DecryptionError,
   EntraFetchError,
   EntraGroupError,
   EntraGroupsFromEmailError,
@@ -29,18 +30,32 @@ 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";
 
 function validateGroupId(groupId: string): boolean {
   const groupIdPattern = /^[a-zA-Z0-9-]+$/; // Adjust the pattern as needed
   return groupIdPattern.test(groupId);
 }
 
-export async function getEntraIdToken(
-  clients: { smClient: SecretsManagerClient; dynamoClient: DynamoDBClient },
-  clientId: string,
-  scopes: string[] = ["https://graph.microsoft.com/.default"],
-  secretName?: string,
-) {
+type GetEntraIdTokenInput = {
+  clients: { smClient: SecretsManagerClient; redisClient: RedisClient.default };
+  encryptionSecret: string;
+  clientId: string;
+  scopes?: string[];
+  secretName?: string;
+  logger: pino.Logger | FastifyBaseLogger;
+};
+export async function getEntraIdToken({
+  clients,
+  encryptionSecret,
+  clientId,
+  scopes = ["https://graph.microsoft.com/.default"],
+  secretName,
+  logger,
+}: GetEntraIdTokenInput) {
   const localSecretName = secretName || genericConfig.EntraSecretName;
   const secretApiConfig =
     (await getSecretValue(clients.smClient, localSecretName)) || {};
@@ -56,12 +71,15 @@ export async function getEntraIdToken(
     secretApiConfig.entra_id_private_key as string,
     "base64",
   ).toString("utf8");
-  const cachedToken = await getItemFromCache(
-    clients.dynamoClient,
-    `entra_id_access_token_${localSecretName}_${clientId}`,
-  );
-  if (cachedToken) {
-    return cachedToken.token as string;
+  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 config = {
     auth: {
@@ -85,13 +103,14 @@ export async function getEntraIdToken(
       });
     }
     date.setTime(date.getTime() - 30000);
-    if (result?.accessToken) {
-      await insertItemIntoCache(
-        clients.dynamoClient,
-        `entra_id_access_token_${localSecretName}`,
-        { token: result?.accessToken },
-        date,
-      );
+    if (result?.accessToken && result?.expiresOn) {
+      await setKey({
+        redisClient: clients.redisClient,
+        key: cacheKey,
+        data: JSON.stringify({ token: result.accessToken }),
+        expiresIn: result.expiresOn.getTime() - new Date().getTime() - 3600,
+        encryptionSecret,
+      });
     }
     return result?.accessToken ?? null;
   } catch (error) {
diff --git a/src/api/functions/redisCache.ts b/src/api/functions/redisCache.ts
@@ -1,12 +1,15 @@
 import { DecryptionError } from "common/errors/index.js";
 import type RedisModule from "ioredis";
 import { z } from "zod";
-import { decrypt, encrypt } from "./encryption.js";
+import { 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;
 };
 
 export type SetInCacheInput = {
@@ -26,6 +29,7 @@ export async function getKey<T extends object>({
   redisClient,
   key,
   encryptionSecret,
+  logger,
 }: GetFromCacheInput): Promise<T | null> {
   const data = await redisClient.get(key);
   if (!data) {
@@ -40,8 +44,25 @@ export async function getKey<T extends object>({
       message: "Encrypted data found but no decryption key provided.",
     });
   }
-  const decryptedData = decrypt({ cipherText: decoded.data, encryptionSecret });
-  return JSON.parse(decryptedData) as T;
+  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
+    ) {
+      logger.info(
+        `Invalid decryption, deleting old Redis key and continuing...`,
+      );
+      await redisClient.del(key);
+      return null;
+    }
+    throw e;
+  }
 }
 
 export async function setKey({
diff --git a/src/api/plugins/auth.ts b/src/api/plugins/auth.ts
@@ -232,6 +232,7 @@ const authPlugin: FastifyPluginAsync = async (fastify, _options) => {
           const cachedJwksSigningKey = await getKey<{ key: string }>({
             redisClient,
             key: `jwksKey:${header.kid}`,
+            logger: request.log,
           });
           if (cachedJwksSigningKey) {
             signingKey = cachedJwksSigningKey.key;
@@ -271,6 +272,7 @@ const authPlugin: FastifyPluginAsync = async (fastify, _options) => {
         const cachedRoles = await getKey<string[]>({
           key: `authCache:${request.username}:roles`,
           redisClient,
+          logger: request.log,
         });
         if (cachedRoles) {
           request.userRoles = new Set(cachedRoles as AppRoles[]);
diff --git a/src/api/routes/iam.ts b/src/api/routes/iam.ts
@@ -61,6 +61,7 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
           region: genericConfig.AwsRegion,
           credentials,
         }),
+        redisClient: fastify.redisClient,
       };
       fastify.log.info(
         `Assumed Entra role ${roleArns.Entra} to get the Entra token.`,
@@ -73,6 +74,7 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
     return {
       smClient: fastify.secretsManagerClient,
       dynamoClient: fastify.dynamoClient,
+      redisClient: fastify.redisClient,
     };
   };
   fastify.withTypeProvider<FastifyZodOpenApiTypeProvider>().patch(
@@ -94,12 +96,13 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
         });
       }
       const userOid = request.tokenPayload.oid;
-      const entraIdToken = await getEntraIdToken(
-        await getAuthorizedClients(),
-        fastify.environmentConfig.AadValidClientId,
-        undefined,
-        genericConfig.EntraSecretName,
-      );
+      const entraIdToken = await getEntraIdToken({
+        clients: await getAuthorizedClients(),
+        clientId: fastify.environmentConfig.AadValidClientId,
+        secretName: genericConfig.EntraSecretName,
+        encryptionSecret: fastify.secretConfig.encryption_key,
+        logger: request.log,
+      });
       await patchUserProfile(
         entraIdToken,
         request.username,
@@ -213,10 +216,13 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
     },
     async (request, reply) => {
       const emails = request.body.emails;
-      const entraIdToken = await getEntraIdToken(
-        await getAuthorizedClients(),
-        fastify.environmentConfig.AadValidClientId,
-      );
+      const entraIdToken = await getEntraIdToken({
+        clients: await getAuthorizedClients(),
+        clientId: fastify.environmentConfig.AadValidClientId,
+        secretName: genericConfig.EntraSecretName,
+        encryptionSecret: fastify.secretConfig.encryption_key,
+        logger: request.log,
+      });
       if (!entraIdToken) {
         throw new InternalServerError({
           message: "Could not get Entra ID token to perform task.",
@@ -306,10 +312,13 @@ const iamRoutes: FastifyPluginAsync = async (fastify, _options) => {
           group: groupId,
         });
       }
-      const entraIdToken = await getEntraIdToken(
-        await getAuthorizedClients(),
-        fastify.environmentConfig.AadValidClientId,
-      );
+      const entraIdToken = await getEntraIdToken({
+        clients: await getAuthorizedClients(),
+        clientId: fastify.environmentConfig.AadValidClientId,
+        secretName: genericConfig.EntraSecretName,
+        encryptionSecret: fastify.secretConfig.encryption_key,
+        logger: request.log,
+      });
       const groupMetadataPromise = getGroupMetadata(entraIdToken, groupId);
       const addResults = await Promise.allSettled(
         request.body.add.map((email) =>
@@ -550,12 +559,13 @@ No action is required from you at this time.
           group: groupId,
         });
       }
-      const entraIdToken = await getEntraIdToken(
-        await getAuthorizedClients(),
-        fastify.environmentConfig.AadValidReadOnlyClientId,
-        undefined,
-        genericConfig.EntraReadOnlySecretName,
-      );
+      const entraIdToken = await getEntraIdToken({
+        clients: await getAuthorizedClients(),
+        clientId: fastify.environmentConfig.AadValidClientId,
+        secretName: genericConfig.EntraSecretName,
+        encryptionSecret: fastify.secretConfig.encryption_key,
+        logger: request.log,
+      });
       const response = await listGroupMembers(entraIdToken, groupId);
       reply.status(200).send(response);
     },
@@ -572,16 +582,17 @@ No action is required from you at this time.
       onRequest: fastify.authorizeFromSchema,
     },
     async (request, reply) => {
-      const entraIdToken = await getEntraIdToken(
-        await getAuthorizedClients(),
-        fastify.environmentConfig.AadValidClientId,
-        undefined,
-        genericConfig.EntraSecretName,
-      );
+      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 },
+        { redisClient, key, logger: request.log },
       );
       if (redisResponse) {
         request.log.debug("Got manageable groups from Redis cache.");
diff --git a/src/api/routes/membership.ts b/src/api/routes/membership.ts
@@ -51,6 +51,7 @@ const membershipPlugin: FastifyPluginAsync = async (fastify, _options) => {
           region: genericConfig.AwsRegion,
           credentials,
         }),
+        redisClient: fastify.redisClient,
       };
       fastify.log.info(
         `Assumed Entra role ${roleArns.Entra} to get the Entra token.`,
@@ -63,6 +64,7 @@ const membershipPlugin: FastifyPluginAsync = async (fastify, _options) => {
     return {
       smClient: fastify.secretsManagerClient,
       dynamoClient: fastify.dynamoClient,
+      redisClient: fastify.redisClient,
     };
   };
   const limitedRoutes: FastifyPluginAsync = async (fastify) => {
@@ -106,10 +108,13 @@ const membershipPlugin: FastifyPluginAsync = async (fastify, _options) => {
             message: `${netId} is already a paid member!`,
           });
         }
-        const entraIdToken = await getEntraIdToken(
-          await getAuthorizedClients(),
-          fastify.environmentConfig.AadValidClientId,
-        );
+        const entraIdToken = await getEntraIdToken({
+          clients: await getAuthorizedClients(),
+          clientId: fastify.environmentConfig.AadValidClientId,
+          secretName: genericConfig.EntraSecretName,
+          encryptionSecret: fastify.secretConfig.encryption_key,
+          logger: request.log,
+        });
         const paidMemberGroup = fastify.environmentConfig.PaidMemberGroupId;
         const isAadMember = await checkPaidMembershipFromEntra(
           netId,
@@ -224,10 +229,13 @@ const membershipPlugin: FastifyPluginAsync = async (fastify, _options) => {
             .header("X-ACM-Data-Source", "dynamo")
             .send({ netId, isPaidMember: true });
         }
-        const entraIdToken = await getEntraIdToken(
-          await getAuthorizedClients(),
-          fastify.environmentConfig.AadValidClientId,
-        );
+        const entraIdToken = await getEntraIdToken({
+          clients: await getAuthorizedClients(),
+          clientId: fastify.environmentConfig.AadValidClientId,
+          secretName: genericConfig.EntraSecretName,
+          encryptionSecret: fastify.secretConfig.encryption_key,
+          logger: request.log,
+        });
         const paidMemberGroup = fastify.environmentConfig.PaidMemberGroupId;
         const isAadMember = await checkPaidMembershipFromEntra(
           netId,
diff --git a/src/api/sqs/handlers/emailMembershipPassHandler.ts b/src/api/sqs/handlers/emailMembershipPassHandler.ts
@@ -4,23 +4,37 @@ import {
   runEnvironment,
   SQSHandlerFunction,
 } from "../index.js";
-import { environmentConfig, genericConfig } from "common/config.js";
-import { getAuthorizedClients } from "../utils.js";
+import {
+  environmentConfig,
+  genericConfig,
+  SecretConfig,
+} from "common/config.js";
+import { getAuthorizedClients, getSecretConfig } from "../utils.js";
 import { getEntraIdToken, getUserProfile } from "api/functions/entraId.js";
 import { issueAppleWalletMembershipCard } from "api/functions/mobileWallet.js";
 import { generateMembershipEmailCommand } from "api/functions/ses.js";
 import { SESClient } from "@aws-sdk/client-ses";
+import RedisModule from "ioredis";
+
+let secretConfig: SecretConfig;
 
 export const emailMembershipPassHandler: SQSHandlerFunction<
   AvailableSQSFunctions.EmailMembershipPass
 > = async (payload, metadata, logger) => {
   const email = payload.email;
   const commonConfig = { region: genericConfig.AwsRegion };
   const clients = await getAuthorizedClients(logger, commonConfig);
-  const entraIdToken = await getEntraIdToken(
-    clients,
-    currentEnvironmentConfig.AadValidClientId,
-  );
+  if (!secretConfig) {
+    secretConfig = await getSecretConfig({ logger, commonConfig });
+  }
+  const redisClient = new RedisModule.default(secretConfig.redis_url);
+  const entraIdToken = await getEntraIdToken({
+    clients: { ...clients, redisClient },
+    clientId: currentEnvironmentConfig.AadValidClientId,
+    secretName: genericConfig.EntraSecretName,
+    encryptionSecret: secretConfig.encryption_key,
+    logger,
+  });
   const userProfile = await getUserProfile(entraIdToken, email);
   const pkpass = await issueAppleWalletMembershipCard(
     clients,
diff --git a/src/api/sqs/handlers/provisionNewMember.ts b/src/api/sqs/handlers/provisionNewMember.ts
diff --git a/src/api/sqs/index.ts b/src/api/sqs/index.ts
diff --git a/src/api/sqs/utils.ts b/src/api/sqs/utils.ts
