fix(webapp): tighten ingestion credential checks

Mono-RevId: 8df337b2de2c4312a39b734e29a840b846628f98

Author: pullfrog[bot]

apps/webapp/app/routes/api.v1.tasks.$taskId.batch.ts

@@ -29,7 +29,17 @@ const { action } = createActionApiRoute(
       resource: (params) => ({ type: "tasks", id: params.taskId }),
     },
   },
-  async ({ params, headers, body, authentication, ability }) => {
+  async ({ params, headers, body, authentication, restrictedApiKey, ability }) => {
+    if (
+      (authentication.type !== "PRIVATE" || restrictedApiKey) &&
+      body.items.some((item) => item.options?.payloadType === "application/store")
+    ) {
+      return json(
+        { error: "Stored payloads require an unrestricted private API key" },
+        { status: 403 }
+      );
+    }
+
     const { taskId } = params;
     const {
       "idempotency-key": idempotencyKey,

apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts

@@ -62,9 +62,15 @@ const { action, loader } = createActionApiRoute(
     },
     corsStrategy: "all",
   },
-  async ({ body, headers, params, authentication, ability }) => {
-    if (authentication.type !== "PRIVATE" && body.options?.payloadType === "application/store") {
-      return json({ error: "Stored payloads require a private API key" }, { status: 403 });
+  async ({ body, headers, params, authentication, restrictedApiKey, ability }) => {
+    if (
+      (authentication.type !== "PRIVATE" || restrictedApiKey) &&
+      body.options?.payloadType === "application/store"
+    ) {
+      return json(
+        { error: "Stored payloads require an unrestricted private API key" },
+        { status: 403 }
+      );
     }
 
     const {

apps/webapp/app/routes/api.v1.tasks.batch.ts

@@ -40,12 +40,15 @@ const { action, loader } = createActionApiRoute(
     },
     corsStrategy: "all",
   },
-  async ({ body, headers, params, authentication }) => {
+  async ({ body, headers, params, authentication, restrictedApiKey }) => {
     if (
-      authentication.type !== "PRIVATE" &&
+      (authentication.type !== "PRIVATE" || restrictedApiKey) &&
       body.items.some((item) => item.options?.payloadType === "application/store")
     ) {
-      return json({ error: "Stored payloads require a private API key" }, { status: 403 });
+      return json(
+        { error: "Stored payloads require an unrestricted private API key" },
+        { status: 403 }
+      );
     }
 
     if (!body.items.length) {

apps/webapp/app/routes/api.v2.tasks.batch.ts

@@ -48,12 +48,15 @@ const { action, loader } = createActionApiRoute(
     },
     corsStrategy: "all",
   },
-  async ({ body, headers, params, authentication, ability }) => {
+  async ({ body, headers, params, authentication, restrictedApiKey, ability }) => {
     if (
-      authentication.type !== "PRIVATE" &&
+      (authentication.type !== "PRIVATE" || restrictedApiKey) &&
       body.items.some((item) => item.options?.payloadType === "application/store")
     ) {
-      return json({ error: "Stored payloads require a private API key" }, { status: 403 });
+      return json(
+        { error: "Stored payloads require an unrestricted private API key" },
+        { status: 403 }
+      );
     }
 
     if (!body.items.length) {

apps/webapp/app/routes/api.v3.batches.$batchId.items.ts

@@ -72,6 +72,8 @@ export async function action({ request, params }: ActionFunctionArgs) {
     return json({ error: authResult.error }, { status: authResult.status });
   }
 
+  const restrictedApiKey = authResult.subject.type === "apiKey" && authResult.subject.restricted;
+
   // Get the request body stream
   const body = request.body;
   if (!body) {
@@ -98,7 +100,7 @@ export async function action({ request, params }: ActionFunctionArgs) {
       streamToAsyncIterable(parsedStream),
       authResult.ability,
       batchId,
-      authResult.subject.type !== "publicJWT"
+      authResult.subject.type !== "publicJWT" && !restrictedApiKey
     );
 
     // Process the stream

apps/webapp/app/services/routeBuilders/apiBuilder.server.ts

@@ -1207,6 +1207,7 @@ type ApiKeyActionHandlerFunction<
     ? z.infer<TBodySchema>
     : undefined;
   authentication: ApiAuthenticationResultSuccess;
+  restrictedApiKey: boolean;
   ability: RbacAbility;
   request: Request;
   resource?: TResource;
@@ -1444,6 +1445,7 @@ export function createActionApiRoute<
             headers: parsedHeaders,
             body: parsedBody,
             authentication: authenticationResult,
+            restrictedApiKey: authResult.restrictedApiKey,
             ability,
             request,
             resource,

apps/webapp/test/api-auth.e2e.test.ts

@@ -10,6 +10,8 @@ import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 import type { TestServer } from "@internal/testcontainers/webapp";
 import { startTestServer } from "@internal/testcontainers/webapp";
 import { generateJWT } from "@trigger.dev/core/v3/jwt";
+import { createHash } from "node:crypto";
+import { generateAdditionalApiKey } from "~/utils/apiKeys";
 import { seedTestEnvironment } from "./helpers/seedTestEnvironment";
 import { seedTestPAT, seedTestUser } from "./helpers/seedTestPAT";
 import { seedTestRun } from "./helpers/seedTestRun";
@@ -146,41 +148,159 @@ describe("Packet download authorization", () => {
     expect(inlineOutput.status).toBe(404);
   });
 
-  it("rejects stored-payload pointers from public JWTs on trigger", async () => {
-    const { environment } = await seedTestEnvironment(server.prisma);
-    const jwt = await generateTestJWT(environment, { scopes: ["write:tasks:test-task"] });
-
-    const res = await server.webapp.fetch("/api/v1/tasks/test-task/trigger", {
-      method: "POST",
-      headers: { Authorization: `Bearer ${jwt}`, "Content-Type": "application/json" },
-      body: JSON.stringify({
-        payload: "run_victim/payload.json",
-        options: { payloadType: "application/store" },
-      }),
-    });
-
-    expect(res.status).toBe(403);
-  });
-
-  it("rejects stored-payload pointers from public JWTs on batch trigger", async () => {
-    const { environment } = await seedTestEnvironment(server.prisma);
-    const jwt = await generateTestJWT(environment, { scopes: ["write:tasks:test-task"] });
+  describe.each([
+    "public JWT",
+    "restricted API key",
+    "root API key",
+    "unrestricted API key",
+  ] as const)("%s stored pointers", (credential) => {
+    it.each([
+      "trigger",
+      "v1 batch",
+      "v2 batch",
+      "v3 items",
+      // This legacy route does not accept public JWTs.
+      ...(credential === "public JWT" ? [] : (["task-scoped v1 batch"] as const)),
+    ] as const)("enforces stored-pointer permission on %s", async (route) => {
+      const { environment, project } = await seedTestEnvironment(server.prisma);
+      const { batchFriendlyId: seededBatchFriendlyId, run } = await seedTestRun(server.prisma, {
+        environmentId: environment.id,
+        projectId: project.id,
+        withBatch: route !== "v3 items",
+      });
+      let batchFriendlyId = seededBatchFriendlyId;
+      if (route === "v3 items") {
+        // Phase 1 initializes both the database record and Redis queue metadata.
+        const response = await server.webapp.fetch("/api/v3/batches", {
+          method: "POST",
+          headers: {
+            Authorization: `Bearer ${environment.apiKey}`,
+            "Content-Type": "application/json",
+          },
+          body: JSON.stringify({ runCount: 1, taskIdentifiers: ["test-task"] }),
+        });
+        expect(response.status).toBe(202);
+        batchFriendlyId = (await response.json()).id;
+      } else {
+        await server.prisma.batchTaskRun.update({
+          where: { id: run.batchId! },
+          data: { status: "PENDING", runCount: 1 },
+        });
+      }
+      // Keep JSON-route controls independent of worker/deployment setup.
+      // The payload gate must run even when a matching run is cached.
+      const idempotencyKey = "stored-payload-authorization";
+      await server.prisma.taskRun.update({
+        where: { id: run.id },
+        data: {
+          idempotencyKey,
+          idempotencyKeyExpiresAt: new Date(Date.now() + 3_600_000),
+        },
+      });
 
-    const res = await server.webapp.fetch("/api/v2/tasks/batch", {
-      method: "POST",
-      headers: { Authorization: `Bearer ${jwt}`, "Content-Type": "application/json" },
-      body: JSON.stringify({
-        items: [
-          {
-            task: "test-task",
-            payload: "run_victim/payload.json",
-            options: { payloadType: "application/store" },
+      let token: string;
+      if (credential === "public JWT") {
+        // Include a batch-wide grant to ensure it cannot bypass the payload gate.
+        token = await generateTestJWT(environment, {
+          scopes: ["write:tasks:test-task", `write:batch:${batchFriendlyId}`],
+        });
+      } else if (credential === "root API key") {
+        token = environment.apiKey;
+      } else {
+        const user = await seedTestUser(server.prisma);
+        token = generateAdditionalApiKey("DEVELOPMENT").apiKey;
+        await server.prisma.apiKey.create({
+          data: {
+            name: "Task-scoped ingestion key",
+            keyHash: createHash("sha256").update(token).digest("hex"),
+            lastFour: token.slice(-4),
+            runtimeEnvironmentId: environment.id,
+            createdByUserId: user.id,
+            presetId: credential === "restricted API key" ? "TRIGGER_ONLY" : null,
+            scopes:
+              credential === "restricted API key"
+                ? ["trigger:tasks:test-task", "batchTrigger:tasks:test-task"]
+                : ["admin"],
           },
-        ],
-      }),
-    });
+        });
+      }
+
+      const item = {
+        task: "test-task",
+        index: 0,
+        payload: "run_other/payload.json",
+        options: { payloadType: "application/store", idempotencyKey },
+      };
+      const path =
+        route === "trigger"
+          ? "/api/v1/tasks/test-task/trigger"
+          : route === "v3 items"
+            ? `/api/v3/batches/${batchFriendlyId}/items`
+            : route === "task-scoped v1 batch"
+              ? "/api/v1/tasks/test-task/batch"
+              : `/api/${route === "v1 batch" ? "v1" : "v2"}/tasks/batch`;
+      const res = await server.webapp.fetch(path, {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${token}`,
+          "idempotency-key": idempotencyKey,
+          "Content-Type": route === "v3 items" ? "application/x-ndjson" : "application/json",
+        },
+        body:
+          route === "v3 items"
+            ? `${JSON.stringify(item)}\n`
+            : JSON.stringify(route === "trigger" ? item : { items: [item] }),
+      });
 
-    expect(res.status).toBe(403);
+      if (credential === "public JWT" || credential === "restricted API key") {
+        expect(res.status).toBe(403);
+        expect(await res.json()).toEqual({
+          error:
+            route === "v3 items"
+              ? "Unauthorized"
+              : "Stored payloads require an unrestricted private API key",
+        });
+        expect(
+          await server.prisma.taskRun.count({ where: { runtimeEnvironmentId: environment.id } })
+        ).toBe(1);
+        expect(
+          await server.prisma.batchTaskRun.count({
+            where: { runtimeEnvironmentId: environment.id },
+          })
+        ).toBe(1);
+        expect(
+          await server.prisma.batchTaskRun.findFirst({ where: { friendlyId: batchFriendlyId } })
+        ).toMatchObject({ status: "PENDING", sealed: false });
+        if (route === "task-scoped v1 batch" && credential === "restricted API key") {
+          const inlineResponse = await server.webapp.fetch(path, {
+            method: "POST",
+            headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
+            body: JSON.stringify({
+              items: [{ payload: { message: "inline" }, options: { idempotencyKey } }],
+            }),
+          });
+          expect(inlineResponse.status).toBe(200);
+          expect(await inlineResponse.json()).toMatchObject({
+            batchId: expect.stringMatching(/^batch_/),
+            runs: [run.friendlyId],
+          });
+        }
+      } else {
+        expect(res.status).toBe(route === "v1 batch" || route === "v2 batch" ? 202 : 200);
+        if (route === "trigger") {
+          expect(await res.json()).toMatchObject({ id: run.friendlyId, isCached: true });
+        } else if (route === "v3 items") {
+          expect(await res.json()).toMatchObject({ itemsAccepted: 1, sealed: true });
+        } else if (route === "task-scoped v1 batch") {
+          expect(await res.json()).toMatchObject({
+            batchId: expect.stringMatching(/^batch_/),
+            runs: [run.friendlyId],
+          });
+        } else {
+          expect(await res.json()).toMatchObject({ id: expect.stringMatching(/^batch_/) });
+        }
+      }
+    });
   });
 });
 

apps/webapp/test/streamBatchItemsAuthorization.test.ts

@@ -60,6 +60,46 @@ describe("streaming batch item authorization", () => {
     ).rejects.toThrow();
   });
 
+  it.each([
+    { name: "public JWT", scopes: ["write:batch:batch_123"], allowStoredPayloads: false },
+    {
+      name: "restricted API key",
+      scopes: ["batchTrigger:tasks:task-a"],
+      allowStoredPayloads: false,
+    },
+    { name: "unrestricted API key", scopes: ["admin"], allowStoredPayloads: true },
+  ])(
+    "enforces stored-payload permission for $name on every item",
+    async ({ scopes, allowStoredPayloads }) => {
+      const ability = withActionAliases(buildJwtAbility(scopes));
+      const inline = { task: "task-a", payload: {} };
+      const stored = {
+        task: "task-a",
+        payload: "run_other/payload.json",
+        options: { payloadType: "application/store" },
+      };
+      async function* items() {
+        yield inline;
+        yield stored;
+      }
+
+      const stream = await authorizedBatchItemStream(
+        items(),
+        ability,
+        "batch_123",
+        allowStoredPayloads
+      );
+      const iterator = stream[Symbol.asyncIterator]();
+      await expect(iterator.next()).resolves.toMatchObject({ value: inline, done: false });
+      if (allowStoredPayloads) {
+        await expect(iterator.next()).resolves.toMatchObject({ value: stored, done: false });
+        await expect(iterator.next()).resolves.toMatchObject({ done: true });
+      } else {
+        await expect(iterator.next()).rejects.toThrow();
+      }
+    }
+  );
+
   it("does not delegate batch-wide writes from selected-task credentials", () => {
     const ability = withActionAliases(buildJwtAbility(["batchTrigger:tasks:task-a"]));