diff --git a/tests/web/attachment-staging.test.ts b/tests/web/attachment-staging.test.ts new file mode 100644 index 00000000..2c9feb88 --- /dev/null +++ b/tests/web/attachment-staging.test.ts @@ -0,0 +1,224 @@ +import assert from "node:assert/strict"; +import { + lstat, + mkdtemp, + readdir, + readFile, + rm, + symlink, +} from "node:fs/promises"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import test from "node:test"; +import { + WebAttachmentStagingError, + WebAttachmentStagingStore, + type WebAttachmentBinding, + type WebAttachmentStagingLimits, +} from "../../web/runtime/attachment-staging.ts"; + +const limits: WebAttachmentStagingLimits = { + maxAttachments: 2, + maxAttachmentBytes: 8, + maxTotalBytes: 12, + maxStagedBytes: 16, + maxSettledReceipts: 2, +}; + +const binding: WebAttachmentBinding = { + workspace: process.cwd(), + sessionId: "session-1", + commandId: "command-1", +}; + +async function fixture() { + const parent = await mkdtemp(join(tmpdir(), "openpi-attachment-test-")); + const store = await WebAttachmentStagingStore.create(parent, limits); + return { + parent, + store, + async cleanup() { + await store.dispose(); + await rm(parent, { recursive: true, force: true }); + }, + }; +} + +test("stages with server-owned paths and consumes once for the exact binding", async () => { + const value = await fixture(); + try { + const batch = await value.store.stage(binding, [ + { + name: "../browser-name-is-metadata.txt", + mime: "text/plain", + bytes: Buffer.from("hello"), + }, + ]); + assert.equal(batch.count, 1); + assert.equal(batch.totalBytes, 5); + + const [storeDirectory] = await readdir(value.parent); + const [batchDirectory] = await readdir(join(value.parent, storeDirectory)); + const [payloadName] = await readdir( + join(value.parent, storeDirectory, batchDirectory), + ); + assert.equal(payloadName.includes("browser-name"), false); + assert.equal( + await readFile( + join(value.parent, storeDirectory, batchDirectory, payloadName), + "utf8", + ), + "hello", + ); + + assert.deepEqual( + await value.store.consume(batch.id, { + ...binding, + commandId: "wrong-command", + }), + { status: "stale" }, + ); + const consumed = await value.store.consume(batch.id, binding); + assert.equal(consumed.status, "consumed"); + if (consumed.status === "consumed") { + assert.equal( + consumed.attachments[0]?.name, + "../browser-name-is-metadata.txt", + ); + assert.equal( + Buffer.from(consumed.attachments[0]?.bytes ?? []).toString(), + "hello", + ); + } + assert.deepEqual(await value.store.consume(batch.id, binding), { + status: "settled", + outcome: "consumed", + }); + } finally { + await value.cleanup(); + } +}); + +test("enforces count, per-file, aggregate, and store-wide byte bounds", async () => { + const value = await fixture(); + try { + await assert.rejects( + value.store.stage(binding, []), + (error) => + error instanceof WebAttachmentStagingError && + error.code === "ATTACHMENT_LIMIT", + ); + await assert.rejects( + value.store.stage(binding, [ + { name: "large", mime: "text/plain", bytes: Buffer.alloc(9) }, + ]), + (error) => + error instanceof WebAttachmentStagingError && + error.code === "BYTE_LIMIT", + ); + await assert.rejects( + value.store.stage(binding, [ + { name: "a", mime: "text/plain", bytes: Buffer.alloc(7) }, + { name: "b", mime: "text/plain", bytes: Buffer.alloc(6) }, + ]), + (error) => + error instanceof WebAttachmentStagingError && + error.code === "BYTE_LIMIT", + ); + const first = await value.store.stage(binding, [ + { name: "a", mime: "text/plain", bytes: Buffer.alloc(8) }, + ]); + const second = await value.store.stage( + { ...binding, commandId: "command-2" }, + [{ name: "b", mime: "text/plain", bytes: Buffer.alloc(8) }], + ); + await assert.rejects( + value.store.stage({ ...binding, commandId: "command-3" }, [ + { name: "c", mime: "text/plain", bytes: Buffer.alloc(1) }, + ]), + (error) => + error instanceof WebAttachmentStagingError && + error.code === "STORE_LIMIT", + ); + assert.deepEqual(await value.store.discard(first.id, binding), { + status: "discarded", + }); + assert.equal( + ( + await value.store.stage({ ...binding, commandId: "command-3" }, [ + { name: "c", mime: "text/plain", bytes: Buffer.alloc(1) }, + ]) + ).totalBytes, + 1, + ); + assert.deepEqual( + await value.store.discard(second.id, { + ...binding, + commandId: "command-2", + }), + { status: "discarded" }, + ); + } finally { + await value.cleanup(); + } +}); + +test("fails closed when a staged payload is replaced by a symlink", async (t) => { + if (process.platform === "win32") { + t.skip("symlink creation requires host-specific privileges on Windows"); + return; + } + const value = await fixture(); + try { + const batch = await value.store.stage(binding, [ + { name: "safe.txt", mime: "text/plain", bytes: Buffer.from("safe") }, + ]); + const [storeDirectory] = await readdir(value.parent); + const [batchDirectory] = await readdir(join(value.parent, storeDirectory)); + const batchPath = join(value.parent, storeDirectory, batchDirectory); + const [payloadName] = await readdir(batchPath); + const payloadPath = join(batchPath, payloadName); + await rm(payloadPath); + await symlink("/etc/hosts", payloadPath); + + assert.deepEqual(await value.store.consume(batch.id, binding), { + status: "failed", + error: "staged attachment integrity check failed", + }); + assert.equal( + (await lstat(join(value.parent, storeDirectory))).isDirectory(), + true, + ); + await assert.rejects(lstat(batchPath)); + } finally { + await value.cleanup(); + } +}); + +test("discard and host disposal remove private staged artifacts", async () => { + const value = await fixture(); + const first = await value.store.stage(binding, [ + { name: "a", mime: "text/plain", bytes: Buffer.from("a") }, + ]); + const [storeDirectory] = await readdir(value.parent); + const storePath = join(value.parent, storeDirectory); + assert.deepEqual(await value.store.discard(first.id, binding), { + status: "discarded", + }); + assert.deepEqual(await readdir(storePath), []); + + await value.store.stage({ ...binding, commandId: "command-2" }, [ + { name: "b", mime: "text/plain", bytes: Buffer.from("b") }, + ]); + await value.store.dispose(); + await assert.rejects(lstat(storePath)); + await assert.rejects( + value.store.stage(binding, [ + { name: "c", mime: "text/plain", bytes: Buffer.from("c") }, + ]), + (error) => + error instanceof WebAttachmentStagingError && + error.code === "STORE_CLOSED", + ); + await rm(value.parent, { recursive: true, force: true }); +}); diff --git a/web/runtime/attachment-staging.ts b/web/runtime/attachment-staging.ts new file mode 100644 index 00000000..69501d6b --- /dev/null +++ b/web/runtime/attachment-staging.ts @@ -0,0 +1,396 @@ +import { randomUUID } from "node:crypto"; +import { + chmod, + lstat, + mkdir, + mkdtemp, + open, + readFile, + realpath, + rm, +} from "node:fs/promises"; +import { isAbsolute, join, relative } from "node:path"; + +export interface WebAttachmentStagingLimits { + readonly maxAttachments: number; + readonly maxAttachmentBytes: number; + readonly maxTotalBytes: number; + readonly maxStagedBytes: number; + readonly maxSettledReceipts: number; +} + +export interface WebAttachmentBinding { + readonly workspace: string; + readonly sessionId: string; + readonly commandId: string; +} + +export interface WebAttachmentPayload { + /** Display metadata only. It is never used as a filesystem path. */ + readonly name: string; + readonly mime: string; + readonly bytes: Uint8Array; +} + +export interface WebStagedAttachmentBatch { + readonly id: string; + readonly count: number; + readonly totalBytes: number; +} + +export type WebAttachmentConsumeReceipt = + | { + readonly status: "consumed"; + readonly attachments: readonly WebAttachmentPayload[]; + } + | { readonly status: "missing" } + | { readonly status: "stale" } + | { + readonly status: "settled"; + readonly outcome: "consumed" | "discarded" | "failed"; + } + | { readonly status: "failed"; readonly error: string }; + +export type WebAttachmentDiscardReceipt = + | { readonly status: "discarded" } + | { readonly status: "missing" } + | { readonly status: "stale" } + | { + readonly status: "settled"; + readonly outcome: "consumed" | "discarded" | "failed"; + }; + +export type WebAttachmentStagingErrorCode = + | "INVALID_BINDING" + | "INVALID_LIMITS" + | "ATTACHMENT_LIMIT" + | "BYTE_LIMIT" + | "STORE_LIMIT" + | "STORE_CLOSED"; + +export class WebAttachmentStagingError extends Error { + readonly code: WebAttachmentStagingErrorCode; + + constructor(code: WebAttachmentStagingErrorCode, message: string) { + super(message); + this.name = "WebAttachmentStagingError"; + this.code = code; + } +} + +interface StoredAttachment { + readonly path: string; + readonly name: string; + readonly mime: string; + readonly size: number; +} + +interface StagingRecord { + readonly id: string; + readonly binding: WebAttachmentBinding; + readonly directory: string; + readonly attachments: readonly StoredAttachment[]; + readonly totalBytes: number; + status: "staged" | "consumed" | "discarded" | "failed"; + settledAt?: number; +} + +function assertPositiveInteger(value: number, name: string) { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new WebAttachmentStagingError( + "INVALID_LIMITS", + `${name} must be a positive integer`, + ); + } +} + +function validateLimits(limits: WebAttachmentStagingLimits) { + assertPositiveInteger(limits.maxAttachments, "maxAttachments"); + assertPositiveInteger(limits.maxAttachmentBytes, "maxAttachmentBytes"); + assertPositiveInteger(limits.maxTotalBytes, "maxTotalBytes"); + assertPositiveInteger(limits.maxStagedBytes, "maxStagedBytes"); + assertPositiveInteger(limits.maxSettledReceipts, "maxSettledReceipts"); + if (limits.maxAttachmentBytes > limits.maxTotalBytes) { + throw new WebAttachmentStagingError( + "INVALID_LIMITS", + "maxAttachmentBytes cannot exceed maxTotalBytes", + ); + } + if (limits.maxTotalBytes > limits.maxStagedBytes) { + throw new WebAttachmentStagingError( + "INVALID_LIMITS", + "maxTotalBytes cannot exceed maxStagedBytes", + ); + } +} + +function validateBinding(binding: WebAttachmentBinding) { + if ( + !isAbsolute(binding.workspace) || + binding.sessionId.length === 0 || + binding.sessionId.length > 160 || + binding.commandId.length === 0 || + binding.commandId.length > 160 + ) { + throw new WebAttachmentStagingError( + "INVALID_BINDING", + "attachment binding must identify an absolute workspace and bounded Session and command ids", + ); + } +} + +function sameBinding( + expected: WebAttachmentBinding, + actual: WebAttachmentBinding, +) { + return ( + expected.workspace === actual.workspace && + expected.sessionId === actual.sessionId && + expected.commandId === actual.commandId + ); +} + +function isContained(parent: string, child: string) { + const path = relative(parent, child); + return path.length > 0 && path !== ".." && !path.startsWith(`..${process.platform === "win32" ? "\\" : "/"}`) && !isAbsolute(path); +} + +function settledReceipt(record: StagingRecord) { + return { + status: "settled" as const, + outcome: record.status as "consumed" | "discarded" | "failed", + }; +} + +export class WebAttachmentStagingStore { + private readonly records = new Map(); + private readonly directory: string; + private readonly limits: WebAttachmentStagingLimits; + private mutation = Promise.resolve(); + private stagedBytes = 0; + private closing = false; + private disposePromise?: Promise; + + private constructor(directory: string, limits: WebAttachmentStagingLimits) { + this.directory = directory; + this.limits = limits; + } + + static async create( + parentDirectory: string, + limits: WebAttachmentStagingLimits, + ) { + validateLimits(limits); + await mkdir(parentDirectory, { recursive: true, mode: 0o700 }); + const parent = await realpath(parentDirectory); + const directory = await mkdtemp(join(parent, ".openpi-web-attachments-")); + await chmod(directory, 0o700); + return new WebAttachmentStagingStore(directory, limits); + } + + stage(binding: WebAttachmentBinding, payloads: readonly WebAttachmentPayload[]) { + if (this.closing) { + return Promise.reject( + new WebAttachmentStagingError("STORE_CLOSED", "attachment store is closed"), + ); + } + return this.exclusive(async () => { + this.assertOpen(); + validateBinding(binding); + const totalBytes = this.validatePayloads(payloads); + if (this.stagedBytes + totalBytes > this.limits.maxStagedBytes) { + throw new WebAttachmentStagingError( + "STORE_LIMIT", + "attachment staging store byte limit exceeded", + ); + } + + const id = randomUUID(); + const directory = join(this.directory, id); + const attachments: StoredAttachment[] = []; + await mkdir(directory, { mode: 0o700 }); + try { + for (const [index, payload] of payloads.entries()) { + const path = join(directory, `${index}-${randomUUID()}.payload`); + const handle = await open(path, "wx", 0o600); + try { + await handle.writeFile(payload.bytes); + } finally { + await handle.close(); + } + attachments.push({ + path, + name: payload.name, + mime: payload.mime, + size: payload.bytes.byteLength, + }); + } + } catch (error) { + await rm(directory, { recursive: true, force: true }); + throw error; + } + + this.records.set(id, { + id, + binding: { ...binding }, + directory, + attachments, + totalBytes, + status: "staged", + }); + this.stagedBytes += totalBytes; + return { id, count: attachments.length, totalBytes }; + }); + } + + consume(id: string, binding: WebAttachmentBinding) { + return this.exclusive(async (): Promise => { + this.assertOpen(); + validateBinding(binding); + const record = this.records.get(id); + if (!record) return { status: "missing" }; + if (!sameBinding(record.binding, binding)) return { status: "stale" }; + if (record.status !== "staged") return settledReceipt(record); + + record.status = "consumed"; + record.settledAt = Date.now(); + try { + const directory = await realpath(record.directory); + const directoryStat = await lstat(directory); + if (!directoryStat.isDirectory() || directoryStat.isSymbolicLink()) { + throw new Error("invalid staging directory"); + } + const attachments: WebAttachmentPayload[] = []; + for (const attachment of record.attachments) { + const fileStat = await lstat(attachment.path); + const canonicalPath = await realpath(attachment.path); + if ( + !fileStat.isFile() || + fileStat.isSymbolicLink() || + !isContained(directory, canonicalPath) || + fileStat.size !== attachment.size + ) { + throw new Error("invalid staged attachment"); + } + const bytes = await readFile(canonicalPath); + if (bytes.byteLength !== attachment.size) { + throw new Error("staged attachment changed while reading"); + } + attachments.push({ + name: attachment.name, + mime: attachment.mime, + bytes, + }); + } + await this.cleanupRecord(record); + this.trimSettledReceipts(); + return { status: "consumed", attachments }; + } catch { + record.status = "failed"; + try { + await this.cleanupRecord(record); + } catch { + // dispose() retries removal of the store-owned root. + } + this.trimSettledReceipts(); + return { + status: "failed", + error: "staged attachment integrity check failed", + }; + } + }); + } + + discard(id: string, binding: WebAttachmentBinding) { + return this.exclusive(async (): Promise => { + this.assertOpen(); + validateBinding(binding); + const record = this.records.get(id); + if (!record) return { status: "missing" }; + if (!sameBinding(record.binding, binding)) return { status: "stale" }; + if (record.status !== "staged") return settledReceipt(record); + record.status = "discarded"; + record.settledAt = Date.now(); + await this.cleanupRecord(record); + this.trimSettledReceipts(); + return { status: "discarded" }; + }); + } + + dispose() { + if (this.disposePromise) return this.disposePromise; + this.closing = true; + this.disposePromise = this.exclusive(async () => { + await rm(this.directory, { recursive: true, force: true }); + this.records.clear(); + this.stagedBytes = 0; + }); + return this.disposePromise; + } + + private validatePayloads(payloads: readonly WebAttachmentPayload[]) { + if (payloads.length === 0 || payloads.length > this.limits.maxAttachments) { + throw new WebAttachmentStagingError( + "ATTACHMENT_LIMIT", + `attachment count must be between 1 and ${this.limits.maxAttachments}`, + ); + } + let totalBytes = 0; + for (const payload of payloads) { + if (payload.bytes.byteLength > this.limits.maxAttachmentBytes) { + throw new WebAttachmentStagingError( + "BYTE_LIMIT", + "attachment exceeds the staging per-file byte limit", + ); + } + totalBytes += payload.bytes.byteLength; + if (totalBytes > this.limits.maxTotalBytes) { + throw new WebAttachmentStagingError( + "BYTE_LIMIT", + "attachments exceed the staging aggregate byte limit", + ); + } + } + return totalBytes; + } + + private async cleanupRecord(record: StagingRecord) { + await rm(record.directory, { recursive: true, force: true }); + this.stagedBytes = Math.max(0, this.stagedBytes - record.totalBytes); + } + + private trimSettledReceipts() { + const settled = [...this.records.values()] + .filter((record) => record.status !== "staged") + .sort((left, right) => (left.settledAt ?? 0) - (right.settledAt ?? 0)); + const removeCount = settled.length - this.limits.maxSettledReceipts; + for (const record of settled.slice(0, Math.max(0, removeCount))) { + this.records.delete(record.id); + } + } + + private assertOpen() { + if (this.closing) { + throw new WebAttachmentStagingError( + "STORE_CLOSED", + "attachment store is closed", + ); + } + } + + private exclusive(operation: () => Promise) { + const previous = this.mutation; + let release: () => void = () => undefined; + this.mutation = new Promise((resolve) => { + release = resolve; + }); + return (async () => { + await previous; + try { + return await operation(); + } finally { + release(); + } + })(); + } +}