Harden custom image setup terminal fallback

server/src/services/environment-custom-image-terminal-sessions.test.ts

@@ -1,7 +1,9 @@
 import { describe, expect, it } from "vitest";
 import {
+  EnvironmentCustomImageTerminalConnectionRegistry,
   EnvironmentCustomImageTerminalSessionStore,
   parseCustomImageSetupSshCommand,
+  validateCustomImageSetupSshPayload,
 } from "./environment-custom-image-terminal-sessions.js";
 
 describe("parseCustomImageSetupSshCommand", () => {
@@ -23,6 +25,14 @@ describe("parseCustomImageSetupSshCommand", () => {
     });
   });
 
+  it("parses the redacted Daytona createSshAccess command shape", () => {
+    expect(parseCustomImageSetupSshCommand("ssh dtca_redacted-token.123@ssh.app.daytona.io")).toEqual({
+      username: "dtca_redacted-token.123",
+      host: "ssh.app.daytona.io",
+      port: 22,
+    });
+  });
+
   it("rejects unsupported or ambiguous SSH command shapes", () => {
     expect(parseCustomImageSetupSshCommand("scp user@example.test")).toBeNull();
     expect(parseCustomImageSetupSshCommand("ssh user@example.test -i key")).toBeNull();
@@ -33,6 +43,70 @@ describe("parseCustomImageSetupSshCommand", () => {
   });
 });
 
+describe("validateCustomImageSetupSshPayload", () => {
+  it("returns parsed SSH connection details and a valid payload expiry", () => {
+    const result = validateCustomImageSetupSshPayload({
+      type: "ssh",
+      command: "ssh dtca_redacted-token.123@ssh.app.daytona.io",
+      expiresAt: "2026-06-25T20:15:00.000Z",
+    }, new Date("2026-06-25T20:00:00.000Z"));
+
+    expect(result).toEqual({
+      ok: true,
+      ssh: {
+        username: "dtca_redacted-token.123",
+        host: "ssh.app.daytona.io",
+        port: 22,
+      },
+      connectionExpiresAt: new Date("2026-06-25T20:15:00.000Z"),
+    });
+  });
+
+  it("returns redacted fallback failures for unsupported payloads and parser failures", () => {
+    expect(validateCustomImageSetupSshPayload({
+      type: "browser_terminal",
+      command: "ssh secret-token@203.0.113.10",
+    }, new Date("2026-06-25T20:00:00.000Z"))).toMatchObject({
+      ok: false,
+      status: 422,
+      code: "unsupported_payload",
+      message: "Setup session terminal connections require an SSH connection payload.",
+    });
+    expect(validateCustomImageSetupSshPayload({
+      type: "ssh",
+      command: "ssh secret-token@203.0.113.10 -i /tmp/private-key",
+    }, new Date("2026-06-25T20:00:00.000Z"))).toMatchObject({
+      ok: false,
+      status: 422,
+      code: "unsupported_command",
+      message: "Setup session SSH payload uses an unsupported command shape.",
+    });
+  });
+
+  it("returns clear failures for invalid and expired payload expiries", () => {
+    expect(validateCustomImageSetupSshPayload({
+      type: "ssh",
+      command: "ssh token@example.test",
+      expiresAt: "not-a-date",
+    }, new Date("2026-06-25T20:00:00.000Z"))).toMatchObject({
+      ok: false,
+      status: 422,
+      code: "invalid_expiry",
+      message: "Setup session SSH payload has an invalid expiry.",
+    });
+    expect(validateCustomImageSetupSshPayload({
+      type: "ssh",
+      command: "ssh token@example.test",
+      expiresAt: "2026-06-25T19:59:59.000Z",
+    }, new Date("2026-06-25T20:00:00.000Z"))).toMatchObject({
+      ok: false,
+      status: 409,
+      code: "expired_payload",
+      message: "Setup session SSH connection payload has expired.",
+    });
+  });
+});
+
 describe("EnvironmentCustomImageTerminalSessionStore", () => {
   it("mints opaque tokens and caps expiry by setup and provider payload expiry", () => {
     const store = new EnvironmentCustomImageTerminalSessionStore();
@@ -49,12 +123,75 @@ describe("EnvironmentCustomImageTerminalSessionStore", () => {
     });
 
     expect(minted.token).toHaveLength(43);
+    expect(minted.session.id).toMatch(/^[0-9a-f-]{36}$/);
     expect(minted.session.expiresAt.toISOString()).toBe("2026-06-25T20:02:00.000Z");
-    expect(store.get(minted.token, new Date("2026-06-25T20:01:59.000Z"))?.ssh).toEqual({
+    expect(store.get({
+      id: minted.session.id,
+      token: minted.token,
+    }, new Date("2026-06-25T20:01:59.000Z"))?.ssh).toEqual({
       username: "ssh-token-secret",
       host: "203.0.113.10",
       port: 2222,
     });
-    expect(store.get(minted.token, new Date("2026-06-25T20:02:00.000Z"))).toBeNull();
+    expect(store.get({
+      id: minted.session.id,
+      token: "wrong-token",
+    }, new Date("2026-06-25T20:01:59.000Z"))).toBeNull();
+    expect(store.get({
+      id: minted.session.id,
+      token: minted.token,
+    }, new Date("2026-06-25T20:02:00.000Z"))).toBeNull();
+  });
+
+  it("deletes all tokens for a setup session", () => {
+    const store = new EnvironmentCustomImageTerminalSessionStore();
+    const now = new Date("2026-06-25T20:00:00.000Z");
+    const first = store.create({
+      setupSessionId: "session-1",
+      companyId: "company-1",
+      environmentId: "env-1",
+      provider: "daytona",
+      ssh: { username: "one", host: "example.test", port: 22 },
+      setupExpiresAt: new Date("2026-06-25T20:30:00.000Z"),
+      now,
+    });
+    const second = store.create({
+      setupSessionId: "session-1",
+      companyId: "company-1",
+      environmentId: "env-1",
+      provider: "daytona",
+      ssh: { username: "two", host: "example.test", port: 22 },
+      setupExpiresAt: new Date("2026-06-25T20:30:00.000Z"),
+      now,
+    });
+
+    expect(store.deleteBySetupSessionId("session-1")).toBe(2);
+    expect(store.get({ id: first.session.id, token: first.token }, now)).toBeNull();
+    expect(store.get({ id: second.session.id, token: second.token }, now)).toBeNull();
+  });
+});
+
+describe("EnvironmentCustomImageTerminalConnectionRegistry", () => {
+  it("closes active terminal connections for a setup session", () => {
+    const registry = new EnvironmentCustomImageTerminalConnectionRegistry();
+    const firstReasons: string[] = [];
+    const secondReasons: string[] = [];
+    const removeFirst = registry.add({
+      setupSessionId: "session-1",
+      close: (reason) => firstReasons.push(reason),
+    });
+    registry.add({
+      setupSessionId: "session-1",
+      close: (reason) => secondReasons.push(reason),
+    });
+
+    expect(registry.closeBySetupSessionId("session-1", "setup_finished")).toBe(2);
+    expect(firstReasons).toEqual(["setup_finished"]);
+    expect(secondReasons).toEqual(["setup_finished"]);
+
+    removeFirst();
+    expect(registry.closeBySetupSessionId("session-1", "setup_cancelled")).toBe(1);
+    expect(firstReasons).toEqual(["setup_finished"]);
+    expect(secondReasons).toEqual(["setup_finished", "setup_cancelled"]);
   });
 });

server/src/services/environment-custom-image-terminal-sessions.ts

@@ -1,4 +1,4 @@
-import { createHash, randomBytes } from "node:crypto";
+import { createHash, randomBytes, randomUUID } from "node:crypto";
 
 const DEFAULT_TERMINAL_SESSION_TOKEN_TTL_MS = 5 * 60 * 1000;
 const TERMINAL_SESSION_TOKEN_BYTES = 32;
@@ -9,7 +9,28 @@ export interface ParsedCustomImageSetupSshCommand {
   port: number;
 }
 
+export type EnvironmentCustomImageTerminalPayloadValidationFailureCode =
+  | "unsupported_payload"
+  | "missing_command"
+  | "unsupported_command"
+  | "invalid_expiry"
+  | "expired_payload";
+
+export type EnvironmentCustomImageTerminalPayloadValidationResult =
+  | {
+      ok: true;
+      ssh: ParsedCustomImageSetupSshCommand;
+      connectionExpiresAt: Date | null;
+    }
+  | {
+      ok: false;
+      status: 409 | 422;
+      code: EnvironmentCustomImageTerminalPayloadValidationFailureCode;
+      message: string;
+    };
+
 export interface EnvironmentCustomImageTerminalSessionRecord {
+  id: string;
   setupSessionId: string;
   companyId: string;
   environmentId: string;
@@ -25,6 +46,11 @@ export interface MintedEnvironmentCustomImageTerminalSession {
   session: EnvironmentCustomImageTerminalSessionRecord;
 }
 
+interface StoredEnvironmentCustomImageTerminalSession {
+  tokenHash: string;
+  session: EnvironmentCustomImageTerminalSessionRecord;
+}
+
 function parsePort(value: string): number | null {
   if (!/^\d{1,5}$/.test(value)) return null;
   const port = Number(value);
@@ -68,6 +94,73 @@ export function parseCustomImageSetupSshCommand(command: string): ParsedCustomIm
   return null;
 }
 
+function readConnectionPayload(payload: unknown): Record<string, unknown> | null {
+  return payload && typeof payload === "object" && !Array.isArray(payload)
+    ? payload as Record<string, unknown>
+    : null;
+}
+
+function readDate(value: unknown): Date | null {
+  if (!value) return null;
+  const date = value instanceof Date ? value : typeof value === "string" ? new Date(value) : null;
+  return date && !Number.isNaN(date.getTime()) ? date : null;
+}
+
+export function validateCustomImageSetupSshPayload(
+  payload: unknown,
+  now: Date,
+): EnvironmentCustomImageTerminalPayloadValidationResult {
+  const record = readConnectionPayload(payload);
+  if (!record || record.type !== "ssh") {
+    return {
+      ok: false,
+      status: 422,
+      code: "unsupported_payload",
+      message: "Setup session terminal connections require an SSH connection payload.",
+    };
+  }
+
+  const command = typeof record.command === "string" ? record.command.trim() : "";
+  if (!command) {
+    return {
+      ok: false,
+      status: 422,
+      code: "missing_command",
+      message: "Setup session SSH payload is missing a supported command.",
+    };
+  }
+
+  const ssh = parseCustomImageSetupSshCommand(command);
+  if (!ssh) {
+    return {
+      ok: false,
+      status: 422,
+      code: "unsupported_command",
+      message: "Setup session SSH payload uses an unsupported command shape.",
+    };
+  }
+
+  const connectionExpiresAt = readDate(record.expiresAt);
+  if (record.expiresAt != null && !connectionExpiresAt) {
+    return {
+      ok: false,
+      status: 422,
+      code: "invalid_expiry",
+      message: "Setup session SSH payload has an invalid expiry.",
+    };
+  }
+  if (connectionExpiresAt && connectionExpiresAt.getTime() <= now.getTime()) {
+    return {
+      ok: false,
+      status: 409,
+      code: "expired_payload",
+      message: "Setup session SSH connection payload has expired.",
+    };
+  }
+
+  return { ok: true, ssh, connectionExpiresAt };
+}
+
 function hashTerminalSessionToken(token: string): string {
   return createHash("sha256").update(token).digest("hex");
 }
@@ -84,7 +177,7 @@ function toValidFutureDate(value: Date | string | null | undefined, now: Date):
 }
 
 export class EnvironmentCustomImageTerminalSessionStore {
-  private readonly sessions = new Map<string, EnvironmentCustomImageTerminalSessionRecord>();
+  private readonly sessionsById = new Map<string, StoredEnvironmentCustomImageTerminalSession>();
 
   create(input: {
     setupSessionId: string;
@@ -106,7 +199,9 @@ export class EnvironmentCustomImageTerminalSessionStore {
     ].filter((date): date is Date => date !== null);
     const expiresAt = minDate(candidateExpirations);
     const token = randomBytes(TERMINAL_SESSION_TOKEN_BYTES).toString("base64url");
+    const id = randomUUID();
     const session: EnvironmentCustomImageTerminalSessionRecord = {
+      id,
       setupSessionId: input.setupSessionId,
       companyId: input.companyId,
       environmentId: input.environmentId,
@@ -116,32 +211,57 @@ export class EnvironmentCustomImageTerminalSessionStore {
       createdAt: now,
       expiresAt,
     };
-    this.sessions.set(hashTerminalSessionToken(token), session);
+    this.sessionsById.set(id, {
+      tokenHash: hashTerminalSessionToken(token),
+      session,
+    });
     return { token, session };
   }
 
-  get(token: string, now = new Date()): EnvironmentCustomImageTerminalSessionRecord | null {
-    if (!token) return null;
-    const key = hashTerminalSessionToken(token);
-    const session = this.sessions.get(key) ?? null;
-    if (!session) return null;
-    if (session.expiresAt.getTime() <= now.getTime()) {
-      this.sessions.delete(key);
+  get(input: { id: string; token: string }, now = new Date()): EnvironmentCustomImageTerminalSessionRecord | null {
+    if (!input.id || !input.token) return null;
+    const stored = this.sessionsById.get(input.id) ?? null;
+    if (!stored) return null;
+    if (stored.tokenHash !== hashTerminalSessionToken(input.token)) return null;
+    if (stored.session.expiresAt.getTime() <= now.getTime()) {
+      this.sessionsById.delete(input.id);
       return null;
     }
-    return session;
+    return stored.session;
   }
 
-  delete(token: string): boolean {
-    if (!token) return false;
-    return this.sessions.delete(hashTerminalSessionToken(token));
+  getById(id: string, now = new Date()): EnvironmentCustomImageTerminalSessionRecord | null {
+    if (!id) return null;
+    const stored = this.sessionsById.get(id) ?? null;
+    if (!stored) return null;
+    if (stored.session.expiresAt.getTime() <= now.getTime()) {
+      this.sessionsById.delete(id);
+      return null;
+    }
+    return stored.session;
+  }
+
+  delete(id: string): boolean {
+    if (!id) return false;
+    return this.sessionsById.delete(id);
+  }
+
+  deleteBySetupSessionId(setupSessionId: string): number {
+    if (!setupSessionId) return 0;
+    let removed = 0;
+    for (const [id, stored] of this.sessionsById) {
+      if (stored.session.setupSessionId !== setupSessionId) continue;
+      this.sessionsById.delete(id);
+      removed += 1;
+    }
+    return removed;
   }
 
   cleanupExpired(now = new Date()): number {
     let removed = 0;
-    for (const [key, session] of this.sessions) {
-      if (session.expiresAt.getTime() <= now.getTime()) {
-        this.sessions.delete(key);
+    for (const [id, stored] of this.sessionsById) {
+      if (stored.session.expiresAt.getTime() <= now.getTime()) {
+        this.sessionsById.delete(id);
         removed += 1;
       }
     }
@@ -149,9 +269,60 @@ export class EnvironmentCustomImageTerminalSessionStore {
   }
 
   clear(): void {
-    this.sessions.clear();
+    this.sessionsById.clear();
   }
 }
 
 export const environmentCustomImageTerminalSessionStore =
   new EnvironmentCustomImageTerminalSessionStore();
+
+export type EnvironmentCustomImageTerminalConnectionClose = (reason: string) => void;
+
+export class EnvironmentCustomImageTerminalConnectionRegistry {
+  private readonly connectionsBySetupSessionId = new Map<string, Set<EnvironmentCustomImageTerminalConnectionClose>>();
+
+  add(input: {
+    setupSessionId: string;
+    close: EnvironmentCustomImageTerminalConnectionClose;
+  }): () => void {
+    const existing = this.connectionsBySetupSessionId.get(input.setupSessionId);
+    const connections = existing ?? new Set<EnvironmentCustomImageTerminalConnectionClose>();
+    connections.add(input.close);
+    if (!existing) {
+      this.connectionsBySetupSessionId.set(input.setupSessionId, connections);
+    }
+
+    return () => {
+      connections.delete(input.close);
+      if (connections.size === 0) {
+        this.connectionsBySetupSessionId.delete(input.setupSessionId);
+      }
+    };
+  }
+
+  closeBySetupSessionId(setupSessionId: string, reason: string): number {
+    const connections = this.connectionsBySetupSessionId.get(setupSessionId);
+    if (!connections) return 0;
+    let closed = 0;
+    for (const close of [...connections]) {
+      close(reason);
+      closed += 1;
+    }
+    return closed;
+  }
+
+  closeAll(reason: string): number {
+    let closed = 0;
+    for (const setupSessionId of [...this.connectionsBySetupSessionId.keys()]) {
+      closed += this.closeBySetupSessionId(setupSessionId, reason);
+    }
+    return closed;
+  }
+
+  clear(): void {
+    this.connectionsBySetupSessionId.clear();
+  }
+}
+
+export const environmentCustomImageTerminalConnectionRegistry =
+  new EnvironmentCustomImageTerminalConnectionRegistry();

server/src/services/index.ts

@@ -90,9 +90,12 @@ export {
   environmentCustomImageService,
 } from "./environment-custom-images.js";
 export {
+  environmentCustomImageTerminalConnectionRegistry,
   environmentCustomImageTerminalSessionStore,
+  EnvironmentCustomImageTerminalConnectionRegistry,
   EnvironmentCustomImageTerminalSessionStore,
   parseCustomImageSetupSshCommand,
+  type EnvironmentCustomImageTerminalConnectionClose,
   type EnvironmentCustomImageTerminalSessionRecord,
   type MintedEnvironmentCustomImageTerminalSession,
   type ParsedCustomImageSetupSshCommand,