fix(webapp): reject cross-environment waitpoint registration

Waitpoint registration now verifies that the target run belongs to the
authenticated environment.

Mono-RevId: 0998236ba421350cd59c5c90fa09e918df0074cb

Author: carderne

.server-changes/restrict-waitpoint-run-scope.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Reject waitpoint registrations that target a run outside the authenticated environment

apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts

@@ -2,11 +2,12 @@ import { json } from "@remix-run/server-runtime";
 import { type WaitForWaitpointTokenResponseBody } from "@trigger.dev/core/v3";
 import { RunId, WaitpointId } from "@trigger.dev/core/v3/isomorphic";
 import { z } from "zod";
-import type { PrismaReplicaClient } from "~/db.server";
+import { prisma, type PrismaReplicaClient } from "~/db.server";
 import { resolveWaitpointThroughReadThrough } from "~/runEngine/concerns/resolveWaitpointThroughReadThrough.server";
 import { logger } from "~/services/logger.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { engine } from "~/v3/runEngine.server";
+import { runStore } from "~/v3/runStore.server";
 
 const { action } = createActionApiRoute(
   {
@@ -23,7 +24,19 @@ const { action } = createActionApiRoute(
     const runId = RunId.toId(params.runFriendlyId);
 
     try {
-      //check permissions
+      const run = await runStore.findRun(
+        {
+          id: runId,
+          runtimeEnvironmentId: authentication.environment.id,
+        },
+        { select: { id: true } },
+        prisma
+      );
+
+      if (!run) {
+        throw new Response("You don't have permissions for this run", { status: 401 });
+      }
+
       const waitpoint = await resolveWaitpointThroughReadThrough({
         waitpointId,
         environmentId: authentication.environment.id,
@@ -49,7 +62,7 @@ const { action } = createActionApiRoute(
       }
 
       const _result = await engine.blockRunWithWaitpoint({
-        runId,
+        runId: run.id,
         waitpoints: [waitpointId],
         projectId: authentication.environment.project.id,
         organizationId: authentication.environment.organization.id,

apps/webapp/test/waitpointTokenWaitRunAuthorization.test.ts

@@ -0,0 +1,221 @@
+import { PostgresRunStore } from "@internal/run-store";
+import { postgresTest } from "@internal/testcontainers";
+import type { PrismaClient } from "@trigger.dev/database";
+import { RunId, WaitpointId } from "@trigger.dev/core/v3/isomorphic";
+import { beforeEach, describe, expect, vi } from "vitest";
+import { seedTestEnvironment } from "./helpers/seedTestEnvironment";
+
+vi.setConfig({ testTimeout: 60_000, hookTimeout: 60_000 });
+
+const H = vi.hoisted(() => ({
+  handlers: [] as Array<{ config: any; handler: any }>,
+  store: undefined as any,
+  prisma: undefined as any,
+  engineCalls: [] as Array<Record<string, unknown>>,
+}));
+
+vi.mock("~/db.server", () => {
+  const prisma = new Proxy(
+    {},
+    {
+      get(_target, property) {
+        const value = H.prisma[property];
+        return typeof value === "object" && value !== null
+          ? new Proxy(value, { get: (_delegate, method) => H.prisma[property][method] })
+          : value;
+      },
+    }
+  );
+
+  return {
+    prisma,
+    $replica: prisma,
+    runOpsLegacyReplica: prisma,
+    runOpsNewPrisma: prisma,
+    runOpsNewReplica: prisma,
+    runOpsSplitReadEnabled: false,
+  };
+});
+
+vi.mock("~/env.server", () => ({
+  env: {
+    RUN_OPS_SPLIT_ENABLED: false,
+    RUN_OPS_SHARDS: [],
+  },
+}));
+
+vi.mock("~/v3/runOpsMigration/shardHandles.server", () => ({
+  runOpsShardReplicas: new Map(),
+  runOpsShardWriters: new Map(),
+}));
+
+vi.mock("~/v3/runStore.server", () => ({
+  runStore: new Proxy(
+    {},
+    {
+      get(_target, property) {
+        const value = H.store[property];
+        return typeof value === "function" ? value.bind(H.store) : value;
+      },
+    }
+  ),
+}));
+
+vi.mock("~/services/routeBuilders/apiBuilder.server", () => ({
+  createActionApiRoute: (config: any, handler: any) => {
+    H.handlers.push({ config, handler });
+    return { action: vi.fn() };
+  },
+}));
+
+vi.mock("~/v3/runEngine.server", () => ({
+  engine: {
+    blockRunWithWaitpoint: vi.fn(async (args: Record<string, unknown>) => {
+      H.engineCalls.push(args);
+      return {};
+    }),
+  },
+}));
+
+vi.mock("~/services/logger.server", () => ({
+  logger: { error: vi.fn() },
+}));
+
+async function routeHandler() {
+  await import("~/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait");
+  const entry = H.handlers.find(
+    ({ config }) =>
+      config?.params?.shape?.runFriendlyId && config?.params?.shape?.waitpointFriendlyId
+  );
+  if (!entry) throw new Error("waitpoint wait route handler was not captured");
+  return entry.handler as (args: any) => Promise<Response>;
+}
+
+async function seedRun(
+  prisma: PrismaClient,
+  environment: { id: string; projectId: string; organizationId: string }
+) {
+  const run = RunId.generate();
+  await prisma.taskRun.create({
+    data: {
+      id: run.id,
+      friendlyId: run.friendlyId,
+      engine: "V2",
+      taskIdentifier: "authorization-test",
+      payload: "{}",
+      traceId: `trace_${run.id}`,
+      spanId: `span_${run.id}`,
+      queue: "task/authorization-test",
+      projectId: environment.projectId,
+      runtimeEnvironmentId: environment.id,
+      organizationId: environment.organizationId,
+      environmentType: "DEVELOPMENT",
+      runTags: [],
+    },
+  });
+  return run;
+}
+
+async function seedWaitpoint(prisma: PrismaClient, environmentId: string, projectId: string) {
+  const waitpoint = WaitpointId.generate();
+  await prisma.waitpoint.create({
+    data: {
+      id: waitpoint.id,
+      friendlyId: waitpoint.friendlyId,
+      type: "MANUAL",
+      status: "PENDING",
+      idempotencyKey: waitpoint.id,
+      userProvidedIdempotencyKey: false,
+      environmentId,
+      projectId,
+    },
+  });
+  return waitpoint;
+}
+
+beforeEach(() => {
+  H.engineCalls = [];
+});
+
+describe("waitpoint-token wait run authorization", () => {
+  postgresTest("rejects a run from another runtime environment", async ({ prisma }) => {
+    const seed = await seedTestEnvironment(prisma);
+    const foreignEnvironment = await prisma.runtimeEnvironment.create({
+      data: {
+        slug: "staging",
+        type: "DEVELOPMENT",
+        apiKey: `tr_dev_foreign_${seed.environment.id}`,
+        pkApiKey: `pk_dev_foreign_${seed.environment.id}`,
+        shortcode: `foreign_${seed.environment.id}`,
+        projectId: seed.project.id,
+        organizationId: seed.organization.id,
+      },
+    });
+    const run = await seedRun(prisma, foreignEnvironment);
+    const waitpoint = await seedWaitpoint(prisma, seed.environment.id, seed.project.id);
+
+    H.prisma = prisma;
+    H.store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+    const handler = await routeHandler();
+
+    const response = await handler({
+      authentication: {
+        environment: {
+          ...seed.environment,
+          project: seed.project,
+          organization: seed.organization,
+        },
+      },
+      params: {
+        runFriendlyId: run.friendlyId,
+        waitpointFriendlyId: waitpoint.friendlyId,
+      },
+    }).then(
+      (result) => result,
+      (error) => error
+    );
+
+    expect(response).toBeInstanceOf(Response);
+    expect(response.status).toBe(401);
+    expect(await response.text()).toBe("You don't have permissions for this run");
+    expect(H.engineCalls).toHaveLength(0);
+  });
+
+  postgresTest(
+    "attaches a waitpoint to a run in the authenticated environment",
+    async ({ prisma }) => {
+      const seed = await seedTestEnvironment(prisma);
+      const run = await seedRun(prisma, seed.environment);
+      const waitpoint = await seedWaitpoint(prisma, seed.environment.id, seed.project.id);
+
+      H.prisma = prisma;
+      H.store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+      const handler = await routeHandler();
+
+      const response = await handler({
+        authentication: {
+          environment: {
+            ...seed.environment,
+            project: seed.project,
+            organization: seed.organization,
+          },
+        },
+        params: {
+          runFriendlyId: run.friendlyId,
+          waitpointFriendlyId: waitpoint.friendlyId,
+        },
+      });
+
+      expect(response.status).toBe(200);
+      expect(await response.json()).toEqual({ success: true });
+      expect(H.engineCalls).toEqual([
+        {
+          runId: run.id,
+          waitpoints: [waitpoint.id],
+          projectId: seed.project.id,
+          organizationId: seed.organization.id,
+        },
+      ]);
+    }
+  );
+});