From d58b5fc4b6dffd1594dd2e743379bb76433fd12b Mon Sep 17 00:00:00 2001
From: John Smith <john@johnjcsmith.com>
Date: Wed, 11 Dec 2024 08:16:13 +1030
Subject: [PATCH] chore: Cache customer auth context

---
 control-plane/src/modules/auth/api-secret.ts    | 16 ++++++----------
 control-plane/src/modules/auth/auth.ts          |  1 +
 control-plane/src/modules/auth/customer-auth.ts | 17 ++++++++++++++++-
 control-plane/src/utilities/cache.ts            |  6 ++++++
 4 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/control-plane/src/modules/auth/api-secret.ts b/control-plane/src/modules/auth/api-secret.ts
index 0c5cbf6c..239f77fa 100644
--- a/control-plane/src/modules/auth/api-secret.ts
+++ b/control-plane/src/modules/auth/api-secret.ts
@@ -1,21 +1,17 @@
 import * as data from "../data";
 import { eq, and, isNull } from "drizzle-orm";
-import { createHash, randomBytes } from "crypto";
+import { randomBytes } from "crypto";
 import { logger } from "../observability/logger";
-import { createCache } from "../../utilities/cache";
+import { createCache, hashFromSecret } from "../../utilities/cache";
 
-const authContextCache = createCache<{
+const apiKeyContextCache = createCache<{
   clusterId: string;
   id: string;
   organizationId: string;
 }>(
-  Symbol("authContextCach"),
+  Symbol("apiKeyContextCache"),
 );
 
-const hashFromSecret = (secret: string): string => {
-  return createHash("sha256").update(secret).digest("hex");
-};
-
 export const isApiSecret = (authorization: string): boolean =>
   authorization.startsWith("sk_");
 
@@ -26,7 +22,7 @@ export const verifyApiKey = async (
 > => {
   const secretHash = hashFromSecret(secret);
 
-  const cached = await authContextCache.get(secretHash);
+  const cached = await apiKeyContextCache.get(secretHash);
 
   if (cached) {
     return cached;
@@ -60,7 +56,7 @@ export const verifyApiKey = async (
     return undefined;
   }
 
-  await authContextCache.set(secretHash, {
+  await apiKeyContextCache.set(secretHash, {
     clusterId: result.clusterId,
     id: result.id,
     organizationId: result.organizationId,
diff --git a/control-plane/src/modules/auth/auth.ts b/control-plane/src/modules/auth/auth.ts
index c904223a..c02a8470 100644
--- a/control-plane/src/modules/auth/auth.ts
+++ b/control-plane/src/modules/auth/auth.ts
@@ -373,6 +373,7 @@ export const extractCustomerAuthState = async (
   if (!cluster.enable_customer_auth) {
     throw new AuthenticationError(
       "Customer auth is not enabled for this cluster",
+      "https://docs.inferable.ai/pages/auth#customer-provided-secrets"
     );
   }
 
diff --git a/control-plane/src/modules/auth/customer-auth.ts b/control-plane/src/modules/auth/customer-auth.ts
index 98baf325..ebeeb357 100644
--- a/control-plane/src/modules/auth/customer-auth.ts
+++ b/control-plane/src/modules/auth/customer-auth.ts
@@ -7,11 +7,16 @@ import { packer } from "../packer";
 import * as jobs from "../jobs/jobs";
 import { getJobStatusSync } from "../jobs/jobs";
 import { getServiceDefinition } from "../service-definitions";
+import { createCache, hashFromSecret } from "../../utilities/cache";
 
 export const VERIFY_FUNCTION_NAME = "handleCustomerAuth";
 export const VERIFY_FUNCTION_SERVICE = "default";
 const VERIFY_FUNCTION_ID = `${VERIFY_FUNCTION_SERVICE}_${VERIFY_FUNCTION_NAME}`;
 
+const customerAuthContextCache = createCache<unknown>(
+  Symbol("customerAuthContextCache"),
+);
+
 /**
  * Calls the customer provided verify function and returns the result
  */
@@ -22,6 +27,14 @@ export const verifyCustomerProvidedAuth = async ({
   token: string;
   clusterId: string;
 }): Promise<unknown> => {
+
+  const secretHash = hashFromSecret(token);
+
+  const cached = await customerAuthContextCache.get(secretHash);
+  if (cached) {
+    return cached;
+  }
+
   try {
     const serviceDefinition = await getServiceDefinition({
       service: VERIFY_FUNCTION_SERVICE,
@@ -56,7 +69,7 @@ export const verifyCustomerProvidedAuth = async ({
     const result = await getJobStatusSync({
       jobId: id,
       owner: { clusterId },
-      ttl: 5_000,
+      ttl: 15_000,
     });
 
     if (
@@ -69,6 +82,8 @@ export const verifyCustomerProvidedAuth = async ({
       );
     }
 
+    await customerAuthContextCache.set(secretHash, result, 300);
+
     return packer.unpack(result.result);
   } catch (e) {
     if (e instanceof JobPollTimeoutError) {
diff --git a/control-plane/src/utilities/cache.ts b/control-plane/src/utilities/cache.ts
index 2ca4bf02..810c9b7d 100644
--- a/control-plane/src/utilities/cache.ts
+++ b/control-plane/src/utilities/cache.ts
@@ -1,3 +1,4 @@
+import { createHash } from "crypto";
 import NodeCache from "node-cache";
 import { redisClient } from "../modules/redis";
 
@@ -37,3 +38,8 @@ export const createCache = <T>(namespace: symbol) => {
 // const cache = createCache(Symbol("cache"));
 // cache.set("key", "value");
 // const value = cache.get("key");
+
+export const hashFromSecret = (secret: string): string => {
+  return createHash("sha256").update(secret).digest("hex");
+};
+