Use central redis cache function for IAM (#183)

Also enables optional encryption for redis cache content (ex: storing
tokens)

Author: devksingh4

src/api/functions/apiKey.ts

@@ -7,8 +7,7 @@ import {
   DynamoDBClient,
   GetItemCommand,
 } from "@aws-sdk/client-dynamodb";
-import { genericConfig } from "common/config.js";
-import { AUTH_DECISION_CACHE_SECONDS as API_KEY_DATA_CACHE_SECONDS } from "./authorization.js";
+import { genericConfig, GENERIC_CACHE_SECONDS } from "common/config.js";
 import { unmarshall } from "@aws-sdk/util-dynamodb";
 import { ApiKeyMaskedEntry, DecomposedApiKey } from "common/types/apiKey.js";
 import { AvailableAuthorizationPolicy } from "common/policies/definition.js";
@@ -105,7 +104,7 @@ export const getApiKeyData = async ({
   });
   const result = await dynamoClient.send(getCommand);
   if (!result || !result.Item) {
-    nodeCache.set(cacheKey, null, API_KEY_DATA_CACHE_SECONDS);
+    nodeCache.set(cacheKey, null, GENERIC_CACHE_SECONDS);
     return undefined;
   }
   const unmarshalled = unmarshall(result.Item) as ApiKeyDynamoEntry;
@@ -124,7 +123,7 @@ export const getApiKeyData = async ({
   if (!("keyHash" in unmarshalled)) {
     return undefined; // bad data, don't cache it
   }
-  let cacheTime = API_KEY_DATA_CACHE_SECONDS;
+  let cacheTime = GENERIC_CACHE_SECONDS;
   if (unmarshalled.expiresAt) {
     const currentEpoch = Date.now();
     cacheTime = min(cacheTime, unmarshalled.expiresAt - currentEpoch);

src/api/functions/authorization.ts

@@ -3,9 +3,6 @@ import { unmarshall } from "@aws-sdk/util-dynamodb";
 import { genericConfig } from "../../common/config.js";
 import { DatabaseFetchError } from "../../common/errors/index.js";
 import { allAppRoles, AppRoles } from "../../common/roles.js";
-import { FastifyInstance } from "fastify";
-
-export const AUTH_DECISION_CACHE_SECONDS = 180;
 
 export async function getUserRoles(
   dynamoClient: DynamoDBClient,

src/api/functions/encryption.ts

@@ -0,0 +1,81 @@
+import { DecryptionError } from "common/errors/index.js";
+import crypto, { createDecipheriv, pbkdf2Sync } from "node:crypto";
+
+const VALID_PREFIX = "VALID:";
+const ITERATIONS = 100000;
+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 const CORRUPTED_DATA_MESSAGE = "Encrypted data is corrupted.";
+
+export function encrypt({
+  plaintext,
+  encryptionSecret,
+}: {
+  plaintext: string;
+  encryptionSecret: string;
+}) {
+  const salt = crypto.randomBytes(16);
+  const iv = crypto.randomBytes(12);
+  const key = crypto.pbkdf2Sync(
+    encryptionSecret,
+    salt,
+    ITERATIONS,
+    KEY_LEN,
+    HASH_FUNCTION,
+  );
+  const cipher = crypto.createCipheriv(ALGORITHM, key, iv);
+  const encrypted = Buffer.concat([
+    cipher.update(`${VALID_PREFIX}${plaintext}`, "utf8"),
+    cipher.final(),
+  ]);
+  const tag = cipher.getAuthTag();
+  return Buffer.concat([salt, iv, tag, encrypted]).toString("hex");
+}
+
+export function decrypt({
+  cipherText,
+  encryptionSecret,
+}: {
+  cipherText: string;
+  encryptionSecret: string;
+}): string {
+  const data = Buffer.from(cipherText, "hex");
+  const salt = data.subarray(0, 16);
+  const iv = data.subarray(16, 28);
+  const tag = data.subarray(28, 44);
+  const encryptedText = data.subarray(44);
+
+  const key = pbkdf2Sync(
+    encryptionSecret,
+    salt,
+    ITERATIONS,
+    KEY_LEN,
+    HASH_FUNCTION,
+  );
+  let decipher, decryptedBuffer;
+  try {
+    decipher = createDecipheriv(ALGORITHM, key, iv);
+    decipher.setAuthTag(tag);
+    decryptedBuffer = Buffer.concat([
+      decipher.update(encryptedText),
+      decipher.final(),
+    ]);
+  } catch (e) {
+    throw new DecryptionError({
+      message: INVALID_DECRYPTION_MESSAGE,
+    });
+  }
+
+  const candidate = decryptedBuffer.toString("utf8");
+  if (candidate.substring(0, VALID_PREFIX.length) !== VALID_PREFIX) {
+    throw new DecryptionError({
+      message: CORRUPTED_DATA_MESSAGE,
+    });
+  }
+  return candidate.substring(VALID_PREFIX.length, candidate.length);
+}

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) {

src/api/functions/redisCache.ts

@@ -1,34 +1,92 @@
-import { type Redis } from "ioredis";
+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 async function getRedisKey<T>({
+export type GetFromCacheInput = {
+  redisClient: RedisModule.default;
+  key: string;
+  encryptionSecret?: string;
+  logger: pino.Logger | FastifyBaseLogger;
+};
+
+export type SetInCacheInput = {
+  redisClient: RedisModule.default;
+  key: string;
+  data: string;
+  expiresIn?: number;
+  encryptionSecret?: string;
+};
+
+const redisEntrySchema = z.object({
+  isEncrypted: z.boolean(),
+  data: z.string(),
+});
+
+export async function getKey<T extends object>({
   redisClient,
   key,
-  parseJson = false,
-}: {
-  redisClient: Redis;
-  key: string;
-  parseJson?: boolean;
-}) {
-  const resp = await redisClient.get(key);
-  if (!resp) {
+  encryptionSecret,
+  logger,
+}: GetFromCacheInput): Promise<T | null> {
+  const data = await redisClient.get(key);
+  if (!data) {
     return null;
   }
-  return parseJson ? (JSON.parse(resp) as T) : (resp as string);
+  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;
+  }
 }
 
-export async function setRedisKey({
+export async function setKey({
   redisClient,
   key,
-  value,
-  expiresSec,
-}: {
-  redisClient: Redis;
-  key: string;
-  value: string;
-  expiresSec?: number;
-}) {
-  if (expiresSec) {
-    return await redisClient.set(key, value, "EX", expiresSec);
-  }
-  return await redisClient.set(key, value);
+  encryptionSecret,
+  data,
+  expiresIn,
+}: SetInCacheInput) {
+  const realData = encryptionSecret
+    ? encrypt({ plaintext: data, encryptionSecret })
+    : data;
+  const redisPayload: z.infer<typeof redisEntrySchema> = {
+    isEncrypted: !!encryptionSecret,
+    data: realData,
+  };
+  const strRedisPayload = JSON.stringify(redisPayload);
+  return expiresIn
+    ? await redisClient.set(key, strRedisPayload, "EX", expiresIn)
+    : await redisClient.set(key, strRedisPayload);
 }