diff --git a/packages/harness/src/attachments/attachment-store.ts b/packages/harness/src/attachments/attachment-store.ts new file mode 100644 index 0000000..483bf67 --- /dev/null +++ b/packages/harness/src/attachments/attachment-store.ts @@ -0,0 +1,19 @@ +import type { ImageAttachment } from "@agent-studio/shared"; +import type { ImageContent } from "@earendil-works/pi-ai"; +import type { Readable } from "node:stream"; +import type { SupportedImageMediaType } from "./image-signatures.js"; + +export type ImageUpload = { + name: string; + mediaType: SupportedImageMediaType; + stream: Readable; +}; + +export interface AttachmentStore { + add(sessionId: string, upload: ImageUpload): Promise; + get(sessionId: string, attachmentId: string): ImageAttachment | undefined; + openNormalized(sessionId: string, attachmentId: string): Promise; + remove(sessionId: string, attachmentId: string): Promise; + removeAll(sessionId: string): Promise; + toPiImages(sessionId: string, attachmentIds: string[]): Promise; +} diff --git a/packages/harness/src/attachments/file-attachment-store.ts b/packages/harness/src/attachments/file-attachment-store.ts new file mode 100644 index 0000000..d1154ab --- /dev/null +++ b/packages/harness/src/attachments/file-attachment-store.ts @@ -0,0 +1,262 @@ +import { randomUUID } from "node:crypto"; +import { createWriteStream } from "node:fs"; +import { chmod, lstat, mkdir, readFile, realpath, rename, rm, stat } from "node:fs/promises"; +import { isAbsolute, join, relative, resolve, sep } from "node:path"; +import { Transform, type TransformCallback } from "node:stream"; +import { pipeline } from "node:stream/promises"; +import { ImageAttachmentSchema, type ImageAttachment } from "@agent-studio/shared"; +import type { ImageContent } from "@earendil-works/pi-ai"; +import sharp from "sharp"; +import type { EventStore } from "../events/event-store.js"; +import { HarnessError } from "../errors/harness-error.js"; +import type { AttachmentStore, ImageUpload } from "./attachment-store.js"; +import { assertImageType, extensionForMediaType, type SupportedImageMediaType } from "./image-signatures.js"; + +type StoredAttachment = { metadata: ImageAttachment; filePath: string }; + +export type FileAttachmentStoreOptions = { + sessionsRoot: string; + eventStore?: EventStore; + maxBytes?: number; + maxWidth?: number; + maxHeight?: number; + maxPixels?: number; + maxAttachmentsPerMessage?: number; +}; + +export class FileAttachmentStore implements AttachmentStore { + readonly #sessionsRoot: string; + readonly #eventStore: EventStore | undefined; + readonly #maxBytes: number; + readonly #maxWidth: number; + readonly #maxHeight: number; + readonly #maxPixels: number; + readonly #maxAttachmentsPerMessage: number; + readonly #bySession = new Map>(); + + public constructor(options: FileAttachmentStoreOptions) { + this.#sessionsRoot = resolve(options.sessionsRoot); + this.#eventStore = options.eventStore; + this.#maxBytes = options.maxBytes ?? 10 * 1024 * 1024; + this.#maxWidth = options.maxWidth ?? 8_192; + this.#maxHeight = options.maxHeight ?? 8_192; + this.#maxPixels = options.maxPixels ?? 25_000_000; + this.#maxAttachmentsPerMessage = options.maxAttachmentsPerMessage ?? 4; + for (const [name, value] of Object.entries({ + maxBytes: this.#maxBytes, + maxWidth: this.#maxWidth, + maxHeight: this.#maxHeight, + maxPixels: this.#maxPixels, + maxAttachmentsPerMessage: this.#maxAttachmentsPerMessage, + })) assertPositiveInteger(value, name); + } + + public async add(sessionId: string, upload: ImageUpload): Promise { + let uploadPath: string | undefined; + let normalizedPath: string | undefined; + try { + const safeName = validateUploadMetadata(sessionId, upload); + const attachmentDirectory = await this.#attachmentDirectory(sessionId); + uploadPath = join(attachmentDirectory, `${randomUUID()}.upload`); + normalizedPath = join(attachmentDirectory, `${randomUUID()}.normalized`); + await pipeline( + upload.stream, + new ByteLimitTransform(this.#maxBytes), + createWriteStream(uploadPath, { flags: "wx", mode: 0o600 }), + ); + const headerAndBody = await readFile(uploadPath); + const mediaType = assertImageType(safeName, upload.mediaType, headerAndBody.subarray(0, 12)); + const sourceMetadata = await sharp(uploadPath, { failOn: "error", limitInputPixels: false }).metadata(); + assertDimensions(sourceMetadata.width, sourceMetadata.height, this.#maxWidth, this.#maxHeight, this.#maxPixels); + await normalizeImage(uploadPath, normalizedPath, mediaType, this.#maxPixels); + const imageMetadata = await sharp(normalizedPath, { failOn: "error", limitInputPixels: this.#maxPixels }).metadata(); + const width = imageMetadata.width; + const height = imageMetadata.height; + assertDimensions(width, height, this.#maxWidth, this.#maxHeight, this.#maxPixels); + const normalizedStat = await stat(normalizedPath); + if (normalizedStat.size > this.#maxBytes) { + throw new HarnessError("IMAGE_TOO_LARGE", "Normalized image exceeds the configured byte limit"); + } + + const attachmentId = `att_${randomUUID()}`; + const finalPath = join(attachmentDirectory, `${attachmentId}${extensionForMediaType(mediaType)}`); + await rename(normalizedPath, finalPath); + normalizedPath = undefined; + const metadata = ImageAttachmentSchema.parse({ + attachmentId, + sessionId, + name: safeName, + mediaType, + byteSize: normalizedStat.size, + width, + height, + status: "ready", + }); + this.#sessionIndex(sessionId).set(attachmentId, { metadata, filePath: finalPath }); + this.#eventStore?.append(sessionId, { + type: "attachment.uploaded", + payload: { + attachmentId, + name: safeName, + mediaType, + byteSize: normalizedStat.size, + width, + height, + }, + }); + return { ...metadata }; + } catch (error) { + const safeName = typeof upload.name === "string" && upload.name.trim() !== "" + ? upload.name.trim().slice(0, 255) + : "unnamed-image"; + this.#eventStore?.append(sessionId, { + type: "attachment.rejected", + payload: { name: safeName, reason: safeErrorCode(error) }, + }); + if (error instanceof HarnessError) throw error; + throw new HarnessError("IMAGE_INVALID", "Image could not be decoded"); + } finally { + if (uploadPath !== undefined) await rm(uploadPath, { force: true }).catch(() => undefined); + if (normalizedPath !== undefined) await rm(normalizedPath, { force: true }).catch(() => undefined); + } + } + + public get(sessionId: string, attachmentId: string): ImageAttachment | undefined { + const stored = this.#bySession.get(sessionId)?.get(attachmentId); + return stored === undefined ? undefined : { ...stored.metadata }; + } + + public async openNormalized(sessionId: string, attachmentId: string): Promise { + const stored = this.#stored(sessionId, attachmentId); + return readFile(stored.filePath); + } + + public async remove(sessionId: string, attachmentId: string): Promise { + const session = this.#bySession.get(sessionId); + const stored = session?.get(attachmentId); + if (stored === undefined) return false; + session!.delete(attachmentId); + await rm(stored.filePath, { force: true }); + return true; + } + + public async removeAll(sessionId: string): Promise { + const session = this.#bySession.get(sessionId); + if (session === undefined) return; + this.#bySession.delete(sessionId); + await Promise.all([...session.values()].map(({ filePath }) => rm(filePath, { force: true }))); + } + + public async toPiImages(sessionId: string, attachmentIds: string[]): Promise { + if (attachmentIds.length > this.#maxAttachmentsPerMessage) { + throw new HarnessError("TOO_MANY_ATTACHMENTS", "Message contains too many attachments"); + } + return Promise.all(attachmentIds.map(async (attachmentId) => { + const stored = this.#stored(sessionId, attachmentId); + const bytes = await readFile(stored.filePath); + return { type: "image", data: bytes.toString("base64"), mimeType: stored.metadata.mediaType }; + })); + } + + async #attachmentDirectory(sessionId: string): Promise { + if (!/^ses_[a-zA-Z0-9_-]+$/.test(sessionId)) { + throw new HarnessError("INVALID_SESSION_ID", "Session ID is invalid"); + } + const canonicalRoot = await realpath(this.#sessionsRoot); + const sessionDirectory = join(canonicalRoot, sessionId); + const sessionMetadata = await lstat(sessionDirectory).catch(() => undefined); + if (sessionMetadata === undefined || !sessionMetadata.isDirectory() || sessionMetadata.isSymbolicLink()) { + throw new HarnessError("SESSION_NOT_FOUND", "Session workspace does not exist"); + } + assertContained(canonicalRoot, await realpath(sessionDirectory)); + const directory = join(sessionDirectory, "attachments"); + await mkdir(directory, { recursive: true, mode: 0o700 }); + const canonicalDirectory = await realpath(directory); + assertContained(sessionDirectory, canonicalDirectory); + return canonicalDirectory; + } + + #stored(sessionId: string, attachmentId: string): StoredAttachment { + const stored = this.#bySession.get(sessionId)?.get(attachmentId); + if (stored === undefined) throw new HarnessError("ATTACHMENT_NOT_FOUND", "Attachment was not found"); + return stored; + } + + #sessionIndex(sessionId: string): Map { + const current = this.#bySession.get(sessionId); + if (current !== undefined) return current; + const created = new Map(); + this.#bySession.set(sessionId, created); + return created; + } +} + +class ByteLimitTransform extends Transform { + #bytes = 0; + + public constructor(private readonly maxBytes: number) { + super(); + } + + public override _transform(chunk: Buffer, _encoding: BufferEncoding, callback: TransformCallback): void { + this.#bytes += chunk.byteLength; + if (this.#bytes > this.maxBytes) { + callback(new HarnessError("IMAGE_TOO_LARGE", "Image exceeds the configured byte limit")); + return; + } + callback(null, chunk); + } +} + +async function normalizeImage( + inputPath: string, + outputPath: string, + mediaType: SupportedImageMediaType, + maxPixels: number, +): Promise { + let pipeline = sharp(inputPath, { failOn: "error", limitInputPixels: maxPixels }).rotate(); + if (mediaType === "image/jpeg") pipeline = pipeline.jpeg({ quality: 90 }); + if (mediaType === "image/png") pipeline = pipeline.png(); + if (mediaType === "image/webp") pipeline = pipeline.webp({ quality: 90 }); + await pipeline.toFile(outputPath); + await chmod(outputPath, 0o600); +} + +function validateUploadMetadata(sessionId: string, upload: ImageUpload): string { + if (!/^ses_[a-zA-Z0-9_-]+$/.test(sessionId)) { + throw new HarnessError("INVALID_SESSION_ID", "Session ID is invalid"); + } + if (typeof upload.name !== "string" || upload.name.trim().length === 0 || upload.name.trim().length > 255) { + throw new HarnessError("INVALID_ATTACHMENT_NAME", "Attachment name is invalid"); + } + return upload.name.trim(); +} + +function safeErrorCode(error: unknown): string { + return error instanceof HarnessError ? error.code : "IMAGE_INVALID"; +} + +function assertContained(root: string, candidate: string): void { + const fromRoot = relative(root, candidate); + if (fromRoot === "" || (!fromRoot.startsWith(`..${sep}`) && fromRoot !== ".." && !isAbsolute(fromRoot))) return; + throw new HarnessError("PATH_OUTSIDE_WORKSPACE", "Attachment path is outside the Session workspace"); +} + +function assertPositiveInteger(value: number, name: string): void { + if (!Number.isSafeInteger(value) || value < 1) throw new RangeError(`${name} must be a positive integer`); +} + +function assertDimensions( + width: number | undefined, + height: number | undefined, + maxWidth: number, + maxHeight: number, + maxPixels: number, +): asserts width is number { + if (width === undefined || height === undefined) { + throw new HarnessError("IMAGE_INVALID", "Image dimensions could not be determined"); + } + if (width > maxWidth || height > maxHeight || width * height > maxPixels) { + throw new HarnessError("IMAGE_DIMENSIONS_EXCEEDED", "Image dimensions exceed the configured limit"); + } +} diff --git a/packages/harness/src/attachments/image-signatures.ts b/packages/harness/src/attachments/image-signatures.ts new file mode 100644 index 0000000..3dc0363 --- /dev/null +++ b/packages/harness/src/attachments/image-signatures.ts @@ -0,0 +1,53 @@ +import { extname } from "node:path"; +import { HarnessError } from "../errors/harness-error.js"; + +export type SupportedImageMediaType = "image/jpeg" | "image/png" | "image/webp"; + +const extensions: Record> = { + "image/jpeg": new Set([".jpg", ".jpeg"]), + "image/png": new Set([".png"]), + "image/webp": new Set([".webp"]), +}; + +export function detectImageMediaType(bytes: Uint8Array): SupportedImageMediaType | undefined { + if ( + bytes.length >= 8 && + bytes[0] === 0x89 && bytes[1] === 0x50 && bytes[2] === 0x4e && bytes[3] === 0x47 && + bytes[4] === 0x0d && bytes[5] === 0x0a && bytes[6] === 0x1a && bytes[7] === 0x0a + ) return "image/png"; + if (bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) { + return "image/jpeg"; + } + if ( + bytes.length >= 12 && + Buffer.from(bytes.subarray(0, 4)).toString("ascii") === "RIFF" && + Buffer.from(bytes.subarray(8, 12)).toString("ascii") === "WEBP" + ) return "image/webp"; + return undefined; +} + +export function assertImageType( + name: string, + declaredType: string, + bytes: Uint8Array, +): SupportedImageMediaType { + if (!(declaredType in extensions)) { + throw new HarnessError("IMAGE_TYPE_NOT_ALLOWED", "Only JPEG, PNG, and WebP images are allowed"); + } + const mediaType = declaredType as SupportedImageMediaType; + const extension = extname(name).toLowerCase(); + if (!extensions[mediaType].has(extension)) { + throw new HarnessError("IMAGE_TYPE_NOT_ALLOWED", "Filename extension does not match an allowed image type"); + } + const detected = detectImageMediaType(bytes); + if (detected === undefined || detected !== mediaType) { + throw new HarnessError("IMAGE_SIGNATURE_MISMATCH", "Image bytes do not match the declared type"); + } + return mediaType; +} + +export function extensionForMediaType(mediaType: SupportedImageMediaType): string { + if (mediaType === "image/jpeg") return ".jpg"; + if (mediaType === "image/png") return ".png"; + return ".webp"; +} diff --git a/packages/harness/src/index.ts b/packages/harness/src/index.ts index dd19787..aa1ad20 100644 --- a/packages/harness/src/index.ts +++ b/packages/harness/src/index.ts @@ -1,5 +1,8 @@ export * from "./approvals/approval-broker.js"; export * from "./approvals/in-memory-approval-broker.js"; +export * from "./attachments/attachment-store.js"; +export * from "./attachments/file-attachment-store.js"; +export * from "./attachments/image-signatures.js"; export * from "./events/event-store.js"; export * from "./events/in-memory-event-store.js"; export * from "./errors/harness-error.js"; diff --git a/packages/harness/tests/attachments.test.ts b/packages/harness/tests/attachments.test.ts new file mode 100644 index 0000000..33ea670 --- /dev/null +++ b/packages/harness/tests/attachments.test.ts @@ -0,0 +1,121 @@ +import { mkdir, mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { Readable } from "node:stream"; +import sharp from "sharp"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import * as harness from "../src/index.js"; + +type MediaType = "image/jpeg" | "image/png" | "image/webp"; +type Upload = { name: string; mediaType: MediaType; stream: Readable }; +type Attachment = { attachmentId: string; sessionId: string; name: string; mediaType: MediaType; byteSize: number; width: number; height: number; status: "ready" }; +type Store = { + add(sessionId: string, upload: Upload): Promise; + get(sessionId: string, attachmentId: string): Attachment | undefined; + openNormalized(sessionId: string, attachmentId: string): Promise; + remove(sessionId: string, attachmentId: string): Promise; + removeAll(sessionId: string): Promise; + toPiImages(sessionId: string, attachmentIds: string[]): Promise>; +}; + +let root = ""; + +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), "agent-studio-attachments-")); + await mkdir(join(root, "ses_a"), { recursive: true }); + await mkdir(join(root, "ses_b"), { recursive: true }); +}); + +afterEach(async () => { + await rm(root, { recursive: true, force: true }); +}); + +function createStore(options: Record = {}): Store { + const Constructor = Reflect.get(harness, "FileAttachmentStore") as unknown as + | (new (options: Record) => Store) + | undefined; + expect(Constructor, "FileAttachmentStore must be exported").toBeDefined(); + return new Constructor!({ sessionsRoot: root, ...options }); +} + +function upload(name: string, mediaType: MediaType, bytes: Buffer): Upload { + return { name, mediaType, stream: Readable.from(bytes) }; +} + +async function image(format: "jpeg" | "png" | "webp", width = 4, height = 3): Promise { + return sharp({ create: { width, height, channels: 3, background: "#7c3aed" } })[format]().toBuffer(); +} + +describe("FileAttachmentStore", () => { + it.each([ + ["photo.jpg", "image/jpeg", "jpeg"], + ["chart.png", "image/png", "png"], + ["creative.webp", "image/webp", "webp"], + ] as const)("accepts and normalizes %s", async (name, mediaType, format) => { + const store = createStore(); + const item = await store.add("ses_a", upload(name, mediaType, await image(format))); + + expect(item).toMatchObject({ sessionId: "ses_a", name, mediaType, width: 4, height: 3, status: "ready" }); + expect(item.attachmentId).toMatch(/^att_/); + expect((await store.openNormalized("ses_a", item.attachmentId)).byteLength).toBe(item.byteSize); + }); + + it("rejects declared PNG with JPEG bytes", async () => { + await expect(createStore().add("ses_a", upload("fake.png", "image/png", await image("jpeg")))) + .rejects.toMatchObject({ code: "IMAGE_SIGNATURE_MISMATCH" }); + }); + + it("rejects unsupported SVG before decoding", async () => { + const svg = Buffer.from(''); + await expect( + createStore().add("ses_a", { name: "vector.svg", mediaType: "image/png", stream: Readable.from(svg) }), + ).rejects.toMatchObject({ code: "IMAGE_TYPE_NOT_ALLOWED" }); + }); + + it("enforces stream byte and decoded dimension limits", async () => { + const png = await image("png", 10, 10); + await expect(createStore({ maxBytes: 12 }).add("ses_a", upload("large.png", "image/png", png))) + .rejects.toMatchObject({ code: "IMAGE_TOO_LARGE" }); + await expect(createStore({ maxWidth: 5 }).add("ses_a", upload("wide.png", "image/png", png))) + .rejects.toMatchObject({ code: "IMAGE_DIMENSIONS_EXCEEDED" }); + await expect(createStore({ maxPixels: 50 }).add("ses_a", upload("pixels.png", "image/png", png))) + .rejects.toMatchObject({ code: "IMAGE_DIMENSIONS_EXCEEDED" }); + }); + + it("does not expose cross-session attachment metadata or bytes", async () => { + const store = createStore(); + const item = await store.add("ses_a", upload("chart.png", "image/png", await image("png"))); + + expect(store.get("ses_b", item.attachmentId)).toBeUndefined(); + await expect(store.openNormalized("ses_b", item.attachmentId)).rejects.toMatchObject({ code: "ATTACHMENT_NOT_FOUND" }); + }); + + it("maps ready attachments to Pi images, enforces count, and removes idempotently", async () => { + const store = createStore({ maxAttachmentsPerMessage: 1 }); + const first = await store.add("ses_a", upload("one.png", "image/png", await image("png"))); + const second = await store.add("ses_a", upload("two.png", "image/png", await image("png"))); + + const images = await store.toPiImages("ses_a", [first.attachmentId]); + expect(images[0]).toMatchObject({ type: "image", mimeType: "image/png" }); + expect(Buffer.from(images[0]!.data, "base64").byteLength).toBe(first.byteSize); + await expect(store.toPiImages("ses_a", [first.attachmentId, second.attachmentId])).rejects.toMatchObject({ code: "TOO_MANY_ATTACHMENTS" }); + await expect(store.remove("ses_a", first.attachmentId)).resolves.toBe(true); + await expect(store.remove("ses_a", first.attachmentId)).resolves.toBe(false); + await store.removeAll("ses_a"); + expect(store.get("ses_a", second.attachmentId)).toBeUndefined(); + }); + + it("emits safe metadata events without image data", async () => { + const Store = Reflect.get(harness, "InMemoryEventStore") as new () => { + listAfter(sessionId: string): { events: unknown[] }; + }; + const events = new Store(); + const store = createStore({ eventStore: events }); + await store.add("ses_a", upload("chart.png", "image/png", await image("png"))); + + const serialized = JSON.stringify(events.listAfter("ses_a")); + expect(serialized).toContain("attachment.uploaded"); + expect(serialized).not.toContain("base64"); + expect(serialized).not.toContain("iVBOR"); + }); +});