route additional api keys separately

Author: carderne

apps/webapp/app/models/runtimeEnvironment.server.ts

@@ -6,6 +6,7 @@ import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.
 import { logger } from "~/services/logger.server";
 import { getUsername } from "~/utils/username";
 import { hashApiKey } from "~/utils/apiKeys";
+import { isAdditionalApiKey } from "@trigger.dev/core/v3/apiKeys";
 import { isDefaultDevBranch, sanitizeBranchName } from "@trigger.dev/core/v3/utils/gitBranch";
 import { scopesGrantFullAccess } from "@trigger.dev/rbac";
 
@@ -125,16 +126,18 @@ async function resolveEnvironmentByApiKey(
   } satisfies Prisma.RuntimeEnvironmentInclude;
 
   const now = new Date();
-  let additionalApiKey: { id: string; lastUsedAt: Date | null } | null = null;
-  let environment = await tx.runtimeEnvironment.findFirst({
-    where: {
-      apiKey,
-    },
-    include,
-  });
+  const routesToAdditionalKey = isAdditionalApiKey(apiKey);
+  let rootEnvironment = routesToAdditionalKey
+    ? null
+    : await tx.runtimeEnvironment.findFirst({
+        where: {
+          apiKey,
+        },
+        include,
+      });
 
   // Fall back to root keys that were rotated within the grace window.
-  if (!environment) {
+  if (!routesToAdditionalKey && !rootEnvironment) {
     const revokedApiKey = await tx.revokedApiKey.findFirst({
       where: {
         apiKey,
@@ -145,34 +148,34 @@ async function resolveEnvironmentByApiKey(
       },
     });
 
-    environment = revokedApiKey?.runtimeEnvironment ?? null;
+    rootEnvironment = revokedApiKey?.runtimeEnvironment ?? null;
   }
 
   // Additional keys are host-owned credentials. Legacy routes cannot apply a
   // scoped ability, so only an explicit full-access scope is accepted.
-  if (!environment) {
-    const match = await tx.apiKey.findFirst({
-      where: {
-        keyHash: hashApiKey(apiKey),
-        revokedAt: null,
-        OR: [{ expiresAt: null }, { expiresAt: { gt: now } }],
-      },
-      select: {
-        id: true,
-        lastUsedAt: true,
-        scopes: true,
-        runtimeEnvironment: { include },
-      },
-    });
-
-    if (match && !scopesGrantFullAccess(match.scopes)) {
-      return { ok: false, reason: "restricted" };
-    }
+  const match = routesToAdditionalKey
+    ? await tx.apiKey.findFirst({
+        where: {
+          keyHash: hashApiKey(apiKey),
+          revokedAt: null,
+          OR: [{ expiresAt: null }, { expiresAt: { gt: now } }],
+        },
+        select: {
+          id: true,
+          lastUsedAt: true,
+          scopes: true,
+          runtimeEnvironment: { include },
+        },
+      })
+    : null;
 
-    additionalApiKey = match ? { id: match.id, lastUsedAt: match.lastUsedAt } : null;
-    environment = match?.runtimeEnvironment ?? null;
+  if (match && !scopesGrantFullAccess(match.scopes)) {
+    return { ok: false, reason: "restricted" };
   }
 
+  const additionalApiKey = match ? { id: match.id, lastUsedAt: match.lastUsedAt } : null;
+  let environment = rootEnvironment ?? match?.runtimeEnvironment ?? null;
+
   if (!environment) {
     return { ok: false, reason: "not-found" };
   }

apps/webapp/test/findEnvironmentByApiKey.test.ts

@@ -1,6 +1,6 @@
 import { postgresTest } from "@internal/testcontainers";
 import { type PrismaClient } from "@trigger.dev/database";
-import { describe, expect, vi } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { findEnvironmentByApiKey } from "~/models/runtimeEnvironment.server";
 import { generateAdditionalApiKey, hashApiKey } from "~/utils/apiKeys";
 import { createTestOrgProjectWithMember, uniqueId } from "./fixtures/environmentVariablesFixtures";
@@ -162,6 +162,43 @@ describe("findEnvironmentByApiKey — non-branchable", () => {
     const resolved = await findEnvironmentByApiKey("tr_dev_nonexistent", undefined, prisma);
     expect(resolved).toBeNull();
   });
+
+  it("queries only the additional-key store for a valid additional-key format", async () => {
+    const runtimeEnvironmentFind = vi.fn();
+    const revokedApiKeyFind = vi.fn();
+    const apiKeyFind = vi.fn(async () => null);
+    const tx = {
+      runtimeEnvironment: { findFirst: runtimeEnvironmentFind },
+      revokedApiKey: { findFirst: revokedApiKeyFind },
+      apiKey: { findFirst: apiKeyFind },
+    } as unknown as PrismaClient;
+
+    await expect(
+      findEnvironmentByApiKey("tr_prod_sk_0123456789abcdefghijklmn", undefined, tx)
+    ).resolves.toBeNull();
+    expect(apiKeyFind).toHaveBeenCalledOnce();
+    expect(runtimeEnvironmentFind).not.toHaveBeenCalled();
+    expect(revokedApiKeyFind).not.toHaveBeenCalled();
+  });
+
+  it.each(["tr_prod_ak_0123456789abcdefghijklmn", "tr_prod_sk_too-short"])(
+    "keeps malformed additional-key formats on the root lookup path: %s",
+    async (apiKey) => {
+      const runtimeEnvironmentFind = vi.fn(async () => null);
+      const revokedApiKeyFind = vi.fn(async () => null);
+      const apiKeyFind = vi.fn();
+      const tx = {
+        runtimeEnvironment: { findFirst: runtimeEnvironmentFind },
+        revokedApiKey: { findFirst: revokedApiKeyFind },
+        apiKey: { findFirst: apiKeyFind },
+      } as unknown as PrismaClient;
+
+      await expect(findEnvironmentByApiKey(apiKey, undefined, tx)).resolves.toBeNull();
+      expect(runtimeEnvironmentFind).toHaveBeenCalledOnce();
+      expect(revokedApiKeyFind).toHaveBeenCalledOnce();
+      expect(apiKeyFind).not.toHaveBeenCalled();
+    }
+  );
 });
 
 describe("findEnvironmentByApiKey — additional and disabled keys", () => {

internal-packages/rbac/src/apiKeyPolicies.test.ts

@@ -19,6 +19,8 @@ type LazyControllerInternals = {
 };
 
 const prismaPlaceholder = {} as PrismaClient;
+const ADDITIONAL_API_KEY = "tr_prod_sk_0123456789abcdefghijklmn";
+const ROOT_API_KEY = "tr_prod_0123456789abcdefghijklmn";
 const environment = {
   id: "env_123",
   organizationId: "org_123",
@@ -53,6 +55,20 @@ function additionalKeyResult(scopes: string[]): AuthSuccess {
   };
 }
 
+function rootKeyResult(): AuthSuccess {
+  return {
+    ok: true,
+    environment,
+    subject: {
+      type: "user",
+      userId: "user_123",
+      organizationId: environment.organizationId,
+      projectId: environment.projectId,
+    },
+    ability: buildJwtAbility(["admin"]),
+  };
+}
+
 function publicJwtResult(): AuthSuccess {
   return {
     ok: true,
@@ -71,6 +87,12 @@ function publicJwt(payload: Record<string, unknown>) {
   return `header.${Buffer.from(JSON.stringify(payload)).toString("base64url")}.signature`;
 }
 
+function bearerRequest(token: string) {
+  return new Request("https://api.trigger.dev/test", {
+    headers: { Authorization: `Bearer ${token}` },
+  });
+}
+
 describe("API-key policy controller composition", () => {
   it("routes public JWTs directly to the host without calling the plugin authenticator", async () => {
     const pluginAuthenticate = vi.fn();
@@ -94,35 +116,126 @@ describe("API-key policy controller composition", () => {
     expect(pluginAuthenticate).not.toHaveBeenCalled();
   });
 
-  it("uses the scoped host result unchanged when plugin auth falls back to an additional key", async () => {
-    const pluginAuthenticate = vi.fn(async () => ({
-      ok: false as const,
-      status: 401 as const,
-      error: "Invalid API key",
-    }));
+  it("routes valid additional keys directly to the host and preserves their scopes", async () => {
+    const pluginAuthenticate = vi.fn();
+    const hostAuthenticate = vi.fn(async () => additionalKeyResult(["write:tasks:send-email"]));
     const plugin = {
       isUsingPlugin: vi.fn(async () => true),
       authenticateBearer: pluginAuthenticate,
     } as unknown as RoleBaseAccessController;
-    const controller = installPlugin(
-      plugin,
-      vi.fn(async () => additionalKeyResult(["write:tasks:send-email"]))
-    );
+    const controller = installPlugin(plugin, hostAuthenticate);
 
-    const result = await controller.authenticateBearer(
-      new Request("https://api.trigger.dev/test", {
-        headers: { Authorization: "Bearer tr_additional" },
-      })
-    );
+    const result = await controller.authenticateBearer(bearerRequest(ADDITIONAL_API_KEY));
 
     expect(result.ok).toBe(true);
+    expect(hostAuthenticate).toHaveBeenCalledOnce();
+    expect(pluginAuthenticate).not.toHaveBeenCalled();
     if (!result.ok) return;
     expect(result.subject).toMatchObject({ type: "apiKey", restricted: true });
     expect(result.ability.can("trigger", { type: "tasks", id: "send-email" })).toBe(true);
     expect(result.ability.can("trigger", { type: "tasks", id: "other-task" })).toBe(false);
     expect(result.ability.can("read", { type: "runs" })).toBe(false);
   });
 
+  it("returns an unknown additional key's host 401 without calling the plugin", async () => {
+    const hostFailure = {
+      ok: false as const,
+      status: 401 as const,
+      error: "Invalid API key",
+    };
+    const pluginAuthenticate = vi.fn();
+    const hostAuthenticate = vi.fn(async () => hostFailure);
+    const plugin = {
+      isUsingPlugin: vi.fn(async () => true),
+      authenticateBearer: pluginAuthenticate,
+    } as unknown as RoleBaseAccessController;
+    const controller = installPlugin(plugin, hostAuthenticate);
+
+    await expect(controller.authenticateBearer(bearerRequest(ADDITIONAL_API_KEY))).resolves.toEqual(
+      hostFailure
+    );
+    expect(hostAuthenticate).toHaveBeenCalledOnce();
+    expect(pluginAuthenticate).not.toHaveBeenCalled();
+  });
+
+  it("routes current root keys to the plugin without calling the host", async () => {
+    const pluginAuthenticate = vi.fn(async () => rootKeyResult());
+    const hostAuthenticate = vi.fn();
+    const plugin = {
+      isUsingPlugin: vi.fn(async () => true),
+      authenticateBearer: pluginAuthenticate,
+    } as unknown as RoleBaseAccessController;
+    const controller = installPlugin(plugin, hostAuthenticate);
+
+    await expect(controller.authenticateBearer(bearerRequest(ROOT_API_KEY))).resolves.toMatchObject(
+      {
+        ok: true,
+        subject: { type: "user" },
+      }
+    );
+    expect(pluginAuthenticate).toHaveBeenCalledOnce();
+    expect(hostAuthenticate).not.toHaveBeenCalled();
+  });
+
+  it.each([401, 403] as const)(
+    "returns a plugin %s for root-shaped keys without calling the host",
+    async (status) => {
+      const pluginFailure = {
+        ok: false as const,
+        status,
+        error: "Unauthorized",
+      };
+      const pluginAuthenticate = vi.fn(async () => pluginFailure);
+      const hostAuthenticate = vi.fn();
+      const plugin = {
+        isUsingPlugin: vi.fn(async () => true),
+        authenticateBearer: pluginAuthenticate,
+      } as unknown as RoleBaseAccessController;
+      const controller = installPlugin(plugin, hostAuthenticate);
+
+      await expect(controller.authenticateBearer(bearerRequest(ROOT_API_KEY))).resolves.toEqual(
+        pluginFailure
+      );
+      expect(pluginAuthenticate).toHaveBeenCalledOnce();
+      expect(hostAuthenticate).not.toHaveBeenCalled();
+    }
+  );
+
+  it("fails closed when an additional-key host route resolves a non-apiKey subject", async () => {
+    const pluginAuthenticate = vi.fn();
+    const hostAuthenticate = vi.fn(async () => rootKeyResult());
+    const plugin = {
+      isUsingPlugin: vi.fn(async () => true),
+      authenticateBearer: pluginAuthenticate,
+    } as unknown as RoleBaseAccessController;
+    const controller = installPlugin(plugin, hostAuthenticate);
+
+    await expect(controller.authenticateBearer(bearerRequest(ADDITIONAL_API_KEY))).resolves.toEqual(
+      {
+        ok: false,
+        status: 401,
+        error: "Invalid API key",
+      }
+    );
+    expect(pluginAuthenticate).not.toHaveBeenCalled();
+  });
+
+  it("keeps additional-key authentication in the no-plugin fallback controller", async () => {
+    const fallbackAuthenticate = vi.fn(async () => additionalKeyResult(["admin"]));
+    const hostAuthenticate = vi.fn();
+    const fallback = {
+      isUsingPlugin: vi.fn(async () => false),
+      authenticateBearer: fallbackAuthenticate,
+    } as unknown as RoleBaseAccessController;
+    const controller = installPlugin(fallback, hostAuthenticate);
+
+    await expect(
+      controller.authenticateBearer(bearerRequest(ADDITIONAL_API_KEY))
+    ).resolves.toMatchObject({ ok: true, subject: { type: "apiKey" } });
+    expect(fallbackAuthenticate).toHaveBeenCalledOnce();
+    expect(hostAuthenticate).not.toHaveBeenCalled();
+  });
+
   it("delegates API-key policy catalogue, preparation, and description", async () => {
     const presets = vi.fn(async () => []);
     const prepare = vi.fn(async () => ({