agl 1.0.17 → 21.0.2

package/ts/classes.uploadmanager.ts

@@ -0,0 +1,240 @@
+import * as plugins from './plugins.js';
+import {
+  controllerMaxAttachmentPromptSuffixBytes,
+  controllerMaxDraftAttachmentBytes,
+  controllerMaxDraftAttachmentTotalBytes,
+  controllerMaxDraftAttachments,
+  type IControllerDraftAttachment,
+} from '../ts_interfaces/interfaces.js';
+
+const maximumActiveUploads = 64;
+const maximumActiveUploadBytes = 64 * 1024 * 1024;
+
+export interface IControllerUploadHandle {
+  id: string;
+  directory: string;
+  paths: string[];
+  promptSuffix: string;
+  bytes: number;
+  cleanup(): Promise<void>;
+}
+
+interface IActiveUpload {
+  handle: IControllerUploadHandle;
+  cleanupPromise?: Promise<void>;
+}
+
+export class ControllerUploadManager {
+  private rootPromise?: Promise<string>;
+  private readonly activeUploads = new Map<string, IActiveUpload>();
+  private readonly pendingCreates = new Set<Promise<IControllerUploadHandle>>();
+  private activeUploadCount = 0;
+  private activeBytes = 0;
+  private closed = false;
+  private closePromise?: Promise<void>;
+
+  public create(
+    attachmentsArg: readonly IControllerDraftAttachment[],
+  ): Promise<IControllerUploadHandle | undefined> {
+    if (attachmentsArg.length === 0) return Promise.resolve(undefined);
+    if (this.closed) throw new Error('The controller upload manager is closed.');
+    if (attachmentsArg.length > controllerMaxDraftAttachments) {
+      throw new Error('The attachment count exceeds the controller limit.');
+    }
+    if (this.activeUploadCount >= maximumActiveUploads) {
+      throw new Error('The controller has too many active uploads.');
+    }
+
+    const decoded = attachmentsArg.map((attachment, index) => {
+      if (
+        attachment.name.length === 0
+        || attachment.name !== plugins.path.basename(attachment.name)
+        || Buffer.byteLength(attachment.name, 'utf8') > 255
+        // eslint-disable-next-line no-control-regex
+        || /[\x00-\x1f\x7f/\\]/.test(attachment.name)
+      ) {
+        throw new Error(`Attachment ${index + 1} has an unsafe file name.`);
+      }
+      if (
+        attachment.mediaType.length === 0
+        || Buffer.byteLength(attachment.mediaType, 'utf8') > 129
+        // eslint-disable-next-line no-control-regex
+        || /[\x00-\x20\x7f]/.test(attachment.mediaType)
+      ) throw new Error(`Attachment ${index + 1} has an unsafe media type.`);
+      const data = Buffer.from(attachment.dataBase64, 'base64');
+      if (data.toString('base64') !== attachment.dataBase64) {
+        throw new Error(`Attachment ${index + 1} is not canonical base64.`);
+      }
+      if (data.byteLength !== attachment.size) {
+        throw new Error(`Attachment ${index + 1} size does not match its content.`);
+      }
+      if (data.byteLength > controllerMaxDraftAttachmentBytes) {
+        throw new Error(`Attachment ${index + 1} exceeds the controller byte limit.`);
+      }
+      return { attachment, data };
+    });
+    const bytes = decoded.reduce((sum, entry) => sum + entry.data.byteLength, 0);
+    if (bytes > controllerMaxDraftAttachmentTotalBytes) {
+      throw new Error('The attachments exceed the controller total byte limit.');
+    }
+    if (this.activeBytes + bytes > maximumActiveUploadBytes) {
+      throw new Error('The controller active upload byte limit is exhausted.');
+    }
+
+    this.activeUploadCount += 1;
+    this.activeBytes += bytes;
+    const createPromise = this.createReserved(decoded, bytes);
+    this.pendingCreates.add(createPromise);
+    void createPromise.finally(() => this.pendingCreates.delete(createPromise)).catch(() => undefined);
+    return createPromise;
+  }
+
+  public async close(): Promise<void> {
+    if (this.closePromise) return this.closePromise;
+    this.closed = true;
+    const closePromise = this.closeInternal();
+    this.closePromise = closePromise;
+    try {
+      await closePromise;
+    } catch (errorArg) {
+      if (this.closePromise === closePromise) this.closePromise = undefined;
+      throw errorArg;
+    }
+  }
+
+  private async createReserved(
+    decodedArg: Array<{
+      attachment: IControllerDraftAttachment;
+      data: Buffer;
+    }>,
+    bytesArg: number,
+  ): Promise<IControllerUploadHandle> {
+    const id = plugins.crypto.randomBytes(16).toString('base64url');
+    let directory: string | undefined;
+    let admitted = false;
+    try {
+      const root = await this.ensureRoot();
+      directory = await plugins.fs.promises.mkdtemp(plugins.path.join(root, 'message-'));
+      await plugins.fs.promises.chmod(directory, 0o700);
+      const paths: string[] = [];
+      for (const [index, entry] of decodedArg.entries()) {
+        const filePath = plugins.path.join(directory, `${index + 1}-${entry.attachment.name}`);
+        const relative = plugins.path.relative(directory, filePath);
+        if (
+          relative === '..'
+          || relative.startsWith(`..${plugins.path.sep}`)
+          || plugins.path.isAbsolute(relative)
+        ) throw new Error('An attachment path escaped its upload directory.');
+        const flags = plugins.fs.constants.O_WRONLY
+          | plugins.fs.constants.O_CREAT
+          | plugins.fs.constants.O_EXCL
+          | (plugins.fs.constants.O_NOFOLLOW ?? 0);
+        const file = await plugins.fs.promises.open(filePath, flags, 0o600);
+        try {
+          await file.writeFile(entry.data);
+        } finally {
+          await file.close();
+        }
+        paths.push(filePath);
+      }
+      const promptSuffix = `\n\n[Attached files - stored on disk; read them with filesystem tools]\n${decodedArg
+        .map((entry, index) => `- ${paths[index]} (${entry.attachment.mediaType}, ${entry.data.byteLength} bytes)`)
+        .join('\n')}`;
+      if (Buffer.byteLength(promptSuffix, 'utf8') > controllerMaxAttachmentPromptSuffixBytes) {
+        throw new Error('The generated attachment prompt suffix exceeds its byte limit.');
+      }
+      const handle: IControllerUploadHandle = {
+        id,
+        directory,
+        paths,
+        promptSuffix,
+        bytes: bytesArg,
+        cleanup: () => this.cleanup(id),
+      };
+      this.activeUploads.set(id, { handle });
+      admitted = true;
+      return handle;
+    } finally {
+      if (!admitted) {
+        if (directory) {
+          await plugins.fs.promises.rm(directory, { recursive: true, force: true }).catch(() => undefined);
+        }
+        this.activeUploadCount -= 1;
+        this.activeBytes -= bytesArg;
+      }
+    }
+  }
+
+  private async closeInternal(): Promise<void> {
+    await Promise.allSettled([...this.pendingCreates]);
+    const results = await Promise.allSettled(
+      [...this.activeUploads.values()].map((entry) => entry.handle.cleanup()),
+    );
+    const errors = results
+      .filter((result): result is PromiseRejectedResult => result.status === 'rejected')
+      .map((result) => result.reason);
+    if (this.rootPromise) {
+      try {
+        const root = await this.rootPromise;
+        await plugins.fs.promises.rm(root, { recursive: true, force: true });
+        this.activeUploads.clear();
+        this.activeUploadCount = 0;
+        this.activeBytes = 0;
+        errors.length = 0;
+      } catch (errorArg) {
+        errors.push(errorArg);
+      }
+    }
+    if (errors.length > 0) throw new AggregateError(errors, 'Controller upload cleanup failed.');
+  }
+
+  private async ensureRoot(): Promise<string> {
+    if (!this.rootPromise) {
+      this.rootPromise = (async () => {
+        let root: string | undefined;
+        try {
+          root = await plugins.fs.promises.mkdtemp(
+            plugins.path.join(plugins.os.tmpdir(), 'harness-controller-uploads-'),
+          );
+          await plugins.fs.promises.chmod(root, 0o700);
+          return await plugins.fs.promises.realpath(root);
+        } catch (errorArg) {
+          if (root) {
+            await plugins.fs.promises.rm(root, { recursive: true, force: true })
+              .catch(() => undefined);
+          }
+          throw errorArg;
+        }
+      })();
+    }
+    const rootPromise = this.rootPromise;
+    try {
+      return await rootPromise;
+    } catch (errorArg) {
+      if (this.rootPromise === rootPromise) this.rootPromise = undefined;
+      throw errorArg;
+    }
+  }
+
+  private async cleanup(idArg: string): Promise<void> {
+    const active = this.activeUploads.get(idArg);
+    if (!active) return;
+    if (active.cleanupPromise) return active.cleanupPromise;
+    const cleanupPromise = (async () => {
+      await plugins.fs.promises.rm(active.handle.directory, { recursive: true, force: true });
+      if (this.activeUploads.get(idArg) === active) {
+        this.activeUploads.delete(idArg);
+        this.activeUploadCount -= 1;
+        this.activeBytes -= active.handle.bytes;
+      }
+    })();
+    active.cleanupPromise = cleanupPromise;
+    try {
+      await cleanupPromise;
+    } finally {
+      if (this.activeUploads.get(idArg) === active && active.cleanupPromise === cleanupPromise) {
+        active.cleanupPromise = undefined;
+      }
+    }
+  }
+}