diff --git a/packages/trpc/src/router/attachment/attachment.ts b/packages/trpc/src/router/attachment/attachment.ts index a0b96c3162a..47fadc75c44 100644 --- a/packages/trpc/src/router/attachment/attachment.ts +++ b/packages/trpc/src/router/attachment/attachment.ts @@ -1,25 +1,17 @@ import { db } from "@superset/db/client"; -import { attachments, files } from "@superset/db/schema"; +import { files, type SelectFile } from "@superset/db/schema"; import { fileOriginalKey } from "@superset/shared/usercontent"; import type { TRPCRouterRecord } from "@trpc/server"; -import { and, eq, inArray } from "drizzle-orm"; import { headObject, presignedGetUrl, presignedPutUrl } from "../../lib/r2"; +import type { SandboxCaller } from "../../lib/sandbox/api-credential"; import { protectedProcedure, userError } from "../../trpc"; import { requireActiveOrgMembership } from "../utils/active-org"; +import { resolvableFiles, TRANSFER_ONLY } from "./resolvable"; import { createAttachmentUploadSchema, resolveAttachmentsSchema, } from "./schema"; -/** - * What marks a `files` row as one of these. Every other kind of file is - * identified by its content — page assets deduplicate on the digest — and no - * real file hashes to the empty string, so this both keeps an attachment out - * of any reuse lookup and keeps `resolve` off files it has no business - * handing a URL for. - */ -const TRANSFER_ONLY = ""; - /** * Composer attachments: bytes on their way from a client to the machine an * agent is running on. @@ -29,11 +21,10 @@ const TRANSFER_ONLY = ""; * bodies in a Worker isolate and is refused outright by Cloudflare's edge * above 100 MB. What travels over the relay is a list of file ids. * - * Rows here stay `pending` for their whole life, which is the point: the - * object is a transfer buffer, dead as soon as the host has written it into - * the worktree, and `pending` is exactly what the file sweep reclaims a day - * later. Nothing references these bytes afterwards, so there is no `ready` - * state to move to and no parent to attach them to. + * A row stays `pending` while it is only a transfer buffer, which the file + * sweep reclaims a day later. A file sent to a cloud workspace is kept + * instead: `anchorAttachments` marks it `ready` on that workspace, and it is + * reaped with the workspace's box. */ export const attachmentRouter = { /** @@ -89,14 +80,8 @@ export const attachmentRouter = { * Presigned GETs for a host about to materialize these into a worktree, * with the metadata it needs to name the files. * - * Scoped to the caller's own organization and to attachments, which is - * what keeps a host to the bytes it asked for: `files.id` is the only - * thing the caller presents, and knowing one is not authority to read it. - * Without the second scope this would sign a URL for any file in the org, - * a page's private assets included. - * - * A box calling for itself gets only the files attached to its own - * workspace: access to the box is access to its files, nothing wider. + * Scoped by `resolvableFiles`: `files.id` is the only thing the caller + * presents, and knowing one is not authority to read it. * * A mutation despite reading nothing but rows: the URLs it mints are * short-lived credentials, and a query's response is cacheable. @@ -104,33 +89,7 @@ export const attachmentRouter = { resolve: protectedProcedure .input(resolveAttachmentsSchema) .mutation(async ({ ctx, input }) => { - const organizationId = await requireActiveOrgMembership(ctx); - - const rows = await db - .select() - .from(files) - .where( - and( - inArray(files.id, input.fileIds), - eq(files.organizationId, organizationId), - eq(files.sha256, TRANSFER_ONLY), - ctx.sandboxCaller - ? inArray( - files.id, - db - .select({ id: attachments.fileId }) - .from(attachments) - .where( - and( - eq(attachments.parentKind, "cloud_workspace"), - eq(attachments.parentId, ctx.sandboxCaller.workspaceId), - ), - ), - ) - : undefined, - ), - ); - const byId = new Map(rows.map((row) => [row.id, row])); + const byId = await loadResolvable(ctx, input.fileIds); // In the order asked for: the caller names attachments positionally // and a reordered answer would rename them. @@ -144,29 +103,69 @@ export const attachmentRouter = { i18nKey: "serverError.attachment.notFound", }); } - - // A presigned PUT lands without the API seeing it, so this is - // where an upload is held to what it declared. A short read - // that never finished is the case worth catching — it would - // otherwise reach the agent as a truncated file. - const key = fileOriginalKey(file.id); - const head = await headObject(key); - if (!head || head.sizeBytes !== file.sizeBytes) { + const resolved = await signUploaded(file); + if (!resolved) { throw userError({ code: "BAD_REQUEST", message: "Attachment was not uploaded — send the bytes first", i18nKey: "serverError.attachment.notUploaded", }); } + return resolved; + }), + ); + }), - return { - fileId: file.id, - name: file.name, - contentType: file.contentType, - sizeBytes: file.sizeBytes, - url: await presignedGetUrl(key), - }; + /** + * `resolve` for a reader that shows what it can: a transcript with one + * file gone still renders the rest. Every id gets an answer, in the order + * asked for, and one the caller may not see is reported the same as one + * that does not exist. + */ + resolveEach: protectedProcedure + .input(resolveAttachmentsSchema) + .mutation(async ({ ctx, input }) => { + const byId = await loadResolvable(ctx, input.fileIds); + return await Promise.all( + input.fileIds.map(async (fileId) => { + const file = byId.get(fileId); + const resolved = file ? await signUploaded(file) : null; + return resolved ?? { fileId, missing: true as const }; }), ); }), } satisfies TRPCRouterRecord; + +async function loadResolvable( + ctx: Parameters[0] & { + sandboxCaller: SandboxCaller | null; + }, + fileIds: string[], +) { + const rows = await resolvableFiles({ + organizationId: await requireActiveOrgMembership(ctx), + userId: ctx.session.user.id, + boxWorkspaceId: ctx.sandboxCaller?.workspaceId ?? null, + fileIds, + }); + return new Map(rows.map((row) => [row.id, row])); +} + +/** + * A presigned PUT lands without the API seeing it, so this is where an + * upload is held to what it declared. A short read that never finished is + * the case worth catching — it would otherwise reach the agent as a + * truncated file. Null when the bytes are not there. + */ +async function signUploaded(file: SelectFile) { + const key = fileOriginalKey(file.id); + const head = await headObject(key); + if (!head || head.sizeBytes !== file.sizeBytes) return null; + return { + fileId: file.id, + name: file.name, + contentType: file.contentType, + sizeBytes: file.sizeBytes, + url: await presignedGetUrl(key), + }; +} diff --git a/packages/trpc/src/router/attachment/resolvable.integration.ts b/packages/trpc/src/router/attachment/resolvable.integration.ts new file mode 100644 index 00000000000..d1dd4af50d0 --- /dev/null +++ b/packages/trpc/src/router/attachment/resolvable.integration.ts @@ -0,0 +1,148 @@ +import { afterAll, beforeAll, describe, expect, test } from "bun:test"; +import { db, dbWs } from "@superset/db/client"; +import { + attachments, + cloudWorkspaces, + environments, + files, + organizations, + users, +} from "@superset/db/schema"; +import { eq, inArray } from "drizzle-orm"; +import { resolvableFiles, TRANSFER_ONLY } from "./resolvable"; + +const ORG = crypto.randomUUID(); +const OWNER = crypto.randomUUID(); +const TEAMMATE = crypto.randomUUID(); +const now = Date.now(); + +const insertFile = async ( + status: "pending" | "ready", + sha256 = TRANSFER_ONLY, +) => { + const id = crypto.randomUUID(); + await db.insert(files).values({ + id, + organizationId: ORG, + name: "screenshot.png", + contentType: "image/png", + sizeBytes: 4, + sha256, + status, + createdByUserId: OWNER, + }); + return id; +}; + +const insertWorkspace = async (visibility: "org" | "just_me") => { + const [environment] = await db + .select({ id: environments.id }) + .from(environments) + .where(eq(environments.organizationId, ORG)); + const id = crypto.randomUUID(); + await db.insert(cloudWorkspaces).values({ + id, + organizationId: ORG, + name: "box", + branch: "main", + baseBranch: "main", + providerSandboxId: `box-${id}`, + environmentId: environment?.id as string, + visibility, + createdByUserId: OWNER, + }); + return id; +}; + +const anchor = (fileId: string, workspaceId: string) => + db.insert(attachments).values({ + fileId, + parentKind: "cloud_workspace", + parentId: workspaceId, + }); + +const resolvable = async ( + userId: string, + fileIds: string[], + boxWorkspaceId: string | null = null, +) => + ( + await resolvableFiles({ + organizationId: ORG, + userId, + boxWorkspaceId, + fileIds, + }) + ) + .map((row) => row.id) + .sort(); + +beforeAll(async () => { + await db + .insert(organizations) + .values({ id: ORG, name: "Resolve Org", slug: `resolve-org-${now}` }); + await db.insert(users).values([ + { id: OWNER, name: "Owner", email: `owner-${now}@example.com` }, + { id: TEAMMATE, name: "Teammate", email: `teammate-${now}@example.com` }, + ]); + await db.insert(environments).values({ + organizationId: ORG, + name: "default", + sourceKind: "image", + sourceRef: "image", + bundleSha: "sha", + }); +}); + +afterAll(async () => { + await db.delete(organizations).where(eq(organizations.id, ORG)); + await db.delete(users).where(inArray(users.id, [OWNER, TEAMMATE])); + await dbWs.$client.end?.().catch(() => {}); +}); + +describe("resolvableFiles", () => { + test("an upload in flight resolves for anyone in the org", async () => { + const pending = await insertFile("pending"); + expect(await resolvable(TEAMMATE, [pending])).toEqual([pending]); + }); + + test("a kept file resolves only through a workspace the caller can see", async () => { + const shared = await insertWorkspace("org"); + const hidden = await insertWorkspace("just_me"); + const inShared = await insertFile("ready"); + const inHidden = await insertFile("ready"); + await anchor(inShared, shared); + await anchor(inHidden, hidden); + + expect(await resolvable(OWNER, [inShared, inHidden])).toEqual( + [inShared, inHidden].sort(), + ); + expect(await resolvable(TEAMMATE, [inShared, inHidden])).toEqual([ + inShared, + ]); + }); + + test("a kept file on no cloud workspace does not resolve", async () => { + const loose = await insertFile("ready"); + expect(await resolvable(OWNER, [loose])).toEqual([]); + }); + + test("a box resolves only its own files, never one in flight", async () => { + const box = await insertWorkspace("org"); + const other = await insertWorkspace("org"); + const own = await insertFile("ready"); + const theirs = await insertFile("ready"); + const pending = await insertFile("pending"); + await anchor(own, box); + await anchor(theirs, other); + + expect(await resolvable(OWNER, [own, theirs, pending], box)).toEqual([own]); + }); + + test("a page asset never resolves", async () => { + const box = await insertWorkspace("org"); + const asset = await insertFile("ready", "ab".repeat(32)); + await anchor(asset, box); + expect(await resolvable(OWNER, [asset])).toEqual([]); + }); +}); diff --git a/packages/trpc/src/router/attachment/resolvable.ts b/packages/trpc/src/router/attachment/resolvable.ts new file mode 100644 index 00000000000..5e229aa2031 --- /dev/null +++ b/packages/trpc/src/router/attachment/resolvable.ts @@ -0,0 +1,69 @@ +import { db } from "@superset/db/client"; +import { attachments, cloudWorkspaces, files } from "@superset/db/schema"; +import { and, eq, inArray, or, type SQL } from "drizzle-orm"; +import { visibleTo } from "../cloud-workspace/access"; + +/** + * What marks a `files` row as a composer attachment. Every other kind of file is + * identified by its content — page assets deduplicate on the digest — and no + * real file hashes to the empty string, so this both keeps an attachment out + * of any reuse lookup and keeps `resolve` off files it has no business + * handing a URL for. + */ +export const TRANSFER_ONLY = ""; + +/** + * The attachments among `fileIds` a caller may read. + * + * A box gets only the files anchored to its own workspace. Anyone else gets + * an upload still in flight, and a kept file only through a cloud workspace + * they can see. In flight stays org-wide because a host resolves with its + * owner's credential, and the sender may be a teammate who shares that host. + */ +export function resolvableFiles({ + organizationId, + userId, + boxWorkspaceId, + fileIds, +}: { + organizationId: string; + userId: string; + boxWorkspaceId: string | null; + fileIds: string[]; +}) { + return db + .select() + .from(files) + .where( + and( + inArray(files.id, fileIds), + eq(files.organizationId, organizationId), + eq(files.sha256, TRANSFER_ONLY), + boxWorkspaceId + ? inArray( + files.id, + anchoredTo(eq(cloudWorkspaces.id, boxWorkspaceId)), + ) + : or( + eq(files.status, "pending"), + inArray( + files.id, + anchoredTo( + and( + eq(cloudWorkspaces.organizationId, organizationId), + visibleTo(userId), + ), + ), + ), + ), + ), + ); +} + +function anchoredTo(workspaces: SQL | undefined) { + return db + .select({ id: attachments.fileId }) + .from(attachments) + .innerJoin(cloudWorkspaces, eq(cloudWorkspaces.id, attachments.parentId)) + .where(and(eq(attachments.parentKind, "cloud_workspace"), workspaces)); +} diff --git a/packages/trpc/src/router/cloud-workspace/cloud-workspace.ts b/packages/trpc/src/router/cloud-workspace/cloud-workspace.ts index 0d5f2c7942d..3ce4746bd74 100644 --- a/packages/trpc/src/router/cloud-workspace/cloud-workspace.ts +++ b/packages/trpc/src/router/cloud-workspace/cloud-workspace.ts @@ -46,7 +46,7 @@ import { import { recordCloudWorkspaceActivity } from "./activity"; import { queueSettleWorkspace, settleWorkspaceNow } from "./meter"; import { nextSandboxNameFor } from "./provision"; -import { queueReap } from "./reap"; +import { queueReap, reapWorkspaceFiles } from "./reap"; import { cloudWorkspaceRecordRouter } from "./record"; import { cloudWorkspaceSharingRouter } from "./sharing"; import { queueProvision, startCloudWorkspace } from "./start"; @@ -623,6 +623,7 @@ export const cloudWorkspaceRouter = { ), ); await deleteSandbox(row.providerSandboxId); + await reapWorkspaceFiles(row.id); }); } await recordCloudWorkspaceActivity( diff --git a/packages/trpc/src/router/cloud-workspace/reap.test.ts b/packages/trpc/src/router/cloud-workspace/reap.test.ts index 4247bd6c04f..9a54f220332 100644 --- a/packages/trpc/src/router/cloud-workspace/reap.test.ts +++ b/packages/trpc/src/router/cloud-workspace/reap.test.ts @@ -1,11 +1,13 @@ import { beforeEach, describe, expect, test } from "bun:test"; import { db } from "@superset/db/client"; import { stub } from "../../../test/stub"; +import * as files from "../../lib/files"; import * as sandbox from "../../lib/sandbox"; import * as jobs from "./jobs"; import * as meter from "./meter"; type Row = { + id: string; status: string; deletedAt: Date | null; provider: string; @@ -19,6 +21,7 @@ let queued: Array<{ path: string; body: unknown; delaySeconds?: number }> = []; let failingStage: Stage | null = null; let meterCalls: string[] = []; let meterFailure: Error | null = null; +let reapedFiles: string[] = []; stub(db.query.cloudWorkspaces, { findFirst: () => Promise.resolve(row) }); stub(sandbox, { @@ -49,6 +52,15 @@ stub(jobs, { }, }); +stub(files, { + detachAll: ({ parentIds }: { parentIds: string[] }) => + Promise.resolve(parentIds.map((id) => `file-of-${id}`)), + reapOrphanFiles: (fileIds: string[]) => { + reapedFiles.push(...fileIds); + return Promise.resolve(); + }, +}); + stub(meter, { settleWorkspaceNow: () => { meterCalls.push("settle"); @@ -112,6 +124,7 @@ describe("queueReap", () => { describe("reapArchivedCloudWorkspace", () => { beforeEach(() => { row = { + id: input.cloudWorkspaceId, status: "deleted", deletedAt: archivedAt, provider: "vercel", @@ -121,6 +134,7 @@ describe("reapArchivedCloudWorkspace", () => { queued = []; meterCalls = []; meterFailure = null; + reapedFiles = []; }); test("the stop stage stops the box and queues its final settle", async () => { @@ -130,6 +144,7 @@ describe("reapArchivedCloudWorkspace", () => { expect(sandboxCalls).toEqual(["stop:ws-box"]); expect(meterCalls).toEqual(["queue"]); expect(queued).toEqual([]); + expect(reapedFiles).toEqual([]); }); test("the delete stage settles the box's usage before deleting it", async () => { @@ -140,6 +155,11 @@ describe("reapArchivedCloudWorkspace", () => { expect(sandboxCalls).toEqual(["delete:ws-box"]); }); + test("the delete stage removes the files sent to the workspace", async () => { + await reapArchivedCloudWorkspace({ ...input, stage: "delete" }); + expect(reapedFiles).toEqual([`file-of-${input.cloudWorkspaceId}`]); + }); + test("a delete still deletes when the final settle fails: at most one lock window goes unbilled", async () => { meterFailure = new Error("autumn down"); expect( @@ -155,6 +175,7 @@ describe("reapArchivedCloudWorkspace", () => { "skipped", ); expect(sandboxCalls).toEqual([]); + expect(reapedFiles).toEqual([]); }); test(`the ${stage} stage leaves a later archive's box alone`, async () => { @@ -175,6 +196,7 @@ describe("reapArchivedCloudWorkspace", () => { await reapArchivedCloudWorkspace({ ...input, stage: "delete" }), ).toBe("skipped"); row = { + id: input.cloudWorkspaceId, status: "deleted", deletedAt: archivedAt, provider: "e2b", diff --git a/packages/trpc/src/router/cloud-workspace/reap.ts b/packages/trpc/src/router/cloud-workspace/reap.ts index 5f298a645f2..864d1d5d120 100644 --- a/packages/trpc/src/router/cloud-workspace/reap.ts +++ b/packages/trpc/src/router/cloud-workspace/reap.ts @@ -1,6 +1,7 @@ import { db } from "@superset/db/client"; import { cloudWorkspaces } from "@superset/db/schema"; import { eq } from "drizzle-orm"; +import { detachAll, reapOrphanFiles } from "../../lib/files"; import { deleteSandbox, stopSandbox } from "../../lib/sandbox"; import { publishCloudWorkspaceJob } from "./jobs"; import { queueSettleWorkspace, settleWorkspaceNow } from "./meter"; @@ -36,6 +37,7 @@ export async function reapArchivedCloudWorkspace( console.error(`[cloud-workspace] ${row.id} final settle failed`, error), ); await deleteSandbox(row.providerSandboxId); + await reapWorkspaceFiles(row.id); return "reaped"; } await stopSandbox(row.providerSandboxId); @@ -43,6 +45,21 @@ export async function reapArchivedCloudWorkspace( return "stopped"; } +/** Drops the files sent to a workspace whose box is gone for good. */ +export async function reapWorkspaceFiles(cloudWorkspaceId: string) { + await reapOrphanFiles( + await detachAll({ + parentKind: "cloud_workspace", + parentIds: [cloudWorkspaceId], + }), + ).catch((error) => + console.error( + `[cloud-workspace] ${cloudWorkspaceId} file reap failed`, + error, + ), + ); +} + function queueReapStage( input: ReapArchivedCloudWorkspaceInput, delaySeconds: number,