diff --git a/packages/shared/src/api.ts b/packages/shared/src/api.ts index ace37e41441..e3878097007 100644 --- a/packages/shared/src/api.ts +++ b/packages/shared/src/api.ts @@ -12,6 +12,7 @@ export const API = { environmentCustomImageSetupSessions: `${API_PREFIX}/environments/:environmentId/custom-image-setup-sessions`, environmentCustomImageSetupSession: `${API_PREFIX}/environment-custom-image-setup-sessions/:sessionId`, environmentCustomImageSetupSessionTerminalToken: `${API_PREFIX}/environment-custom-image-setup-sessions/:sessionId/terminal-session-token`, + environmentCustomImageSetupSessionTerminalWs: `${API_PREFIX}/environment-custom-image-setup-sessions/:sessionId/terminal/ws`, environmentCustomImageSetupSessionFinish: `${API_PREFIX}/environment-custom-image-setup-sessions/:sessionId/finish`, environmentCustomImageSetupSessionCancel: `${API_PREFIX}/environment-custom-image-setup-sessions/:sessionId/cancel`, issues: `${API_PREFIX}/issues`, diff --git a/packages/shared/src/validators/environment-custom-images.ts b/packages/shared/src/validators/environment-custom-images.ts index 65e5984ebf9..6c66f89f059 100644 --- a/packages/shared/src/validators/environment-custom-images.ts +++ b/packages/shared/src/validators/environment-custom-images.ts @@ -101,11 +101,13 @@ export type CreateEnvironmentCustomImageTerminalSessionToken = z.infer; export const environmentCustomImageTerminalSessionTokenSchema = z.object({ + id: z.string().min(1), token: z.string().min(32), expiresAt: isoDateTime, setupSessionId: z.string().min(1), environmentId: z.string().min(1), connectionType: z.literal("ssh"), + websocketPath: z.string().min(1), }).strict(); export type EnvironmentCustomImageTerminalSessionToken = z.infer; diff --git a/server/package.json b/server/package.json index cc41dc5818f..ca8fbb6956a 100644 --- a/server/package.json +++ b/server/package.json @@ -77,6 +77,7 @@ "pino-http": "^10.4.0", "pino-pretty": "^13.1.3", "sharp": "^0.35.2", + "ssh2": "^1.17.0", "ws": "^8.19.0", "zod": "^3.24.2" }, diff --git a/server/src/__tests__/environment-custom-image-routes.test.ts b/server/src/__tests__/environment-custom-image-routes.test.ts index ceda94a2425..35aebec21e6 100644 --- a/server/src/__tests__/environment-custom-image-routes.test.ts +++ b/server/src/__tests__/environment-custom-image-routes.test.ts @@ -2,7 +2,10 @@ import express from "express"; import request from "supertest"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { environmentRoutes } from "../routes/environments.js"; -import { environmentCustomImageTerminalSessionStore } from "../services/environment-custom-image-terminal-sessions.js"; +import { + environmentCustomImageTerminalConnectionRegistry, + environmentCustomImageTerminalSessionStore, +} from "../services/environment-custom-image-terminal-sessions.js"; const now = new Date("2026-06-25T20:00:00.000Z"); @@ -232,6 +235,7 @@ describe("environment customImage setup routes", () => { Object.values(mockSecretService).forEach((mock) => mock.mockReset()); mockLogActivity.mockReset(); environmentCustomImageTerminalSessionStore.clear(); + environmentCustomImageTerminalConnectionRegistry.clear(); mockInstanceSettingsService.listCompanyIds.mockResolvedValue(["company-1"]); mockEnvironmentService.getById.mockResolvedValue(createEnvironment()); @@ -355,8 +359,12 @@ describe("environment customImage setup routes", () => { environmentId: "env-1", connectionType: "ssh", }); + expect(typeof res.body.id).toBe("string"); expect(typeof res.body.token).toBe("string"); expect(typeof res.body.expiresAt).toBe("string"); + expect(res.body.websocketPath).toContain( + `/api/environment-custom-image-setup-sessions/session-1/terminal/ws?terminalSessionId=${encodeURIComponent(res.body.id)}`, + ); expect(mockEnvironmentCustomImageService.refreshSetupSession).toHaveBeenCalledWith({ sessionId: "session-1", includeConnectionPayload: true, @@ -472,6 +480,43 @@ describe("environment customImage setup routes", () => { expect(loggedActivityJson()).not.toContain("ssh-token-secret"); }); + it("rejects invalid and expired terminal payload expiries without minting tokens", async () => { + mockEnvironmentCustomImageService.getSessionById.mockResolvedValue(createSession({ + expiresAt: futureDate(), + })); + mockEnvironmentCustomImageService.refreshSetupSession.mockResolvedValueOnce({ + session: createSession({ expiresAt: futureDate() }), + connectionPayload: { + type: "ssh", + command: "ssh ssh-token-secret@ssh.app.daytona.io", + expiresAt: "not-a-date", + }, + }); + + const invalidExpiry = await request(createApp(boardActor())) + .post("/api/environment-custom-image-setup-sessions/session-1/terminal-session-token") + .send({}); + + expect(invalidExpiry.status).toBe(422); + expect(JSON.stringify(invalidExpiry.body)).not.toContain("ssh-token-secret"); + + mockEnvironmentCustomImageService.refreshSetupSession.mockResolvedValueOnce({ + session: createSession({ expiresAt: futureDate() }), + connectionPayload: { + type: "ssh", + command: "ssh ssh-token-secret@ssh.app.daytona.io", + expiresAt: new Date(Date.now() - 60_000).toISOString(), + }, + }); + + const expiredPayload = await request(createApp(boardActor())) + .post("/api/environment-custom-image-setup-sessions/session-1/terminal-session-token") + .send({}); + + expect(expiredPayload.status).toBe(409); + expect(JSON.stringify(expiredPayload.body)).not.toContain("ssh-token-secret"); + }); + it("denies agent API key actors before customImage state or payloads are read", async () => { const app = createApp(agentActor()); const start = await request(app) @@ -517,6 +562,19 @@ describe("environment customImage setup routes", () => { it("finishes and promotes a template while logging redacted template details", async () => { mockEnvironmentCustomImageService.getSessionById.mockResolvedValue(createSession()); + const terminal = environmentCustomImageTerminalSessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "token-secret", host: "203.0.113.10", port: 2222 }, + setupExpiresAt: futureDate(), + }); + const closeReasons: string[] = []; + environmentCustomImageTerminalConnectionRegistry.add({ + setupSessionId: "session-1", + close: (reason) => closeReasons.push(reason), + }); const res = await request(createApp(boardActor())) .post("/api/environment-custom-image-setup-sessions/session-1/finish") @@ -528,6 +586,11 @@ describe("environment customImage setup routes", () => { sessionId: "session-1", metadata: { safeLabel: "done" }, }); + expect(environmentCustomImageTerminalSessionStore.get({ + id: terminal.session.id, + token: terminal.token, + })).toBeNull(); + expect(closeReasons).toEqual(["setup_finished"]); const activity = loggedActivityJson(); expect(activity).not.toContain("captured-template-secret"); expect(activity).not.toContain("snapshot-secret-ref"); diff --git a/server/src/__tests__/environment-custom-image-terminal-ws.test.ts b/server/src/__tests__/environment-custom-image-terminal-ws.test.ts new file mode 100644 index 00000000000..8ffd1742985 --- /dev/null +++ b/server/src/__tests__/environment-custom-image-terminal-ws.test.ts @@ -0,0 +1,427 @@ +import { EventEmitter } from "node:events"; +import { createServer, type Server as HttpServer } from "node:http"; +import { createRequire } from "node:module"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../services/environment-custom-images.js", () => ({ + environmentCustomImageService: vi.fn(() => { + throw new Error("test must inject a custom image service"); + }), +})); + +import { + setupEnvironmentCustomImageTerminalWebSocketServer, + type EnvironmentCustomImageSshConnector, + type EnvironmentCustomImageSshShell, +} from "../realtime/environment-custom-image-terminal-ws.js"; +import { + EnvironmentCustomImageTerminalConnectionRegistry, + EnvironmentCustomImageTerminalSessionStore, +} from "../services/environment-custom-image-terminal-sessions.js"; + +const require = createRequire(import.meta.url); +const { WebSocket } = require("ws") as { + WebSocket: new (url: string) => { + readyState: number; + send(data: string): void; + close(): void; + on(event: "open", listener: () => void): void; + on(event: "message", listener: (data: Buffer | string) => void): void; + on(event: "close", listener: () => void): void; + on(event: "error", listener: (err: Error) => void): void; + }; +}; + +class FakeSshShell extends EventEmitter implements EnvironmentCustomImageSshShell { + writes: string[] = []; + resizes: Array<{ cols: number; rows: number }> = []; + closeCalls = 0; + + write(data: string): void { + this.writes.push(data); + } + + resize(cols: number, rows: number): void { + this.resizes.push({ cols, rows }); + } + + close(): void { + this.closeCalls += 1; + } + + onData(listener: (data: string) => void): void { + this.on("data", listener); + } + + onClose(listener: () => void): void { + this.on("close", listener); + } + + onError(listener: (err: Error) => void): void { + this.on("ssh-error", listener); + } + + emitData(data: string) { + this.emit("data", data); + } + + emitSshClose() { + this.emit("close"); + } + + emitSshError(err: Error) { + this.emit("ssh-error", err); + } +} + +function futureDate(minutes = 60) { + return new Date(Date.now() + minutes * 60 * 1000); +} + +function createSession(overrides: Record = {}) { + return { + id: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + status: "waiting_for_user", + expiresAt: futureDate(), + ...overrides, + }; +} + +async function flushPromises() { + await new Promise((resolve) => setImmediate(resolve)); +} + +async function listen(server: HttpServer) { + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("server did not listen on a TCP port"); + return address.port; +} + +async function closeServer(server: HttpServer) { + if (!server.listening) return; + await new Promise((resolve) => server.close(() => resolve())); +} + +function waitForOpen(ws: InstanceType) { + return new Promise((resolve, reject) => { + ws.on("open", resolve); + ws.on("error", reject); + }); +} + +function waitForError(ws: InstanceType) { + return new Promise((resolve) => { + ws.on("error", resolve); + }); +} + +function waitForClose(ws: InstanceType) { + return new Promise((resolve) => { + ws.on("close", resolve); + }); +} + +function waitForJsonMessage>( + ws: InstanceType, + predicate: (frame: T) => boolean, +) { + return new Promise((resolve) => { + ws.on("message", (data: Buffer | string) => { + const text = typeof data === "string" ? data : data.toString("utf8"); + const parsed = JSON.parse(text) as T; + if (predicate(parsed)) resolve(parsed); + }); + }); +} + +function terminalUrl(port: number, input: { setupSessionId?: string; terminalSessionId: string; token: string }) { + const setupSessionId = input.setupSessionId ?? "session-1"; + return `ws://127.0.0.1:${port}/api/environment-custom-image-setup-sessions/${setupSessionId}/terminal/ws` + + `?terminalSessionId=${encodeURIComponent(input.terminalSessionId)}&token=${encodeURIComponent(input.token)}` + + "&cols=100&rows=30"; +} + +describe("custom image terminal websocket bridge", () => { + let servers: HttpServer[] = []; + let sessionStore: EnvironmentCustomImageTerminalSessionStore; + let connectionRegistry: EnvironmentCustomImageTerminalConnectionRegistry; + let fakeShell: FakeSshShell; + let customImages: { + getSessionById: ReturnType; + refreshSetupSession: ReturnType; + }; + let connector: EnvironmentCustomImageSshConnector & { connect: ReturnType }; + + beforeEach(() => { + servers = []; + sessionStore = new EnvironmentCustomImageTerminalSessionStore(); + connectionRegistry = new EnvironmentCustomImageTerminalConnectionRegistry(); + fakeShell = new FakeSshShell(); + customImages = { + getSessionById: vi.fn(async () => createSession()), + refreshSetupSession: vi.fn(async () => ({ + session: createSession(), + connectionPayload: { + type: "ssh", + command: "ssh fresh-token@fresh.example.test -p 2200", + expiresAt: futureDate(30).toISOString(), + }, + })), + }; + connector = { + connect: vi.fn(async () => fakeShell), + }; + }); + + afterEach(async () => { + for (const server of servers) { + server.emit("close"); + } + await Promise.all(servers.map((server) => closeServer(server))); + }); + + async function startHarness() { + const server = createServer(); + servers.push(server); + setupEnvironmentCustomImageTerminalWebSocketServer(server, {} as never, { + customImageService: customImages, + sessionStore, + connectionRegistry, + sshConnector: connector, + }); + const port = await listen(server); + return { server, port }; + } + + it("rejects invalid and expired terminal credentials before refreshing provider payloads", async () => { + const { port } = await startHarness(); + const invalid = new WebSocket(terminalUrl(port, { + terminalSessionId: "missing", + token: "bad-token", + })); + + const invalidError = await waitForError(invalid); + expect(invalidError.message).toContain("422"); + expect(customImages.refreshSetupSession).not.toHaveBeenCalled(); + + const expired = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "old", host: "old.example.test", port: 22 }, + setupExpiresAt: new Date(Date.now() + 60 * 60 * 1000), + connectionExpiresAt: new Date(Date.now() - 60 * 1000), + now: new Date(Date.now() - 2 * 60 * 1000), + }); + const expiredWs = new WebSocket(terminalUrl(port, { + terminalSessionId: expired.session.id, + token: expired.token, + })); + + const expiredError = await waitForError(expiredWs); + expect(expiredError.message).toContain("422"); + expect(customImages.refreshSetupSession).not.toHaveBeenCalled(); + }); + + it("rejects unsupported refreshed payloads before opening an SSH bridge", async () => { + const { port } = await startHarness(); + customImages.refreshSetupSession.mockResolvedValueOnce({ + session: createSession(), + connectionPayload: { + type: "browser_terminal", + command: "ssh provider-secret@203.0.113.10", + }, + }); + const unsupportedPayload = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "old-token", host: "old.example.test", port: 22 }, + setupExpiresAt: futureDate(), + }); + + const unsupportedWs = new WebSocket(terminalUrl(port, { + terminalSessionId: unsupportedPayload.session.id, + token: unsupportedPayload.token, + })); + const unsupportedError = await waitForError(unsupportedWs); + expect(unsupportedError.message).toContain("422"); + expect(sessionStore.get({ + id: unsupportedPayload.session.id, + token: unsupportedPayload.token, + })).toBeNull(); + + customImages.refreshSetupSession.mockResolvedValueOnce({ + session: createSession(), + connectionPayload: { + type: "ssh", + command: "ssh provider-secret@203.0.113.10 -i /tmp/private-key", + }, + }); + const unsupportedCommand = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "old-token", host: "old.example.test", port: 22 }, + setupExpiresAt: futureDate(), + }); + + const unsupportedCommandWs = new WebSocket(terminalUrl(port, { + terminalSessionId: unsupportedCommand.session.id, + token: unsupportedCommand.token, + })); + const unsupportedCommandError = await waitForError(unsupportedCommandWs); + expect(unsupportedCommandError.message).toContain("422"); + expect(sessionStore.get({ + id: unsupportedCommand.session.id, + token: unsupportedCommand.token, + })).toBeNull(); + expect(connector.connect).not.toHaveBeenCalled(); + }); + + it("bridges websocket input, SSH output, and resize frames through a fake shell", async () => { + const { port } = await startHarness(); + const minted = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "old-token", host: "old.example.test", port: 22 }, + setupExpiresAt: futureDate(), + }); + const ws = new WebSocket(terminalUrl(port, { + terminalSessionId: minted.session.id, + token: minted.token, + })); + const readyPromise = waitForJsonMessage(ws, (frame) => frame.type === "ready"); + + await waitForOpen(ws); + const ready = await readyPromise; + expect(ready).toMatchObject({ + type: "ready", + setupSessionId: "session-1", + terminalSessionId: minted.session.id, + }); + expect(customImages.refreshSetupSession).toHaveBeenCalledWith({ + sessionId: "session-1", + includeConnectionPayload: true, + }); + expect(connector.connect).toHaveBeenCalledWith({ + ssh: { username: "fresh-token", host: "fresh.example.test", port: 2200 }, + term: "xterm-256color", + cols: 100, + rows: 30, + }); + + ws.send(JSON.stringify({ type: "input", data: "echo ok\r" })); + await flushPromises(); + expect(fakeShell.writes).toEqual(["echo ok\r"]); + + const outputPromise = waitForJsonMessage(ws, (frame) => frame.type === "output"); + fakeShell.emitData("shell output\r\n"); + await expect(outputPromise).resolves.toMatchObject({ + type: "output", + data: "shell output\r\n", + }); + + ws.send(JSON.stringify({ type: "resize", cols: 120, rows: 40 })); + await flushPromises(); + expect(fakeShell.resizes).toEqual([{ cols: 120, rows: 40 }]); + + const closePromise = waitForClose(ws); + ws.close(); + await closePromise; + expect(fakeShell.closeCalls).toBeGreaterThan(0); + expect(sessionStore.get({ id: minted.session.id, token: minted.token })).toBeNull(); + }); + + it("sends a redacted fallback error when SSH bridge connection fails", async () => { + const { port } = await startHarness(); + connector.connect.mockRejectedValueOnce(new Error("provider secret should not leak")); + const minted = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "old-token", host: "old.example.test", port: 22 }, + setupExpiresAt: futureDate(), + }); + const ws = new WebSocket(terminalUrl(port, { + terminalSessionId: minted.session.id, + token: minted.token, + })); + const errorPromise = waitForJsonMessage(ws, (frame) => frame.type === "error"); + const closePromise = waitForClose(ws); + + await waitForOpen(ws); + await expect(errorPromise).resolves.toMatchObject({ + type: "error", + message: "SSH terminal connection failed.", + }); + await closePromise; + expect(sessionStore.get({ id: minted.session.id, token: minted.token })).toBeNull(); + }); + + it("cleans up and closes the websocket when the SSH shell errors", async () => { + const { port } = await startHarness(); + const minted = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "token", host: "example.test", port: 22 }, + setupExpiresAt: futureDate(), + }); + const ws = new WebSocket(terminalUrl(port, { + terminalSessionId: minted.session.id, + token: minted.token, + })); + const readyPromise = waitForJsonMessage(ws, (frame) => frame.type === "ready"); + + await waitForOpen(ws); + await readyPromise; + const errorPromise = waitForJsonMessage(ws, (frame) => frame.type === "error"); + const closePromise = waitForClose(ws); + fakeShell.emitSshError(new Error("provider secret should not leak")); + await expect(errorPromise).resolves.toMatchObject({ + type: "error", + message: "SSH terminal connection failed.", + }); + await closePromise; + expect(sessionStore.get({ id: minted.session.id, token: minted.token })).toBeNull(); + expect(fakeShell.closeCalls).toBeGreaterThan(0); + }); + + it("closes active terminal sessions on server shutdown", async () => { + const { server, port } = await startHarness(); + const minted = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "token", host: "example.test", port: 22 }, + setupExpiresAt: futureDate(), + }); + const ws = new WebSocket(terminalUrl(port, { + terminalSessionId: minted.session.id, + token: minted.token, + })); + const readyPromise = waitForJsonMessage(ws, (frame) => frame.type === "ready"); + + await waitForOpen(ws); + await readyPromise; + const closePromise = waitForClose(ws); + server.emit("close"); + await closePromise; + + expect(sessionStore.get({ id: minted.session.id, token: minted.token })).toBeNull(); + expect(fakeShell.closeCalls).toBeGreaterThan(0); + }); +}); diff --git a/server/src/index.ts b/server/src/index.ts index 93e867f8958..7254ee53662 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -33,6 +33,7 @@ import detectPort from "detect-port"; import { createApp } from "./app.js"; import { loadConfig } from "./config.js"; import { logger } from "./middleware/logger.js"; +import { setupEnvironmentCustomImageTerminalWebSocketServer } from "./realtime/environment-custom-image-terminal-ws.js"; import { setupLiveEventsWebSocketServer } from "./realtime/live-events-ws.js"; import { feedbackService, @@ -700,6 +701,9 @@ export async function startServer(): Promise { process.env.PAPERCLIP_RUNTIME_API_CANDIDATES_JSON = JSON.stringify(runtimeApiCandidates); process.env.PAPERCLIP_API_URL = configuredApiUrl; + setupEnvironmentCustomImageTerminalWebSocketServer(server, db as any, { + pluginWorkerManager, + }); setupLiveEventsWebSocketServer(server, db as any, { deploymentMode: config.deploymentMode, resolveSessionFromHeaders, diff --git a/server/src/realtime/environment-custom-image-terminal-ws.ts b/server/src/realtime/environment-custom-image-terminal-ws.ts new file mode 100644 index 00000000000..a7ea9529ace --- /dev/null +++ b/server/src/realtime/environment-custom-image-terminal-ws.ts @@ -0,0 +1,624 @@ +import type { IncomingMessage, Server as HttpServer } from "node:http"; +import { createRequire } from "node:module"; +import type { Duplex } from "node:stream"; +import type { Db } from "@paperclipai/db"; +import { conflict, unprocessable } from "../errors.js"; +import { logger } from "../middleware/logger.js"; +import { environmentCustomImageService } from "../services/environment-custom-images.js"; +import { + environmentCustomImageTerminalConnectionRegistry, + environmentCustomImageTerminalSessionStore, + validateCustomImageSetupSshPayload, + type EnvironmentCustomImageTerminalConnectionRegistry, + type EnvironmentCustomImageTerminalPayloadValidationResult, + type EnvironmentCustomImageTerminalSessionRecord, + type EnvironmentCustomImageTerminalSessionStore, + type ParsedCustomImageSetupSshCommand, +} from "../services/environment-custom-image-terminal-sessions.js"; +import type { PluginWorkerManager } from "../services/plugin-worker-manager.js"; + +interface TerminalWsSocket { + readyState: number; + send(data: string): void; + close(code?: number, reason?: string): void; + terminate(): void; + on(event: "message", listener: (data: unknown) => void): void; + on(event: "close", listener: () => void): void; + on(event: "error", listener: (err: Error) => void): void; +} + +interface TerminalWsServer { + clients: Set; + on(event: "connection", listener: (socket: TerminalWsSocket, req: IncomingMessage) => void): void; + on(event: "close", listener: () => void): void; + close(callback?: (err?: Error) => void): void; + handleUpgrade( + req: IncomingMessage, + socket: Duplex, + head: Buffer, + callback: (ws: TerminalWsSocket) => void, + ): void; + emit(event: "connection", ws: TerminalWsSocket, req: IncomingMessage): boolean; +} + +interface SetupSessionSnapshot { + id: string; + companyId: string; + environmentId: string; + provider: string; + status: string; + expiresAt: Date | string | null; +} + +interface CustomImageTerminalService { + getSessionById(sessionId: string): Promise; + refreshSetupSession(input: { + sessionId: string; + includeConnectionPayload: true; + }): Promise<{ + session: SetupSessionSnapshot; + connectionPayload: unknown; + }>; +} + +export interface EnvironmentCustomImageSshShell { + write(data: string): void; + resize(cols: number, rows: number): void; + close(): void; + onData(listener: (data: string) => void): void; + onClose(listener: () => void): void; + onError(listener: (err: Error) => void): void; +} + +export interface EnvironmentCustomImageSshConnector { + connect(input: { + ssh: ParsedCustomImageSetupSshCommand; + term: string; + cols: number; + rows: number; + }): Promise; +} + +interface TerminalUpgradeContext { + setupSessionId: string; + terminalSession: EnvironmentCustomImageTerminalSessionRecord; + ssh: ParsedCustomImageSetupSshCommand; + initialCols: number; + initialRows: number; +} + +interface IncomingMessageWithTerminalContext extends IncomingMessage { + paperclipWebSocketHandled?: boolean; + paperclipTerminalUpgradeContext?: TerminalUpgradeContext; +} + +const require = createRequire(import.meta.url); +const { WebSocket, WebSocketServer } = require("ws") as { + WebSocket: { OPEN: number }; + WebSocketServer: new (opts: { noServer: boolean }) => TerminalWsServer; +}; + +function isWritableUpgradeSocket(socket: Duplex) { + const maybeWritableState = socket as Duplex & { writable?: boolean; writableEnded?: boolean; writableDestroyed?: boolean }; + return !socket.destroyed && maybeWritableState.writable !== false && !maybeWritableState.writableEnded && !maybeWritableState.writableDestroyed; +} + +function closeUpgradeSocket(socket: Duplex) { + if (!socket.destroyed) { + socket.destroy(); + } +} + +function rejectUpgrade(socket: Duplex, statusLine: string, message: string) { + const safe = message.replace(/[\r\n]+/g, " ").trim(); + if (!isWritableUpgradeSocket(socket)) { + closeUpgradeSocket(socket); + return; + } + + try { + socket.once("finish", () => closeUpgradeSocket(socket)); + socket.end(`HTTP/1.1 ${statusLine}\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n${safe}`); + } catch (err) { + logger.warn({ errorName: err instanceof Error ? err.name : typeof err }, "failed to reject custom image terminal websocket upgrade"); + closeUpgradeSocket(socket); + } +} + +function parseTerminalPath(pathname: string): { setupSessionId: string } | null { + const match = pathname.match(/^\/api\/environment-custom-image-setup-sessions\/([^/]+)\/terminal\/ws$/); + if (!match) return null; + + try { + const setupSessionId = decodeURIComponent(match[1] ?? ""); + return setupSessionId ? { setupSessionId } : null; + } catch { + return null; + } +} + +function parseTerminalDimension(value: string | null, fallback: number) { + if (!value) return fallback; + if (!/^\d{1,4}$/.test(value)) return fallback; + const parsed = Number(value); + return Number.isInteger(parsed) && parsed > 0 && parsed <= 9999 ? parsed : fallback; +} + +function readDate(value: unknown): Date | null { + if (!value) return null; + const date = value instanceof Date ? value : new Date(String(value)); + return Number.isNaN(date.getTime()) ? null : date; +} + +function requireFutureSetupExpiry(session: { expiresAt: Date | string | null }, now: Date): Date { + const expiresAt = readDate(session.expiresAt); + if (!expiresAt || expiresAt.getTime() <= now.getTime()) { + throw conflict("Environment customImage setup session has expired."); + } + return expiresAt; +} + +function statusLineForError(err: unknown) { + const status = typeof err === "object" && err !== null && "status" in err + ? Number((err as { status?: unknown }).status) + : 500; + if (status === 400) return "400 Bad Request"; + if (status === 401) return "401 Unauthorized"; + if (status === 403) return "403 Forbidden"; + if (status === 404) return "404 Not Found"; + if (status === 409) return "409 Conflict"; + if (status === 422) return "422 Unprocessable Entity"; + return "500 Internal Server Error"; +} + +function safeErrorName(err: unknown) { + return err instanceof Error ? err.name : typeof err; +} + +function terminalPayloadValidationError( + failure: Extract, +): Error { + return failure.status === 409 ? conflict(failure.message) : unprocessable(failure.message); +} + +function decodeClientMessage(data: unknown): string { + if (typeof data === "string") return data; + if (Buffer.isBuffer(data)) return data.toString("utf8"); + if (data instanceof ArrayBuffer) return Buffer.from(data).toString("utf8"); + if (ArrayBuffer.isView(data)) return Buffer.from(data.buffer, data.byteOffset, data.byteLength).toString("utf8"); + return ""; +} + +function sendJson(socket: TerminalWsSocket, frame: Record) { + if (socket.readyState !== WebSocket.OPEN) return; + socket.send(JSON.stringify(frame)); +} + +function closeClient(socket: TerminalWsSocket, code: number, reason: string) { + if (socket.readyState !== WebSocket.OPEN) return; + socket.close(code, reason.slice(0, 120)); +} + +function handleClientFrame( + socket: TerminalWsSocket, + shell: EnvironmentCustomImageSshShell | null, + raw: unknown, +) { + const text = decodeClientMessage(raw); + if (!text) return; + + let parsed: unknown; + try { + parsed = JSON.parse(text); + } catch { + shell?.write(text); + return; + } + + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + return; + } + const frame = parsed as Record; + if (frame.type === "input") { + if (typeof frame.data === "string") { + shell?.write(frame.data); + } + return; + } + if (frame.type === "resize") { + const cols = typeof frame.cols === "number" ? frame.cols : null; + const rows = typeof frame.rows === "number" ? frame.rows : null; + if ( + shell + && Number.isInteger(cols) + && Number.isInteger(rows) + && cols > 0 + && rows > 0 + && cols <= 9999 + && rows <= 9999 + ) { + shell.resize(cols, rows); + } + return; + } + + sendJson(socket, { type: "error", message: "Unsupported terminal frame." }); +} + +async function validateTerminalUpgrade(input: { + setupSessionId: string; + terminalSessionId: string; + token: string; + now: Date; + sessionStore: EnvironmentCustomImageTerminalSessionStore; + customImages: CustomImageTerminalService; +}): Promise { + const terminalSession = input.sessionStore.get({ + id: input.terminalSessionId, + token: input.token, + }, input.now); + if (!terminalSession || terminalSession.setupSessionId !== input.setupSessionId) { + throw unprocessable("Invalid terminal session token."); + } + + const storedSetupSession = await input.customImages.getSessionById(input.setupSessionId); + if (!storedSetupSession) { + input.sessionStore.delete(terminalSession.id); + throw unprocessable("Invalid terminal setup session."); + } + if ( + storedSetupSession.companyId !== terminalSession.companyId + || storedSetupSession.environmentId !== terminalSession.environmentId + || storedSetupSession.provider !== terminalSession.provider + ) { + input.sessionStore.delete(terminalSession.id); + throw unprocessable("Invalid terminal setup session."); + } + + const refreshed = await input.customImages.refreshSetupSession({ + sessionId: input.setupSessionId, + includeConnectionPayload: true, + }); + if ( + refreshed.session.id !== terminalSession.setupSessionId + || refreshed.session.companyId !== terminalSession.companyId + || refreshed.session.environmentId !== terminalSession.environmentId + || refreshed.session.provider !== terminalSession.provider + ) { + input.sessionStore.delete(terminalSession.id); + throw unprocessable("Invalid terminal setup session."); + } + if (refreshed.session.status !== "waiting_for_user") { + input.sessionStore.delete(terminalSession.id); + throw conflict(`Cannot open terminal for setup status "${refreshed.session.status}".`); + } + requireFutureSetupExpiry(refreshed.session, input.now); + + const payloadValidation = validateCustomImageSetupSshPayload(refreshed.connectionPayload, input.now); + if (!payloadValidation.ok) { + input.sessionStore.delete(terminalSession.id); + throw terminalPayloadValidationError(payloadValidation); + } + + return { ...terminalSession, ssh: payloadValidation.ssh }; +} + +class Ssh2Shell implements EnvironmentCustomImageSshShell { + constructor( + private readonly client: { + end(): void; + destroy?(): void; + on(event: "close", listener: () => void): void; + on(event: "error", listener: (err: Error) => void): void; + }, + private readonly stream: { + write(data: string): void; + end(): void; + destroy?(): void; + setWindow?(rows: number, cols: number, height: number, width: number): void; + on(event: "data", listener: (data: Buffer | string) => void): void; + on(event: "close", listener: () => void): void; + on(event: "error", listener: (err: Error) => void): void; + }, + ) {} + + write(data: string): void { + this.stream.write(data); + } + + resize(cols: number, rows: number): void { + this.stream.setWindow?.(rows, cols, 0, 0); + } + + close(): void { + try { + this.stream.end(); + } catch { + this.stream.destroy?.(); + } + try { + this.client.end(); + } catch { + this.client.destroy?.(); + } + } + + onData(listener: (data: string) => void): void { + this.stream.on("data", (data: Buffer | string) => { + listener(typeof data === "string" ? data : data.toString("utf8")); + }); + } + + onClose(listener: () => void): void { + this.stream.on("close", listener); + this.client.on("close", listener); + } + + onError(listener: (err: Error) => void): void { + this.stream.on("error", listener); + this.client.on("error", listener); + } +} + +export function createSsh2EnvironmentCustomImageSshConnector(): EnvironmentCustomImageSshConnector { + return { + connect: async ({ ssh, term, cols, rows }) => { + const { Client } = require("ssh2") as { + Client: new () => { + once(event: "ready", listener: () => void): void; + once(event: "error", listener: (err: Error) => void): void; + on(event: "close", listener: () => void): void; + on(event: "error", listener: (err: Error) => void): void; + connect(config: Record): void; + shell( + window: { term: string; cols: number; rows: number }, + callback: (err: Error | undefined, stream: ConstructorParameters[1]) => void, + ): void; + end(): void; + destroy?(): void; + }; + }; + + return await new Promise((resolve, reject) => { + const client = new Client(); + let settled = false; + const fail = (err: Error) => { + if (settled) return; + settled = true; + try { + client.end(); + } catch { + client.destroy?.(); + } + reject(err); + }; + + client.once("ready", () => { + client.shell({ term, cols, rows }, (err, stream) => { + if (err || !stream) { + fail(err ?? new Error("SSH shell failed to open.")); + return; + } + if (settled) return; + settled = true; + resolve(new Ssh2Shell(client, stream)); + }); + }); + client.once("error", fail); + client.connect({ + host: ssh.host, + port: ssh.port, + username: ssh.username, + readyTimeout: 20000, + keepaliveInterval: 15000, + keepaliveCountMax: 3, + }); + }); + }, + }; +} + +export function setupEnvironmentCustomImageTerminalWebSocketServer( + server: HttpServer, + db: Db, + opts: { + pluginWorkerManager?: PluginWorkerManager; + customImageService?: CustomImageTerminalService; + sessionStore?: EnvironmentCustomImageTerminalSessionStore; + connectionRegistry?: EnvironmentCustomImageTerminalConnectionRegistry; + sshConnector?: EnvironmentCustomImageSshConnector; + } = {}, +) { + const wss = new WebSocketServer({ noServer: true }); + const customImages = opts.customImageService ?? environmentCustomImageService(db, { + pluginWorkerManager: opts.pluginWorkerManager, + }); + const sessionStore = opts.sessionStore ?? environmentCustomImageTerminalSessionStore; + const connectionRegistry = opts.connectionRegistry ?? environmentCustomImageTerminalConnectionRegistry; + const sshConnector = opts.sshConnector ?? createSsh2EnvironmentCustomImageSshConnector(); + + wss.on("connection", (socket: TerminalWsSocket, req: IncomingMessage) => { + const context = (req as IncomingMessageWithTerminalContext).paperclipTerminalUpgradeContext; + if (!context) { + socket.close(1008, "missing context"); + return; + } + + let shell: EnvironmentCustomImageSshShell | null = null; + let cleanupRegistry: (() => void) | null = null; + let expiryTimer: ReturnType | null = null; + let cleanedUp = false; + + const cleanup = (reason: string) => { + if (cleanedUp) return; + cleanedUp = true; + if (expiryTimer) clearTimeout(expiryTimer); + cleanupRegistry?.(); + cleanupRegistry = null; + sessionStore.delete(context.terminalSession.id); + if (shell) { + shell.close(); + shell = null; + } + logger.info({ + setupSessionId: context.setupSessionId, + terminalSessionId: context.terminalSession.id, + reason, + }, "custom image terminal websocket closed"); + }; + + const closeTerminal = (reason: string, code = 1000, socketReason = "closed") => { + sendJson(socket, { type: "closed", reason }); + closeClient(socket, code, socketReason); + cleanup(reason); + }; + + cleanupRegistry = connectionRegistry.add({ + setupSessionId: context.setupSessionId, + close: (reason) => { + closeTerminal(reason); + }, + }); + + const expiresInMs = Math.max(0, context.terminalSession.expiresAt.getTime() - Date.now()); + expiryTimer = setTimeout(() => { + closeTerminal("expired", 1008, "expired"); + }, expiresInMs); + + socket.on("message", (data: unknown) => { + handleClientFrame(socket, shell, data); + }); + + socket.on("close", () => { + cleanup("client_closed"); + }); + + socket.on("error", () => { + cleanup("client_error"); + }); + + void sshConnector.connect({ + ssh: context.ssh, + term: "xterm-256color", + cols: context.initialCols, + rows: context.initialRows, + }) + .then((connectedShell) => { + if (cleanedUp) { + connectedShell.close(); + return; + } + shell = connectedShell; + shell.onData((data) => { + sendJson(socket, { type: "output", data }); + }); + shell.onClose(() => { + if (cleanedUp) return; + closeTerminal("ssh_closed"); + }); + shell.onError((err) => { + if (cleanedUp) return; + logger.warn({ + errorName: safeErrorName(err), + setupSessionId: context.setupSessionId, + terminalSessionId: context.terminalSession.id, + }, "custom image terminal ssh stream failed"); + sendJson(socket, { type: "error", message: "SSH terminal connection failed." }); + closeClient(socket, 1011, "ssh error"); + cleanup("ssh_error"); + }); + sendJson(socket, { + type: "ready", + setupSessionId: context.setupSessionId, + terminalSessionId: context.terminalSession.id, + }); + }) + .catch((err) => { + logger.warn({ + errorName: safeErrorName(err), + setupSessionId: context.setupSessionId, + terminalSessionId: context.terminalSession.id, + }, "custom image terminal ssh connection failed"); + sendJson(socket, { type: "error", message: "SSH terminal connection failed." }); + closeClient(socket, 1011, "ssh error"); + cleanup("ssh_connect_error"); + }); + }); + + wss.on("close", () => { + connectionRegistry.closeAll("server_shutdown"); + }); + + server.on("close", () => { + connectionRegistry.closeAll("server_shutdown"); + for (const client of wss.clients) { + client.terminate(); + } + wss.close(); + }); + + server.on("upgrade", (req, socket, head) => { + const reqWithContext = req as IncomingMessageWithTerminalContext; + if (!req.url) return; + + const url = new URL(req.url, "http://localhost"); + const path = parseTerminalPath(url.pathname); + if (!path) return; + + reqWithContext.paperclipWebSocketHandled = true; + + const onRawSocketError = (err: Error) => { + logger.warn({ errorName: safeErrorName(err), path: req.url }, "custom image terminal websocket upgrade socket error"); + }; + const cleanupRawSocketListeners = () => { + socket.off("error", onRawSocketError); + socket.off("close", cleanupRawSocketListeners); + }; + + socket.on("error", onRawSocketError); + socket.once("close", cleanupRawSocketListeners); + + const terminalSessionId = url.searchParams.get("terminalSessionId")?.trim() ?? ""; + const token = url.searchParams.get("token")?.trim() ?? ""; + const initialCols = parseTerminalDimension(url.searchParams.get("cols"), 80); + const initialRows = parseTerminalDimension(url.searchParams.get("rows"), 24); + if (!terminalSessionId || !token) { + rejectUpgrade(socket, "400 Bad Request", "missing terminal credentials"); + return; + } + + void validateTerminalUpgrade({ + setupSessionId: path.setupSessionId, + terminalSessionId, + token, + now: new Date(), + sessionStore, + customImages, + }) + .then((terminalSession) => { + if (!isWritableUpgradeSocket(socket)) { + cleanupRawSocketListeners(); + return; + } + + reqWithContext.paperclipTerminalUpgradeContext = { + setupSessionId: path.setupSessionId, + terminalSession, + ssh: terminalSession.ssh, + initialCols, + initialRows, + }; + + cleanupRawSocketListeners(); + wss.handleUpgrade(req, socket, head, (ws: TerminalWsSocket) => { + wss.emit("connection", ws, reqWithContext); + }); + }) + .catch((err) => { + logger.warn({ errorName: safeErrorName(err), path: req.url }, "custom image terminal websocket upgrade rejected"); + rejectUpgrade(socket, statusLineForError(err), "terminal websocket upgrade rejected"); + }); + }); + + return wss; +} diff --git a/server/src/realtime/live-events-ws.ts b/server/src/realtime/live-events-ws.ts index 07f4b05f4f3..2dbcd0f5ba9 100644 --- a/server/src/realtime/live-events-ws.ts +++ b/server/src/realtime/live-events-ws.ts @@ -47,6 +47,7 @@ interface UpgradeContext { } interface IncomingMessageWithContext extends IncomingMessage { + paperclipWebSocketHandled?: boolean; paperclipUpgradeContext?: UpgradeContext; } @@ -255,6 +256,10 @@ export function setupLiveEventsWebSocketServer( }); server.on("upgrade", (req, socket, head) => { + if ((req as IncomingMessageWithContext).paperclipWebSocketHandled) { + return; + } + const onRawSocketError = (err: Error) => { logger.warn({ err, path: req.url }, "live websocket upgrade socket error"); }; diff --git a/server/src/routes/environments.ts b/server/src/routes/environments.ts index 3847f6eacec..3f66da5c602 100644 --- a/server/src/routes/environments.ts +++ b/server/src/routes/environments.ts @@ -23,8 +23,10 @@ import { projectService, } from "../services/index.js"; import { + environmentCustomImageTerminalConnectionRegistry, environmentCustomImageTerminalSessionStore, - parseCustomImageSetupSshCommand, + validateCustomImageSetupSshPayload, + type EnvironmentCustomImageTerminalPayloadValidationResult, } from "../services/environment-custom-image-terminal-sessions.js"; import { collectEnvironmentSecretRefs, @@ -312,10 +314,13 @@ export function environmentRoutes( return expiresAt; } - function readConnectionPayload(payload: unknown): Record | null { - return payload && typeof payload === "object" && !Array.isArray(payload) - ? payload as Record - : null; + function throwTerminalPayloadValidationFailure( + failure: Extract, + ): never { + if (failure.status === 409) { + throw conflict(failure.message); + } + throw unprocessable(failure.message); } router.get("/companies/:companyId/environments", async (req, res) => { @@ -436,24 +441,9 @@ export function environmentRoutes( throw conflict(`Cannot create terminal session token from setup status "${refreshed.session.status}".`); } const setupExpiresAt = requireFutureSetupExpiry(refreshed.session, now); - const payload = readConnectionPayload(refreshed.connectionPayload); - if (!payload || payload.type !== "ssh") { - throw unprocessable("Setup session terminal connections require an SSH connection payload."); - } - const command = typeof payload.command === "string" ? payload.command.trim() : ""; - if (!command) { - throw unprocessable("Setup session SSH payload is missing a supported command."); - } - const ssh = parseCustomImageSetupSshCommand(command); - if (!ssh) { - throw unprocessable("Setup session SSH payload uses an unsupported command shape."); - } - const connectionExpiresAt = readDate(payload.expiresAt); - if (payload.expiresAt != null && !connectionExpiresAt) { - throw unprocessable("Setup session SSH payload has an invalid expiry."); - } - if (connectionExpiresAt && connectionExpiresAt.getTime() <= now.getTime()) { - throw conflict("Setup session SSH connection payload has expired."); + const payloadValidation = validateCustomImageSetupSshPayload(refreshed.connectionPayload, now); + if (!payloadValidation.ok) { + throwTerminalPayloadValidationFailure(payloadValidation); } const minted = environmentCustomImageTerminalSessionStore.create({ @@ -461,9 +451,9 @@ export function environmentRoutes( companyId: refreshed.session.companyId, environmentId: refreshed.session.environmentId, provider: refreshed.session.provider, - ssh, + ssh: payloadValidation.ssh, setupExpiresAt, - connectionExpiresAt, + connectionExpiresAt: payloadValidation.connectionExpiresAt, now, }); const actor = getActorInfo(req); @@ -481,11 +471,15 @@ export function environmentRoutes( }, }); res.status(201).json({ + id: minted.session.id, token: minted.token, expiresAt: minted.session.expiresAt.toISOString(), setupSessionId: minted.session.setupSessionId, environmentId: minted.session.environmentId, connectionType: "ssh", + websocketPath: + `/api/environment-custom-image-setup-sessions/${encodeURIComponent(minted.session.setupSessionId)}/terminal/ws` + + `?terminalSessionId=${encodeURIComponent(minted.session.id)}&token=${encodeURIComponent(minted.token)}`, }); }, ); @@ -506,6 +500,8 @@ export function environmentRoutes( sessionId: session.id, metadata: req.body.metadata, }); + environmentCustomImageTerminalSessionStore.deleteBySetupSessionId(session.id); + environmentCustomImageTerminalConnectionRegistry.closeBySetupSessionId(session.id, "setup_finished"); await logEnvironmentCustomImageActivity({ actor, companyId: session.companyId, @@ -536,6 +532,8 @@ export function environmentRoutes( sessionId: session.id, reason: req.body.reason ?? null, }); + environmentCustomImageTerminalSessionStore.deleteBySetupSessionId(session.id); + environmentCustomImageTerminalConnectionRegistry.closeBySetupSessionId(session.id, "setup_cancelled"); await logEnvironmentCustomImageActivity({ actor, companyId: session.companyId, diff --git a/server/src/services/environment-custom-image-terminal-sessions.test.ts b/server/src/services/environment-custom-image-terminal-sessions.test.ts index 91d88d058d1..23c1034dcb6 100644 --- a/server/src/services/environment-custom-image-terminal-sessions.test.ts +++ b/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"]); }); }); diff --git a/server/src/services/environment-custom-image-terminal-sessions.ts b/server/src/services/environment-custom-image-terminal-sessions.ts index ca700f4ab76..1a77a3866f1 100644 --- a/server/src/services/environment-custom-image-terminal-sessions.ts +++ b/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 | null { + return payload && typeof payload === "object" && !Array.isArray(payload) + ? payload as Record + : 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(); + private readonly sessionsById = new Map(); 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>(); + + add(input: { + setupSessionId: string; + close: EnvironmentCustomImageTerminalConnectionClose; + }): () => void { + const existing = this.connectionsBySetupSessionId.get(input.setupSessionId); + const connections = existing ?? new Set(); + 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(); diff --git a/server/src/services/index.ts b/server/src/services/index.ts index da185cd24c0..4300bf75eda 100644 --- a/server/src/services/index.ts +++ b/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, diff --git a/ui/src/pages/CompanyEnvironments.test.tsx b/ui/src/pages/CompanyEnvironments.test.tsx index 7305d6d21dc..7ed7f172450 100644 --- a/ui/src/pages/CompanyEnvironments.test.tsx +++ b/ui/src/pages/CompanyEnvironments.test.tsx @@ -172,6 +172,27 @@ function createTemplate(overrides: Record = {}) { }; } +function supportedDaytonaCapabilities() { + return { + adapters: [], + drivers: { local: "supported", ssh: "supported", sandbox: "supported", plugin: "unsupported" }, + sandboxProviders: { + daytona: { + status: "supported", + supportsSavedProbe: true, + supportsUnsavedProbe: true, + supportsRunExecution: true, + supportsReusableLeases: true, + supportsInteractiveSetup: true, + interactiveSetupConnectionTypes: ["ssh"], + supportsTemplateCapture: true, + supportsTemplateDelete: true, + displayName: "Daytona", + }, + }, + }; +} + describe("CompanyEnvironments — test provider button", () => { let container: HTMLDivElement; let root: ReturnType | null; @@ -601,6 +622,84 @@ describe("CompanyEnvironments — test provider button", () => { expect(getOpenDialog()?.textContent).not.toContain(command); }); + it("shows a setup connection refresh fallback without breaking finish or cancel", async () => { + root = createRoot(container); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + mockEnvironmentsApi.list.mockResolvedValue([ + { id: "env-1", name: "Daytona", driver: "sandbox", description: null, config: { provider: "daytona" } }, + ]); + mockEnvironmentsApi.capabilities.mockResolvedValue(supportedDaytonaCapabilities()); + mockEnvironmentsApi.customImageTemplate.mockResolvedValue({ + activeTemplate: null, + activeSession: createSession(), + latestSession: createSession(), + }); + mockEnvironmentsApi.customImageSetupSession.mockRejectedValue(new Error("proxy unavailable")); + + await act(async () => { + root!.render( + + + + + , + ); + }); + await flushReact(); + + await act(async () => click(editButtons(container)[0])); + await waitForAssertion(() => { + expect(getOpenDialog()?.textContent).toContain("Setup connection details could not be refreshed."); + }); + + const dialog = getOpenDialog()!; + expect(findButton(dialog, "Finished")?.disabled).toBe(false); + expect(findButton(dialog, "Cancel")?.disabled).toBe(false); + + await act(async () => click(findButton(dialog, "Finished"))); + await flushReact(); + + expect(mockEnvironmentsApi.finishCustomImageSetupSession).toHaveBeenCalledExactlyOnceWith("session-1", {}); + }); + + it("shows provider fallback messaging for unsupported setup connection payloads", async () => { + root = createRoot(container); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + mockEnvironmentsApi.list.mockResolvedValue([ + { id: "env-1", name: "Daytona", driver: "sandbox", description: null, config: { provider: "daytona" } }, + ]); + mockEnvironmentsApi.capabilities.mockResolvedValue(supportedDaytonaCapabilities()); + mockEnvironmentsApi.customImageTemplate.mockResolvedValue({ + activeTemplate: null, + activeSession: createSession(), + latestSession: createSession(), + }); + mockEnvironmentsApi.customImageSetupSession.mockResolvedValue({ + session: createSession(), + connectionPayload: { type: "browser_terminal", command: null }, + }); + + await act(async () => { + root!.render( + + + + + , + ); + }); + await flushReact(); + + await act(async () => click(editButtons(container)[0])); + await waitForAssertion(() => { + expect(getOpenDialog()?.textContent).toContain("Browser terminal is not available for this provider connection."); + }); + + const dialog = getOpenDialog()!; + expect(findButton(dialog, "Finished")?.disabled).toBe(false); + expect(findButton(dialog, "Cancel")?.disabled).toBe(false); + }); + it("shows active template controls for refresh, rollback, and disable", async () => { root = createRoot(container); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); diff --git a/ui/src/pages/CompanyEnvironments.tsx b/ui/src/pages/CompanyEnvironments.tsx index 2f475fddb38..cc9e915b773 100644 --- a/ui/src/pages/CompanyEnvironments.tsx +++ b/ui/src/pages/CompanyEnvironments.tsx @@ -199,6 +199,27 @@ function readConnectionCommand(payload: EnvironmentCustomImageConnectionPayload : null; } +function setupConnectionFallbackMessage(input: { + payload: EnvironmentCustomImageConnectionPayload | null; + refreshError: unknown; + isLoading: boolean; +}): string | null { + if (input.refreshError) { + return "Setup connection details could not be refreshed. You can still finish or cancel this setup."; + } + if (input.isLoading) return null; + if (!input.payload) { + return "Connection details are not available yet. You can still finish or cancel this setup."; + } + if (input.payload.type !== "ssh") { + return "Browser terminal is not available for this provider connection. Use the provider setup instructions, then finish or cancel here."; + } + if (!readConnectionCommand(input.payload)) { + return "Connection details are not available yet. You can still finish or cancel this setup."; + } + return null; +} + function capabilityState(capability: EnvironmentProviderCapability | null | undefined) { if (!capability || capability.status !== "supported" || !capability.supportsInteractiveSetup) { return { @@ -451,6 +472,13 @@ function EnvironmentImageTemplatePanel({ ? sessionQuery.data?.connectionPayload ?? null : null; const connectionCommand = readConnectionCommand(connectionPayload); + const connectionFallbackMessage = session?.status === "waiting_for_user" + ? setupConnectionFallbackMessage({ + payload: connectionPayload, + refreshError: sessionQuery.isError ? sessionQuery.error : null, + isLoading: sessionQuery.isLoading, + }) + : null; const sessionExpiresAt = formatDateTime(connectionPayload?.expiresAt ?? session?.expiresAt ?? null); const capturedAt = formatDateTime(activeTemplate?.capturedAt ?? activeTemplate?.createdAt ?? null); const lastUsedAt = formatDateTime(activeTemplate?.lastUsedAt ?? null); @@ -500,9 +528,9 @@ function EnvironmentImageTemplatePanel({ ) : null} - {session.status === "waiting_for_user" && !connectionCommand ? ( + {session.status === "waiting_for_user" && connectionFallbackMessage ? (
- Connection details are not available yet. + {connectionFallbackMessage}
) : null} {session.failureReason ? (