@zgeoff/atc 2.24.0 → 2.26.0

package/src/daemon/imp-provider.ts

@@ -1,6 +1,7 @@
 import { DaemonError } from '../protocol/daemon-error';
 import { LineDecoder } from '../protocol/line-decoder';
 import { isCompiledBinary } from '../shared/is-compiled-binary';
+import type { BrokerAuthHost } from './broker-auth-host';
 import { buildImpName } from './build-imp-name';
 import { buildTarArchive } from './build-tar-archive';
 import type {
@@ -96,6 +97,8 @@ export class ImpProvider implements ExecutionProvider {
   // namespace an impd token's imp patterns are checked against.
   readonly impPrefix: string;
 
+  readonly brokerAuth: BrokerAuthHost;
+
   private readonly port: ImpPort;
 
   private readonly target: ImpTargetOptions;
@@ -106,6 +109,14 @@ export class ImpProvider implements ExecutionProvider {
 
   private readonly hosts = new Map<string, ImpHost>();
 
+  // Each imp's readyings, sleeps, and lease returns run one after another,
+  // in the order they were asked for, so none of them gives back the lease
+  // a readying just took.
+  private readonly turns = new Map<string, Promise<void>>();
+
+  // The readyings of each imp still waiting for their turn.
+  private readonly waitingPrepares = new Map<string, number>();
+
   private readonly atcBinary: string | null;
 
   // Whether impd carries output offsets, once a prepare has read its
@@ -132,6 +143,19 @@ export class ImpProvider implements ExecutionProvider {
       dir,
       atc: target.guestATC ?? (this.atcBinary === null ? null : `${dir}/bin/atc`),
     };
+
+    this.brokerAuth = {
+      impPrefix: this.impPrefix,
+      port,
+      getImpName: (hostKey) => this.getImpName(hostKey),
+      createImp: (hostKey) =>
+        port.createImp({
+          name: this.getImpName(hostKey),
+          ...(target.image === undefined ? {} : { image: target.image }),
+          ...(target.memoryMib === undefined ? {} : { memoryMib: target.memoryMib }),
+        }),
+      destroyImp: (hostKey) => this.destroyHost(hostKey),
+    };
   }
 
   // The imp a host key runs on, under this provider's prefix.
@@ -141,7 +165,25 @@ export class ImpProvider implements ExecutionProvider {
 
   // Creates the host's imp when impd holds none, then takes the daemon's
   // lease, which boots or wakes the imp.
-  readonly prepareHost = async (request: HostRequest): Promise<void> => {
+  readonly prepareHost = (request: HostRequest): Promise<void> => {
+    const name = this.getImpName(request.host);
+
+    this.waitingPrepares.set(name, (this.waitingPrepares.get(name) ?? 0) + 1);
+
+    return this.withHostTurn(name, () => {
+      const left = (this.waitingPrepares.get(name) ?? 1) - 1;
+
+      if (left === 0) {
+        this.waitingPrepares.delete(name);
+      } else {
+        this.waitingPrepares.set(name, left);
+      }
+
+      return this.prepareHostInTurn(request);
+    });
+  };
+
+  private readonly prepareHostInTurn = async (request: HostRequest): Promise<void> => {
     const name = this.getImpName(request.host);
     const label = `atc-${request.daemonID}`;
 
@@ -187,9 +229,11 @@ export class ImpProvider implements ExecutionProvider {
       harnesses: 0,
       renewTimer: null,
       suspending: false,
+      isIdle: null,
     };
 
     host.suspending = false;
+    host.isIdle = request.isIdle ?? null;
 
     this.hosts.set(request.host, host);
     this.startRenewal(host, label);
@@ -201,6 +245,7 @@ export class ImpProvider implements ExecutionProvider {
       harnesses: 0,
       renewTimer: null,
       suspending: false,
+      isIdle: null,
     };
 
     this.hosts.set(spec.host, host);
@@ -225,11 +270,13 @@ export class ImpProvider implements ExecutionProvider {
         cwd: spec.cwd,
         cols: spec.cols,
         rows: spec.rows,
+        ...(spec.requireBroker === true ? { require: ['broker'] } : {}),
       },
       {
         offsets: this.offsets,
         reconnectDelaysMs: this.reconnectDelaysMs,
         ...(reporter === null ? {} : { ready: reporter.forward.listening }),
+        ...(spec.admit === undefined ? {} : { admit: spec.admit }),
         isSuspending: () => host.suspending,
         onDone: () => {
           reporter?.forward.stop();
@@ -271,37 +318,57 @@ export class ImpProvider implements ExecutionProvider {
   // then asks impd to sleep the imp without force. Another owner's lease
   // refuses the sleep: the daemon takes its lease back and the refusal
   // carries what impd showed of the other leases.
-  readonly suspendHost = async (hostKey: string): Promise<void> => {
-    const host = this.hosts.get(hostKey) ?? {
-      name: this.getImpName(hostKey),
-      harnesses: 0,
-      renewTimer: null,
-      suspending: false,
-    };
+  readonly suspendHost = (hostKey: string, isIdle?: () => boolean): Promise<void> =>
+    this.withHostTurn(this.getImpName(hostKey), async () => {
+      const host = this.hosts.get(hostKey) ?? {
+        name: this.getImpName(hostKey),
+        harnesses: 0,
+        renewTimer: null,
+        suspending: false,
+        isIdle: null,
+      };
+
+      this.hosts.set(hostKey, host);
+
+      // A host a harness runs on, or that a readying waits for, is busy.
+      const isBusy = () =>
+        isIdle !== undefined &&
+        (!isIdle() || host.harnesses > 0 || (this.waitingPrepares.get(host.name) ?? 0) > 0);
+
+      if (isBusy()) {
+        throw buildBusyRefusal(host.name);
+      }
 
-    this.hosts.set(hostKey, host);
-    this.stopRenewal(host);
+      this.stopRenewal(host);
 
-    host.suspending = true;
+      host.suspending = true;
 
-    try {
-      if (this.label !== null) {
-        await this.port.releaseLease(host.name, this.label);
-      }
+      try {
+        if (this.label !== null) {
+          await this.port.releaseLease(host.name, this.label);
+        }
 
-      await this.port.suspendImp(host.name);
-    } catch (error) {
-      host.suspending = false;
+        if (isBusy()) {
+          host.suspending = false;
 
-      if (error instanceof ImpPortError && error.code === 'LEASED') {
-        await this.restoreLease(host);
+          await this.restoreLease(host);
 
-        throw buildLeasedRefusal(host.name, error.data);
-      }
+          throw buildBusyRefusal(host.name);
+        }
 
-      throw toHostRefusal(error, host.name);
-    }
-  };
+        await this.port.suspendImp(host.name);
+      } catch (error) {
+        host.suspending = false;
+
+        if (error instanceof ImpPortError && error.code === 'LEASED') {
+          await this.restoreLease(host);
+
+          throw buildLeasedRefusal(host.name, error.data);
+        }
+
+        throw toHostRefusal(error, host.name);
+      }
+    });
 
   // A destroy ends every lease on the imp with it.
   readonly destroyHost = async (hostKey: string): Promise<void> => {
@@ -437,13 +504,7 @@ export class ImpProvider implements ExecutionProvider {
 
     host.renewTimer = setInterval(
       () => {
-        if (host.harnesses === 0) {
-          void this.tryReleaseLease(host);
-
-          return;
-        }
-
-        void this.tryRenewLease(host, label);
+        void this.refreshLease(host, label);
       },
       (this.leaseSeconds * 1000) / 3,
     );
@@ -473,6 +534,28 @@ export class ImpProvider implements ExecutionProvider {
     }
   }
 
+  // Runs a readying, a sleep, or a lease return of an imp after the one
+  // before it there. A turn never starts in its caller's tick, so a harness
+  // end the daemon still applies counts before the turn checks the host.
+  private async withHostTurn<T>(name: string, run: () => Promise<T>): Promise<T> {
+    const before = this.turns.get(name) ?? Promise.resolve();
+    const turn = Promise.withResolvers<void>();
+
+    this.turns.set(name, turn.promise);
+
+    try {
+      await before;
+
+      return await run();
+    } finally {
+      turn.resolve();
+
+      if (this.turns.get(name) === turn.promise) {
+        this.turns.delete(name);
+      }
+    }
+  }
+
   private async restoreLease(host: ImpHost): Promise<void> {
     if (this.label === null) {
       return;
@@ -487,19 +570,45 @@ export class ImpProvider implements ExecutionProvider {
     this.startRenewal(host, this.label);
   }
 
-  private async tryReleaseLease(host: ImpHost): Promise<boolean> {
-    this.stopRenewal(host);
+  // A host with no harness gives its lease back, unless a launch readies it,
+  // which keeps the lease renewed.
+  private async refreshLease(host: ImpHost, label: string): Promise<void> {
+    if (host.harnesses === 0) {
+      const released = await this.tryReleaseLease(host);
 
-    if (this.label === null) {
-      return false;
+      if (released) {
+        return;
+      }
     }
 
-    try {
-      return await this.port.releaseLease(host.name, this.label);
-    } catch {
-      return false;
+    if (host.renewTimer !== null) {
+      await this.tryRenewLease(host, label);
     }
   }
+
+  // Gives an imp's lease back once nothing runs or starts there, in the
+  // imp's turn, so it never drops a lease a readying took for a launch.
+  private tryReleaseLease(host: ImpHost): Promise<boolean> {
+    return this.withHostTurn(host.name, async () => {
+      const isBusy =
+        host.harnesses > 0 ||
+        host.suspending ||
+        (this.waitingPrepares.get(host.name) ?? 0) > 0 ||
+        host.isIdle?.() === false;
+
+      if (isBusy || this.label === null) {
+        return false;
+      }
+
+      this.stopRenewal(host);
+
+      try {
+        return await this.port.releaseLease(host.name, this.label);
+      } catch {
+        return false;
+      }
+    });
+  }
 }
 
 // What the provider follows of one imp: its name, how many harnesses run
@@ -510,6 +619,10 @@ interface ImpHost {
   harnesses: number;
   renewTimer: ReturnType<typeof setInterval> | null;
   suspending: boolean;
+
+  // The daemon's check that nothing runs or starts on the host, from its
+  // latest readying; null until one gives it.
+  isIdle: (() => boolean) | null;
 }
 
 function buildImpSessionName(sessionID: string): string {
@@ -571,6 +684,14 @@ function buildGuestEnv(env: Readonly<Record<string, string>>): Record<string, st
   return { TERM: 'xterm-256color', LANG: 'C.UTF-8', PATH: GUEST_PATH, ...own };
 }
 
+function buildBusyRefusal(name: string): DaemonError {
+  return new DaemonError(
+    'host_unavailable',
+    `imp ${name} stays awake: a session runs there or is starting there`,
+    { provider: 'imp', problem: 'busy' },
+  );
+}
+
 function buildLeasedRefusal(name: string, data: unknown): DaemonError {
   const record = typeof data === 'object' && data !== null ? data : {};
   const leases = 'leases' in record && Array.isArray(record.leases) ? record.leases : [];

package/src/daemon/local-pty-provider.ts

@@ -1,5 +1,6 @@
 import { mkdir } from 'node:fs/promises';
 import { spawn } from 'bun-pty';
+import { DaemonError } from '../protocol/daemon-error';
 import { collectCleanEnv } from '../shared/collect-clean-env';
 import type {
   CommandResult,
@@ -39,7 +40,17 @@ export class LocalPTYProvider implements ExecutionProvider {
 
   // The daemon's machine keeps no process the daemon lets go of, so a
   // detach ends the harness as a kill does, after dropping every listener.
+  // It has no credential broker, so a harness that requires one never
+  // starts here.
   readonly spawnHarness = (spec: HarnessSpec): HarnessHandle => {
+    if (spec.requireBroker === true) {
+      throw new DaemonError(
+        'auth_target_unsupported',
+        "the daemon's own machine has no credential broker to start the harness behind",
+        { provider: 'local-pty' },
+      );
+    }
+
     const pty = spawn(spec.bin, [...spec.args], {
       name: 'xterm-256color',
       cols: spec.cols,

package/src/daemon/make-non-blocking-log.ts

@@ -0,0 +1,91 @@
+import { write } from 'node:fs';
+import { promisify } from 'node:util';
+
+// The most bytes of lines that wait behind the write in flight; a line past
+// it is dropped and counted.
+const MAX_QUEUED_BYTES = 64 * 1024;
+
+// The node:fs write without a callback, which Bun runs on its thread pool.
+const writeAsync = promisify(write);
+
+/**
+ * Makes a log that writes each line to a file descriptor without ever
+ * blocking the event loop, so a reader that stops reading stalls nothing
+ * but the log. Each write goes through the asynchronous node:fs write,
+ * which Bun runs on its thread pool, one write at a time. Lines that
+ * arrive meanwhile queue up to 64 KiB; a line past that is dropped, and the
+ * next write after the drops starts with one `atc log dropped=N` line. A
+ * failed write drops its lines and counts them the same way, and the next
+ * line logged writes again.
+ */
+export function makeNonBlockingLog(fd: number): (line: string) => void {
+  let queue: string[] = [];
+  let queuedBytes = 0;
+  let dropped = 0;
+  let writing = false;
+
+  const recordLine = (line: string): boolean => {
+    const text = `${line}\n`;
+    const bytes = Buffer.byteLength(text);
+
+    if (queuedBytes + bytes > MAX_QUEUED_BYTES) {
+      return false;
+    }
+
+    queue.push(text);
+
+    queuedBytes += bytes;
+
+    return true;
+  };
+
+  const writeNext = (): void => {
+    if (queue.length === 0 && dropped === 0) {
+      writing = false;
+
+      return;
+    }
+
+    writing = true;
+
+    const lines = dropped > 0 ? [`atc log dropped=${String(dropped)}\n`, ...queue] : queue;
+
+    // The lines this write stands for, the ones the dropped line counts
+    // included, which a failed write drops again.
+    const count = dropped + queue.length;
+
+    queue = [];
+    queuedBytes = 0;
+    dropped = 0;
+    void writeChunk(Buffer.from(lines.join('')), count);
+  };
+
+  const writeChunk = async (chunk: Buffer, lineCount: number): Promise<void> => {
+    let rest = chunk;
+
+    try {
+      while (rest.length > 0) {
+        const result = await writeAsync(fd, rest);
+
+        rest = rest.subarray(result.bytesWritten);
+      }
+    } catch {
+      dropped += lineCount;
+      writing = false;
+
+      return;
+    }
+
+    writeNext();
+  };
+
+  return (line) => {
+    if (!recordLine(line)) {
+      dropped++;
+    }
+
+    if (!writing) {
+      writeNext();
+    }
+  };
+}

package/src/daemon/refusal-log.ts

@@ -0,0 +1,163 @@
+import { formatLogField } from './format-log-field';
+
+// Why the TCP listener refused a handshake: a missing token, a wrong one, a
+// handshake that would wait while the cap of delayed handshakes is full, a
+// socket that closed while its handshake waited, a line other than one
+// handshake before the handshake passed, or a line over the size cap
+// before it passed.
+type HandshakeRefusalReason =
+  | 'missing_token'
+  | 'unauthorized'
+  | 'delay_cap_full'
+  | 'closed_during_delay'
+  | 'unexpected_line'
+  | 'line_too_long';
+
+export type Refusal =
+  | {
+      readonly event: 'handshake_refused';
+      readonly peer: string;
+      readonly reason: HandshakeRefusalReason;
+    }
+  | {
+      readonly event: 'principal_refused';
+      readonly peer: string;
+    };
+
+interface RefusalLogOptions {
+  readonly log: (line: string) => void;
+  readonly now: () => number;
+
+  // How long a window of repeated refusals lasts before the next one is
+  // logged again.
+  readonly intervalMs: number;
+
+  // How many windows are tracked at once; a refusal that would start one
+  // past the cap counts toward one overflow window instead.
+  readonly maxWindows: number;
+}
+
+interface RefusalWindow {
+  readonly startedAt: number;
+
+  // The refusals in the window after the first, which no line holds yet.
+  pending: number;
+}
+
+// The line of the window that counts refusals past the cap of windows.
+const OVERFLOW_LINE = 'atc tcp event=refused peer=overflow';
+
+/**
+ * Logs TCP listener refusals as `key=value` lines, at most one per window
+ * for each peer and kind of refusal: a handshake refusal per reason, and a
+ * principal refusal, which never holds the principal it gave. The first
+ * refusal of a window logs at once with `count=1`; later ones in the window
+ * are counted, and a line with their `count` follows once the window ends,
+ * so the counts of every line sum to every refusal. While the cap of
+ * windows is full, a refusal that would start a new window counts toward
+ * one overflow window with `peer=overflow` instead, so a flood across many
+ * peers logs no more lines than the cap allows. Ended windows are logged on
+ * the next refusal from any peer and on a drain.
+ */
+export class RefusalLog {
+  private readonly opts: RefusalLogOptions;
+
+  // Keyed by the line without its count, in the order the windows started.
+  private readonly windows = new Map<string, RefusalWindow>();
+
+  private overflow: RefusalWindow | null = null;
+
+  constructor(opts: RefusalLogOptions) {
+    this.opts = opts;
+  }
+
+  record(refusal: Refusal): void {
+    const now = this.opts.now();
+
+    this.drainEnded(now);
+
+    const line = formatRefusal(refusal);
+    const window = this.windows.get(line);
+
+    if (window !== undefined) {
+      window.pending++;
+
+      return;
+    }
+
+    if (this.windows.size >= this.opts.maxWindows) {
+      this.recordOverflow(now);
+
+      return;
+    }
+
+    this.windows.set(line, { startedAt: now, pending: 0 });
+    this.opts.log(`${line} count=1`);
+  }
+
+  /**
+   * Logs the pending count of every window and forgets them all.
+   */
+  drain(): void {
+    for (const [line, window] of this.windows) {
+      this.logPending(line, window);
+    }
+
+    this.windows.clear();
+
+    if (this.overflow !== null) {
+      this.logPending(OVERFLOW_LINE, this.overflow);
+
+      this.overflow = null;
+    }
+  }
+
+  // Windows start in map order and all last the same interval, so the
+  // ended ones lead the map.
+  private drainEnded(now: number): void {
+    if (this.overflow !== null && this.hasEnded(this.overflow, now)) {
+      this.logPending(OVERFLOW_LINE, this.overflow);
+
+      this.overflow = null;
+    }
+
+    for (const [line, window] of this.windows) {
+      if (!this.hasEnded(window, now)) {
+        return;
+      }
+
+      this.windows.delete(line);
+      this.logPending(line, window);
+    }
+  }
+
+  private hasEnded(window: Readonly<RefusalWindow>, now: number): boolean {
+    return now - window.startedAt >= this.opts.intervalMs;
+  }
+
+  private logPending(line: string, window: Readonly<RefusalWindow>): void {
+    if (window.pending > 0) {
+      this.opts.log(`${line} count=${window.pending}`);
+    }
+  }
+
+  private recordOverflow(now: number): void {
+    if (this.overflow !== null) {
+      this.overflow.pending++;
+
+      return;
+    }
+
+    this.overflow = { startedAt: now, pending: 0 };
+
+    this.opts.log(`${OVERFLOW_LINE} count=1`);
+  }
+}
+
+function formatRefusal(refusal: Refusal): string {
+  const peer = formatLogField(refusal.peer);
+
+  return refusal.event === 'handshake_refused'
+    ? `atc tcp event=handshake_refused peer=${peer} reason=${refusal.reason}`
+    : `atc tcp event=principal_refused peer=${peer} principal=unlisted`;
+}

package/src/daemon/restore-fleet.ts

@@ -185,7 +185,8 @@ function isSameSession(s: Session, entry: FleetEntry): boolean {
 
 // The refusals that leave one session without a terminal: its target is
 // misconfigured, gone from the config, changed, has no provider here, or
-// cannot start a terminal, or its agent refuses every start.
+// cannot start a terminal, or its agent refuses every start, or its
+// host's runtime auth binding refuses the launch, its broker included.
 const REFUSED_ADOPT_CODES: ReadonlySet<string> = new Set([
   'unsupported_operation',
   'unknown_target',
@@ -195,6 +196,20 @@ const REFUSED_ADOPT_CODES: ReadonlySet<string> = new Set([
   'host_unavailable',
   'auth_not_configured',
   'auth_target_unsupported',
+  'auth_impd_too_old',
+  'auth_token_scope',
+  'auth_token_too_broad',
+  'auth_imp_out_of_scope',
+  'auth_secret_not_grantable',
+  'auth_secret_mismatch',
+  'auth_runtime_mismatch',
+  'auth_grant_missing',
+  'auth_grants_mismatch',
+  'auth_rebind_required',
+  'auth_blocked',
+  'auth_binding_invalid',
+  'broker_not_ready',
+  'auth_placeholder_unsupported',
 ]);
 
 async function tryAdoptTerminal(