feat: encapsulate keys

apis/cloudflare/.gitignore

@@ -172,3 +172,5 @@ dist
 
 .dev.vars
 wrangler.toml
+
+certs/

apis/cloudflare/e2e/auth.test.ts

@@ -0,0 +1,72 @@
+import { createRandomToken } from "./utils/testAuth";
+import { expect, test } from "bun:test";
+import { PROXY, TOKEN } from "./utils/constants";
+
+const COMMON_HEADERS = {
+  "Content-Type": "application/json",
+  "X-Grit-Api": TOKEN,
+};
+
+const OPENAI_HEADERS_MISSING_GRIT_KEY = {
+  "Content-Type": "application/json",
+};
+
+test("auth__failsWhenMissingGritKey", async () => {
+  const res = await fetch(PROXY, {
+    headers: OPENAI_HEADERS_MISSING_GRIT_KEY,
+    method: "POST",
+    body: JSON.stringify({
+      model: "gpt-3.5-turbo",
+      messages: [
+        {
+          role: "user",
+          content: "What is a proxy?",
+        },
+      ],
+      seed: 8,
+    }),
+  });
+
+  expect(res.status).toBe(400);
+});
+
+test("auth__failsWhenInvalidGritKey", async () => {
+  const res = await fetch(PROXY, {
+    headers: { ...OPENAI_HEADERS_MISSING_GRIT_KEY, "X-Grit-Api": "junk" },
+    method: "POST",
+    body: JSON.stringify({
+      model: "gpt-3.5-turbo",
+      messages: [
+        {
+          role: "user",
+          content: "What is a proxy?",
+        },
+      ],
+      seed: 1,
+    }),
+  });
+
+  expect(res.status).toBe(401);
+});
+
+test("auth__succeedsWhenProvidedGritKey", async () => {
+  const res = await fetch(PROXY, {
+    headers: COMMON_HEADERS,
+    method: "POST",
+    body: JSON.stringify({
+      model: "gpt-3.5-turbo",
+      messages: [
+        {
+          role: "user",
+          content: "What is a proxy?",
+        },
+      ],
+      seed: 1,
+    }),
+  });
+
+  expect(res.status).toBe(200);
+
+  const antCompletions = await res.json();
+  expect(antCompletions.choices.length >= 1).toBe(true);
+});

apis/cloudflare/e2e/basic.test.ts

@@ -0,0 +1,50 @@
+import { PROXY, TOKEN } from "./utils/constants";
+import { createRandomToken } from "./utils/testAuth";
+import { expect, test } from "bun:test";
+
+const COMMON_HEADERS = {
+  "Content-Type": "application/json",
+  "X-Grit-Api": TOKEN,
+};
+
+test("basic__routesBetweenModels", async () => {
+  const res = await fetch(PROXY, {
+    headers: COMMON_HEADERS,
+    method: "POST",
+    body: JSON.stringify({
+      model: "gpt-3.5-turbo",
+      messages: [
+        {
+          role: "user",
+          content: "What is a proxy?",
+        },
+      ],
+      seed: 1,
+    }),
+  });
+
+  expect(res.status).toBe(200);
+
+  const openAiCompletions = await res.json();
+  expect(openAiCompletions.choices.length >= 1).toBe(true);
+
+  const antRes = await fetch(PROXY, {
+    headers: COMMON_HEADERS,
+    method: "POST",
+    body: JSON.stringify({
+      model: "claude-2.1",
+      messages: [
+        {
+          role: "user",
+          content: "What is a proxy?",
+        },
+      ],
+      seed: 1,
+    }),
+  });
+
+  expect(antRes.status).toBe(200);
+
+  const antCompletions = await antRes.json();
+  expect(antCompletions.choices.length >= 1).toBe(true);
+});

apis/cloudflare/e2e/utils/constants.ts

@@ -0,0 +1,6 @@
+const ROOT = process.env.PROXY ?? "localhost:62952";
+
+export const PROXY = `${ROOT}/v1/chat/completions`;
+
+export const TOKEN =
+  "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTUxNjIzOTAyMiwiaXNzIjoiZ3JpdC1sbG0tcm91dGVyIn0.Hx2Z14WQdNB_bG8BwApb3WVKfReQo8krOPxtsTHhjgTjV45pMo197uRy0SdgFHfVg0_t0h705XNgF5zBv-DXqsiqBagNr9yip8vSKGltSAd04PEiFsRUlLxaWuMqHfBEhfeSA0x46u14UytOBV38auz4FPC3IzV_E7_a5fs6X2MGUCm1bALp9vZh5KBHNBxUpb7yNT5T5f5pzxNP7qyMdttEndSJgkLZbjpFOyZqAkgs59MyqXxlasxnk5aWAjEn7S4HPkLhB8btTAwwjRXQpnDZQ7xjtrA9bB2ZnQIq8g6xpI6J_vmHZkLskednoRHOSRt4_BkpOsGxtshn72UHoVqWNZTbspvqZANQlB5CnqiAqO2mlmUk4Mrk4Gg89BioMhYCDOUps-qlyDuXeRSMZn2H81BYh6oT-dai-qE5zKWy2Z1lOAdwWzgndn0qBtPvUvyIqe7wqGtNLEvRccRIMGuSv68GcGWuxhE9rRK1w1eOcM5I7I-sYlELbhX_qJ147azxiFjULzJDS8jDge0npMDUJVwpX_BAaG3A2qLQhZMFduQrCbfM4NqlKQ0qTb6-M1I6dOv-zg7A0Q3jcCLFG0bFZczMACmutYME-m44EK7ruIvc6pWom9bN2wtjsqoXS0FCtkji55MGA_hBGVmB8qgOAHg_JxYssuz2tBnv6e4";

apis/cloudflare/package.json

@@ -8,19 +8,23 @@
     "dev": "wrangler dev --port 8787 --inspector-port 9299",
     "start": "wrangler dev",
     "watch": "tsup --watch --dts",
-    "build": "tsup --clean --dts"
+    "build": "tsup --clean --dts",
+    "test": "bun --env-file=.env test",
+    "test:prod": "export PROXY=https://proxy.admin-a65.workers.dev && bun --env-file=.env test"
   },
   "devDependencies": {
     "@cloudflare/workers-types": "^4.20240512.0",
+    "@types/jsonwebtoken": "^9.0.6",
     "itty-router": "^3.0.12",
+    "tsup": "^8.0.1",
     "typescript": "^5.0.4",
-    "wrangler": "^3.57.1",
-    "tsup": "^8.0.1"
+    "wrangler": "^3.57.1"
   },
   "dependencies": {
     "@braintrust/proxy": "workspace:*",
     "@opentelemetry/resources": "^1.18.1",
     "@opentelemetry/sdk-metrics": "^1.18.1",
-    "dotenv": "^16.3.1"
+    "dotenv": "^16.3.1",
+    "jose": "^5.6.3"
   }
 }

apis/cloudflare/src/auth.ts

@@ -0,0 +1,30 @@
+import { verify } from "jsonwebtoken";
+import * as jose from "jose";
+
+export async function authenticateToken(
+  token: string,
+  env: Env,
+): Promise<boolean> {
+  if (!env.JWT_PUB_KEY) {
+    throw Error("Expected JWT_PUB_KEY in env");
+  }
+
+  const pubKey = await jose.importSPKI(env.JWT_PUB_KEY, "RS256");
+
+  try {
+    await jose.jwtVerify(token, pubKey, {
+      issuer: "grit-llm-router",
+    });
+  } catch (error) {
+    console.error(`Error verifying token ${error}`);
+    return false;
+  }
+
+  return true;
+}
+
+export function parseGritToken(headers: Headers): string | null {
+  const authHeader = headers.get("X-Grit-Api");
+
+  return authHeader;
+}

apis/cloudflare/src/keyPicker.ts

@@ -0,0 +1,19 @@
+type KeyNames = "OPENAI_API_KEY" | "ANT_API_KEY";
+const modelToKeyName: Map<string, KeyNames> = new Map([
+  ["gpt-3.5-turbo", "OPENAI_API_KEY"],
+  ["claude-2.1", "ANT_API_KEY"],
+]);
+
+export function getProviderKey(modelName: string, env: Env): string | null {
+  if (!modelToKeyName.has(modelName)) {
+    return null;
+  }
+
+  const keyName = modelToKeyName.get(modelName);
+
+  if (!keyName) {
+    return null;
+  }
+
+  return env[keyName] ?? null;
+}

apis/cloudflare/src/proxy.ts

@@ -1,6 +1,8 @@
 import { EdgeProxyV1, FlushingExporter } from "@braintrust/proxy/edge";
 import { NOOP_METER_PROVIDER, initMetrics } from "@braintrust/proxy";
 import { PrometheusMetricAggregator } from "./metric-aggregator";
+import { authenticateToken, parseGritToken } from "./auth";
+import { getProviderKey } from "./keyPicker";
 
 export const proxyV1Prefix = "/v1";
 
@@ -12,6 +14,10 @@ declare global {
     PROMETHEUS_SCRAPE_USER?: string;
     PROMETHEUS_SCRAPE_PASSWORD?: string;
     WHITELISTED_ORIGINS?: string;
+    JWT_PUB_KEY?: string;
+    JWT_SECRET?: string;
+    OPENAI_API_KEY?: string;
+    ANT_API_KEY?: string;
   }
 }
 
@@ -64,6 +70,64 @@ export async function handleProxyV1(
     "cloudflare-metrics",
   );
 
+  const gritToken = parseGritToken(request.headers);
+
+  if (!gritToken) {
+    return new Response("Missing X-Grit-Api Header", {
+      status: 400,
+      headers: {
+        "Content-Type": "text/plain",
+      },
+    });
+  }
+
+  const isAuthed = await authenticateToken(gritToken, env);
+
+  if (!isAuthed) {
+    return new Response("Invalid X-Grit-Api", {
+      status: 401,
+      headers: {
+        "Content-Type": "text/plain",
+      },
+    });
+  }
+
+  const clonedRequest = request.clone();
+
+  const body = await clonedRequest.json();
+
+  if (
+    typeof body !== "object" ||
+    !body ||
+    !("model" in body) ||
+    typeof body.model !== "string"
+  ) {
+    return new Response("Expected model in body", {
+      status: 400,
+      headers: {
+        "Content-Type": "text/plain",
+      },
+    });
+  }
+
+  const providerKey = getProviderKey(body.model, env);
+
+  if (!providerKey) {
+    return new Response(`Model ${body.model} not found`, {
+      status: 404,
+      headers: {
+        "Content-Type": "text/plain",
+      },
+    });
+  }
+
+  const headers = new Headers(request.headers);
+  headers.set("Authorization", `Bearer ${providerKey}`);
+
+  const reqWithKey = new Request(request, {
+    headers,
+  });
+
   const whitelist = originWhitelist(env);
 
   const cacheGetLatency = meter.createHistogram("results_cache_get_latency");
@@ -120,7 +184,7 @@ export async function handleProxyV1(
     braintrustApiUrl: braintrustAppUrl(env).toString(),
     meterProvider,
     whitelist,
-  })(request, ctx);
+  })(reqWithKey, ctx);
 }
 
 export async function handlePrometheusScrape(