diff --git a/apps/api/package.json b/apps/api/package.json index 715dc23af5..cfe3165730 100644 --- a/apps/api/package.json +++ b/apps/api/package.json @@ -106,6 +106,7 @@ "ts-results": "^3.3.0", "tsyringe": "^4.10.0", "unleash-client": "^6.6.0", + "ws": "^8.18.2", "zod": "3.*" }, "devDependencies": { @@ -126,6 +127,7 @@ "@types/pg": "^8.11.6", "@types/semver": "^7.5.2", "@types/tar": "^6.1.13", + "@types/ws": "^8.5.4", "@typescript-eslint/eslint-plugin": "^7.12.0", "@vitest/coverage-v8": "^4.1.5", "drizzle-kit": "^0.22.7", diff --git a/apps/api/src/deployment/controllers/shell-exec/shell-exec.controller.spec.ts b/apps/api/src/deployment/controllers/shell-exec/shell-exec.controller.spec.ts new file mode 100644 index 0000000000..5e908a3fb7 --- /dev/null +++ b/apps/api/src/deployment/controllers/shell-exec/shell-exec.controller.spec.ts @@ -0,0 +1,310 @@ +import { Err, Ok } from "ts-results"; +import { container } from "tsyringe"; +import { describe, expect, it } from "vitest"; +import { mock } from "vitest-mock-extended"; + +import { AuthService } from "@src/auth/services/auth.service"; +import type { WalletReaderService } from "@src/billing/services/wallet-reader/wallet-reader.service"; +import type { DeploymentReaderService } from "@src/deployment/services/deployment-reader/deployment-reader.service"; +import type { ShellExecService } from "@src/deployment/services/shell-exec/shell-exec.service"; +import type { ProviderService } from "@src/provider/services/provider/provider.service"; +import { mapExecError, ShellExecController } from "./shell-exec.controller"; + +import { createUser } from "@test/seeders/user.seeder"; + +type DeploymentResponse = NonNullable>>; +type Lease = DeploymentResponse["leases"][number]; +type ProviderInfo = NonNullable>>; + +function createLease(overrides: Partial<{ gseq: number; oseq: number; provider: string; state: string }> = {}): Lease { + return mock({ + id: { owner: "akash1owner", dseq: "1234", gseq: overrides.gseq ?? 1, oseq: overrides.oseq ?? 1, provider: overrides.provider ?? "akash1provider", bseq: 0 }, + state: overrides.state ?? "active", + price: { denom: "uakt", amount: "100" }, + created_at: "12345", + closed_on: "0", + status: null + }); +} + +function createDeployment(overrides: Partial<{ state: string; leases: Lease[] }> = {}): DeploymentResponse { + return mock({ + deployment: { + id: { owner: "akash1owner", dseq: "1234" }, + state: overrides.state ?? "active", + hash: "abc123", + created_at: "12345" + }, + leases: overrides.leases ?? [createLease()], + escrow_account: { + id: { scope: "deployment", xid: "1234" }, + state: { + owner: "akash1owner", + state: "open", + transferred: [], + settled_at: "12345", + funds: [{ denom: "uakt", amount: "1000" }], + deposits: [] + } + } + }); +} + +function createProviderInfo(overrides: Partial<{ hostUri: string }> = {}): ProviderInfo { + return mock({ hostUri: overrides.hostUri ?? "https://provider.example.com" }); +} + +describe("mapExecError", () => { + it.each([ + ["Command timed out", 504, "Command execution timed out"], + ["Auth expired: provider closed connection (code 4001)", 403, "Provider authentication expired"], + ["Invalid provider host: http://provider.example.com", 502, "Invalid provider host"], + ["WebSocket connection failed: ECONNREFUSED", 502, "Failed to connect to provider"], + ["Provider error: internal error", 502, "Provider returned an error"], + ["Connection closed without exit code", 502, "Provider connection closed unexpectedly"], + ["some entirely unexpected failure", 502, "Shell execution failed"] + ])("maps %j to status %i", (errVal, status, message) => { + expect(mapExecError(errVal as string)).toEqual({ status, message }); + }); +}); + +describe(ShellExecController.name, () => { + it("throws 404 when deployment not found", async () => { + const { controller, deploymentReaderService } = setup(); + // Simulate "deployment not found": clear the default stub so the lookup resolves undefined. + deploymentReaderService.findByUserIdAndDseq.mockReset(); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(404); + expect(error.message).toBe("Deployment not found"); + }); + + it("throws 404 when no lease matches the provided gseq and oseq", async () => { + const { controller } = setup(); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 99, oseq: 99, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(404); + expect(error.message).toBe("Lease not found"); + }); + + it("throws 500 when lease provider address is an empty string", async () => { + const { controller } = setup({ provider: "" }); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(500); + expect(error.message).toBe("Lease provider address not found"); + }); + + it("throws 400 when lease state is not active", async () => { + const { controller } = setup({ state: "closed" }); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(400); + expect(error.message).toBe("Lease is not active"); + }); + + it("throws 504 when the command times out", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("Command timed out")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(504); + expect(error.message).toBe("Command execution timed out"); + }); + + it("returns the shell exec result on successful execution", async () => { + const { controller, shellExecService } = setup(); + + const result = await controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 }); + + expect(result).toEqual({ stdout: "output", stderr: "", exitCode: 0, truncated: false }); + expect(shellExecService.execute).toHaveBeenCalledWith({ + providerBaseUrl: "https://provider.example.com", + providerAddress: "akash1provider", + dseq: "1234", + gseq: 1, + oseq: 1, + service: "web", + command: ["ls"], + timeout: 60, + jwtToken: "test-token" + }); + }); + + it("forwards stdin through to the shell exec service", async () => { + const { controller, shellExecService } = setup(); + const secretValue = Math.random().toString(36).substring(2); + + await controller.exec({ + dseq: "1234", + gseq: 1, + oseq: 1, + command: ["sh", "-c", "cat > /run/secrets/.env"], + service: "web", + timeout: 60, + stdin: `SECRET=${secretValue}` + }); + + expect(shellExecService.execute).toHaveBeenCalledWith(expect.objectContaining({ stdin: `SECRET=${secretValue}` })); + }); + + it("throws 404 when provider info lookup returns null", async () => { + const { controller, providerService } = setup(); + providerService.getProvider.mockResolvedValue(null); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(404); + expect(error.message).toBe("Provider not found"); + }); + + it("does not mint a provider JWT when the provider lookup fails (getProvider precedes toProviderAuth)", async () => { + const { controller, providerService } = setup(); + providerService.getProvider.mockResolvedValue(null); + + await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(providerService.toProviderAuth).not.toHaveBeenCalled(); + }); + + it("throws 404 when deployment has an empty leases array", async () => { + const { controller, deploymentReaderService } = setup(); + deploymentReaderService.findByUserIdAndDseq.mockResolvedValue(createDeployment({ leases: [] })); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(404); + expect(error.message).toBe("Lease not found"); + }); + + it("finds the correct lease among multiple leases by gseq and oseq", async () => { + const { controller, deploymentReaderService, shellExecService } = setup(); + + deploymentReaderService.findByUserIdAndDseq.mockResolvedValue( + createDeployment({ + leases: [ + createLease({ gseq: 1, oseq: 1 }), + createLease({ gseq: 2, oseq: 1, provider: "akash1provider2" }), + createLease({ gseq: 1, oseq: 2, provider: "akash1provider3" }) + ] + }) + ); + + const result = await controller.exec({ dseq: "1234", gseq: 2, oseq: 1, command: ["ls"], service: "web", timeout: 60 }); + + expect(result).toEqual({ stdout: "output", stderr: "", exitCode: 0, truncated: false }); + expect(shellExecService.execute).toHaveBeenCalledWith(expect.objectContaining({ providerAddress: "akash1provider2" })); + }); + + it("throws 400 when deployment state is closed", async () => { + const { controller, deploymentReaderService } = setup(); + deploymentReaderService.findByUserIdAndDseq.mockResolvedValue(createDeployment({ state: "closed" })); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(400); + expect(error.message).toBe("Deployment is not active"); + }); + + it("throws 502 with stable message when WS connection fails", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("WebSocket connection failed: ECONNREFUSED")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(502); + expect(error.message).toBe("Failed to connect to provider"); + }); + + it("throws 502 with stable message when provider returns error", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("Provider error: internal error")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(502); + expect(error.message).toBe("Provider returned an error"); + }); + + it("throws 403 when the provider JWT expires mid-run", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("Auth expired: provider closed connection (code 4001)")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(403); + expect(error.message).toBe("Provider authentication expired"); + }); + + it("throws 502 for an invalid (server-derived) provider host", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("Invalid provider host: http://provider.example.com")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(502); + expect(error.message).toBe("Invalid provider host"); + }); + + it("throws 502 with a distinct message when the connection closes without an exit code", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("Connection closed without exit code")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(502); + expect(error.message).toBe("Provider connection closed unexpectedly"); + }); + + it("falls back to a generic 502 for an unrecognized service error", async () => { + const { controller, shellExecService } = setup(); + shellExecService.execute.mockResolvedValue(Err("some entirely unexpected failure")); + + const error = await captureError(() => controller.exec({ dseq: "1234", gseq: 1, oseq: 1, command: ["ls"], service: "web", timeout: 60 })); + + expect(error.status).toBe(502); + expect(error.message).toBe("Shell execution failed"); + }); + + async function captureError(fn: () => Promise): Promise<{ status: number; message: string }> { + try { + await fn(); + throw new Error("Expected function to throw"); + } catch (error) { + return error as { status: number; message: string }; + } + } + + function setup(overrides?: { provider?: string; state?: string }) { + const user = createUser(); + const deploymentReaderService = mock(); + const providerService = mock(); + const shellExecService = mock(); + const authService = mock({ currentUser: user }); + const walletReaderService = mock(); + + container.register(AuthService, { useValue: authService }); + + const controller = new ShellExecController(deploymentReaderService, providerService, shellExecService, authService, walletReaderService); + + const provider = overrides?.provider ?? "akash1provider"; + const state = overrides?.state ?? "active"; + + const deployment = createDeployment({ leases: [createLease({ provider, state })] }); + + deploymentReaderService.findByUserIdAndDseq.mockResolvedValue(deployment); + walletReaderService.getWalletByUserId.mockResolvedValue( + mock>>({ id: 1, address: "akash1wallet" }) + ); + providerService.toProviderAuth.mockResolvedValue({ type: "jwt" as const, token: "test-token" }); + providerService.getProvider.mockResolvedValue(createProviderInfo()); + shellExecService.execute.mockResolvedValue(new Ok({ stdout: "output", stderr: "", exitCode: 0, truncated: false })); + + return { controller, deploymentReaderService, providerService, shellExecService, authService, walletReaderService, user, deployment }; + } +}); diff --git a/apps/api/src/deployment/controllers/shell-exec/shell-exec.controller.ts b/apps/api/src/deployment/controllers/shell-exec/shell-exec.controller.ts new file mode 100644 index 0000000000..ce9f39d156 --- /dev/null +++ b/apps/api/src/deployment/controllers/shell-exec/shell-exec.controller.ts @@ -0,0 +1,95 @@ +import assert from "http-assert"; +import { singleton } from "tsyringe"; + +import { AuthService, Protected } from "@src/auth/services/auth.service"; +import { WalletReaderService } from "@src/billing/services/wallet-reader/wallet-reader.service"; +import { ShellExecRequest, ShellExecResponse } from "@src/deployment/http-schemas/shell-exec.schema"; +import { DeploymentReaderService } from "@src/deployment/services/deployment-reader/deployment-reader.service"; +import { ShellExecService } from "@src/deployment/services/shell-exec/shell-exec.service"; +import { ProviderService } from "@src/provider/services/provider/provider.service"; + +type ExecErrorMapping = { status: number; message: string }; + +/** + * Ordered mapping from a service-layer error sentinel (the stable, prefixed + * strings returned by `ShellExecService.execute`) to a user-facing HTTP status + * and message. Keeping this a table — rather than a chain of ternaries — makes + * the contract readable, testable, and cheap to extend, and guarantees the raw + * internal string is never leaked to the client. Each failure mode gets the + * semantically correct status: + * - timeout → 504 (we are a gateway to the provider, not a slow client) + * - auth expired → 403 + * - invalid host → 502 (`hostUri` is server-derived on-chain data, not + * client input, so a bad host is an upstream fault) + * - connection/provider→ 502 + */ +const EXEC_ERROR_TABLE: ReadonlyArray<{ prefix: string } & ExecErrorMapping> = [ + { prefix: "Command timed out", status: 504, message: "Command execution timed out" }, + { prefix: "Auth expired", status: 403, message: "Provider authentication expired" }, + { prefix: "Invalid provider host", status: 502, message: "Invalid provider host" }, + { prefix: "WebSocket connection failed", status: 502, message: "Failed to connect to provider" }, + { prefix: "Provider error", status: 502, message: "Provider returned an error" }, + { prefix: "Connection closed without exit code", status: 502, message: "Provider connection closed unexpectedly" } +]; + +export function mapExecError(errVal: string): ExecErrorMapping { + const match = EXEC_ERROR_TABLE.find(entry => errVal.startsWith(entry.prefix)); + return match ? { status: match.status, message: match.message } : { status: 502, message: "Shell execution failed" }; +} + +@singleton() +export class ShellExecController { + constructor( + private readonly deploymentReaderService: DeploymentReaderService, + private readonly providerService: ProviderService, + private readonly shellExecService: ShellExecService, + private readonly authService: AuthService, + private readonly walletReaderService: WalletReaderService + ) {} + + @Protected([{ action: "read", subject: "Lease" }]) + async exec(input: ShellExecRequest & { dseq: string; gseq: number; oseq: number }): Promise { + const userId = this.authService.currentUser.id; + + const deployment = await this.deploymentReaderService.findByUserIdAndDseq(userId, input.dseq); + + assert(deployment, 404, "Deployment not found"); + assert(deployment.deployment.state === "active", 400, "Deployment is not active"); + + const lease = deployment.leases.find(l => l.id.gseq === input.gseq && l.id.oseq === input.oseq); + + assert(lease, 404, "Lease not found"); + assert(lease.id.provider, 500, "Lease provider address not found"); + assert(lease.state === "active", 400, "Lease is not active"); + + const providerAddress = lease.id.provider; + + const providerInfo = await this.providerService.getProvider(providerAddress); + + assert(providerInfo, 404, "Provider not found"); + + const wallet = await this.walletReaderService.getWalletByUserId(userId); + + const auth = await this.providerService.toProviderAuth({ walletId: wallet.id, provider: providerAddress }, ["shell"]); + + const result = await this.shellExecService.execute({ + providerBaseUrl: providerInfo.hostUri, + providerAddress: providerAddress, + dseq: input.dseq, + gseq: input.gseq, + oseq: input.oseq, + service: input.service, + command: input.command, + stdin: input.stdin, + timeout: input.timeout, + jwtToken: auth.token + }); + + if (!result.ok) { + const { status, message } = mapExecError(result.val); + assert(false, status, message); + } + + return result.val; + } +} diff --git a/apps/api/src/deployment/http-schemas/shell-exec.schema.spec.ts b/apps/api/src/deployment/http-schemas/shell-exec.schema.spec.ts new file mode 100644 index 0000000000..0bfce980dd --- /dev/null +++ b/apps/api/src/deployment/http-schemas/shell-exec.schema.spec.ts @@ -0,0 +1,79 @@ +import { describe, expect, it } from "vitest"; + +import { + JSON_ENVELOPE_OVERHEAD_BYTES, + MAX_COMMAND_ARG_BYTES, + MAX_COMMAND_ARGS, + MAX_STDIN_BYTES, + SERVICE_NAME_MAX, + SHELL_EXEC_BODY_LIMIT_BYTES, + ShellExecRequestSchema +} from "./shell-exec.schema"; + +describe("ShellExecRequestSchema", () => { + describe("stdin byte cap", () => { + it("accepts stdin at exactly the byte cap", () => { + const result = ShellExecRequestSchema.safeParse(base({ stdin: "A".repeat(MAX_STDIN_BYTES) })); + expect(result.success).toBe(true); + }); + + it("rejects stdin one byte over the cap", () => { + const result = ShellExecRequestSchema.safeParse(base({ stdin: "A".repeat(MAX_STDIN_BYTES + 1) })); + expect(result.success).toBe(false); + }); + + it("counts UTF-8 bytes, not UTF-16 code units, so a short multi-byte string over the byte cap is rejected", () => { + // "😀" is 2 UTF-16 code units but 4 UTF-8 bytes. This many are under MAX_STDIN_BYTES + // by `.length` (code units) yet over it in bytes — a `z.string().max()` would wrongly pass. + const emojiCount = Math.floor(MAX_STDIN_BYTES / 2); + const value = "😀".repeat(emojiCount); + expect(value.length).toBeLessThanOrEqual(MAX_STDIN_BYTES); + expect(Buffer.byteLength(value, "utf8")).toBeGreaterThan(MAX_STDIN_BYTES); + + const result = ShellExecRequestSchema.safeParse(base({ stdin: value })); + expect(result.success).toBe(false); + }); + }); + + describe("command caps", () => { + it("accepts a command arg at exactly the per-arg byte cap", () => { + const result = ShellExecRequestSchema.safeParse(base({ command: ["sh", "-c", "A".repeat(MAX_COMMAND_ARG_BYTES)] })); + expect(result.success).toBe(true); + }); + + it("rejects a command arg one byte over the per-arg cap", () => { + const result = ShellExecRequestSchema.safeParse(base({ command: ["sh", "-c", "A".repeat(MAX_COMMAND_ARG_BYTES + 1)] })); + expect(result.success).toBe(false); + }); + + it("rejects more than the maximum number of command args", () => { + const result = ShellExecRequestSchema.safeParse(base({ command: Array.from({ length: MAX_COMMAND_ARGS + 1 }, () => "x") })); + expect(result.success).toBe(false); + }); + + it("rejects an empty command array", () => { + const result = ShellExecRequestSchema.safeParse(base({ command: [] })); + expect(result.success).toBe(false); + }); + }); + + describe("SHELL_EXEC_BODY_LIMIT_BYTES", () => { + it("is derived arithmetically from the field caps (never hardcoded)", () => { + expect(SHELL_EXEC_BODY_LIMIT_BYTES).toBe(MAX_COMMAND_ARGS * MAX_COMMAND_ARG_BYTES + MAX_STDIN_BYTES + SERVICE_NAME_MAX + JSON_ENVELOPE_OVERHEAD_BYTES); + }); + + it("leaves headroom above the largest single field cap so a schema-valid body is not rejected first", () => { + expect(SHELL_EXEC_BODY_LIMIT_BYTES).toBeGreaterThan(MAX_STDIN_BYTES); + expect(SHELL_EXEC_BODY_LIMIT_BYTES).toBeGreaterThan(MAX_COMMAND_ARGS * MAX_COMMAND_ARG_BYTES); + }); + }); +}); + +function base(overrides: Partial<{ command: string[]; service: string; timeout: number; stdin: string }>) { + return { + command: overrides.command ?? ["ls"], + service: overrides.service ?? "web", + timeout: overrides.timeout ?? 60, + ...(overrides.stdin !== undefined ? { stdin: overrides.stdin } : {}) + }; +} diff --git a/apps/api/src/deployment/http-schemas/shell-exec.schema.ts b/apps/api/src/deployment/http-schemas/shell-exec.schema.ts new file mode 100644 index 0000000000..ee400157bb --- /dev/null +++ b/apps/api/src/deployment/http-schemas/shell-exec.schema.ts @@ -0,0 +1,75 @@ +import { z } from "@hono/zod-openapi"; + +/** + * Request-side size limits for the shell-exec endpoint, exported as named + * constants so the schema, the route body limit, and the tests all derive from + * a single source of truth and cannot drift (per maintainer review on #3097). + * + * Real secret-injection payloads (env files, API keys, DB passwords, TLS + * cert/key chains, service-mesh tokens) are kilobytes — these caps are generous + * for the intended use while keeping per-request memory bounded so concurrent + * execs cannot be used as a cheap DoS. + */ +export const MAX_STDIN_BYTES = 16 * 1024; +export const MAX_COMMAND_ARGS = 64; +export const MAX_COMMAND_ARG_BYTES = 1024; +export const SERVICE_NAME_MAX = 253; + +/** + * Headroom for JSON key names, quotes, and escaping. The route body limit is + * measured on RAW request bytes (before JSON parsing), whereas the per-field + * byte caps above apply to the decoded strings (after parsing) — related but + * not identical. This overhead ensures a schema-valid request is never rejected + * by the body limit first. + */ +export const JSON_ENVELOPE_OVERHEAD_BYTES = 2 * 1024; + +/** + * Explicit request-body limit for the route, DERIVED arithmetically from the + * field caps above (never hardcoded) so it stays in sync automatically if any + * individual limit changes. + */ +export const SHELL_EXEC_BODY_LIMIT_BYTES = MAX_COMMAND_ARGS * MAX_COMMAND_ARG_BYTES + MAX_STDIN_BYTES + SERVICE_NAME_MAX + JSON_ENVELOPE_OVERHEAD_BYTES; + +/** UTF-8 byte length — the true wire size, unlike `.length`/`z.string().max()` which count UTF-16 code units. */ +const utf8Bytes = (value: string): number => Buffer.byteLength(value, "utf8"); + +export const ShellExecParamsSchema = z.object({ + dseq: z.string().regex(/^\d+$/), + gseq: z.coerce.number().int().nonnegative(), + oseq: z.coerce.number().int().nonnegative() +}); + +export const ShellExecRequestSchema = z.object({ + command: z + .array( + z + .string() + .min(1) + .refine(s => utf8Bytes(s) <= MAX_COMMAND_ARG_BYTES, { message: `each command argument must not exceed ${MAX_COMMAND_ARG_BYTES} bytes (UTF-8)` }) + ) + .min(1) + .max(MAX_COMMAND_ARGS), + service: z.string().min(1).max(SERVICE_NAME_MAX), + timeout: z.number().int().min(1).max(120).default(60), + stdin: z + .string() + .refine(s => utf8Bytes(s) <= MAX_STDIN_BYTES, { message: `stdin must not exceed ${MAX_STDIN_BYTES} bytes (UTF-8)` }) + .optional() + .openapi({ + description: + 'Optional raw UTF-8 data streamed to the command\'s standard input (max 16 KiB). Put secrets HERE (env files, tokens, passwords) — never in `command`, whose tokens are placed in the provider-proxy request URL, which is logged. Example command: ["sh","-c","cat > /run/secrets/.env"].', + example: "SECRET=value" + }) +}); + +export const ShellExecResponseSchema = z.object({ + stdout: z.string(), + stderr: z.string(), + exitCode: z.number(), + truncated: z.boolean() +}); + +export type ShellExecParams = z.infer; +export type ShellExecRequest = z.infer; +export type ShellExecResponse = z.infer; diff --git a/apps/api/src/deployment/routes/shell-exec/shell-exec.router.ts b/apps/api/src/deployment/routes/shell-exec/shell-exec.router.ts new file mode 100644 index 0000000000..fa56a72648 --- /dev/null +++ b/apps/api/src/deployment/routes/shell-exec/shell-exec.router.ts @@ -0,0 +1,76 @@ +import { container } from "tsyringe"; + +import { createRoute } from "@src/core/lib/create-route/create-route"; +import { OpenApiHonoHandler } from "@src/core/services/open-api-hono-handler/open-api-hono-handler"; +import { SECURITY_BEARER_OR_API_KEY } from "@src/core/services/openapi-docs/openapi-security"; +import { ShellExecController } from "@src/deployment/controllers/shell-exec/shell-exec.controller"; +import { + SHELL_EXEC_BODY_LIMIT_BYTES, + ShellExecParamsSchema, + ShellExecRequestSchema, + ShellExecResponseSchema +} from "@src/deployment/http-schemas/shell-exec.schema"; + +export const shellExecRouter = new OpenApiHonoHandler(); + +const shellExecRoute = createRoute({ + method: "post", + path: "/v1/deployments/{dseq}/leases/{gseq}/{oseq}/shell-exec", + summary: "Execute a shell command in a deployment container", + tags: ["Shell Exec"], + security: SECURITY_BEARER_OR_API_KEY, + // Explicit body limit derived from the schema field caps (see shell-exec.schema.ts) + // so it can never drift from them. Measured on raw request bytes (pre-parse). + bodyLimit: { maxSize: SHELL_EXEC_BODY_LIMIT_BYTES }, + request: { + params: ShellExecParamsSchema, + body: { + content: { + "application/json": { + schema: ShellExecRequestSchema + } + } + } + }, + responses: { + 200: { + description: "Command executed successfully", + content: { + "application/json": { + schema: ShellExecResponseSchema + } + } + }, + 400: { + description: "Invalid request (e.g., lease not active)" + }, + 401: { + description: "Unauthorized" + }, + 403: { + description: "Forbidden - user does not own this deployment, or the provider authentication expired mid-execution" + }, + 404: { + description: "Deployment or lease not found" + }, + 413: { + description: "Request body exceeds the shell-exec body limit" + }, + 500: { + description: "Internal server error (e.g., lease provider address missing)" + }, + 502: { + description: "Provider proxy error (invalid provider host, connection failure, or provider-reported error)" + }, + 504: { + description: "Command execution timed out" + } + } +}); + +shellExecRouter.openapi(shellExecRoute, async function routeShellExec(c) { + const params = c.req.valid("param"); + const body = c.req.valid("json"); + const result = await container.resolve(ShellExecController).exec({ ...params, ...body }); + return c.json(result, 200); +}); diff --git a/apps/api/src/deployment/services/shell-exec/shell-exec.service.spec.ts b/apps/api/src/deployment/services/shell-exec/shell-exec.service.spec.ts new file mode 100644 index 0000000000..f3228ea14b --- /dev/null +++ b/apps/api/src/deployment/services/shell-exec/shell-exec.service.spec.ts @@ -0,0 +1,632 @@ +import { faker } from "@faker-js/faker"; +import { describe, expect, it, vi } from "vitest"; +import { mock } from "vitest-mock-extended"; +import { WebSocket } from "ws"; + +import type { LoggerService } from "@src/core"; +import type { DeploymentConfig } from "@src/deployment/config/config.provider"; +import { + buildShellUrl, + isStrictBase64, + isValidProviderHost, + MAX_OUTPUT_SIZE, + parseExitCode, + type ShellExecInput, + type ShellExecOutput, + ShellExecService, + toProxyWebSocketUrl +} from "./shell-exec.service"; + +vi.mock("ws", () => ({ + WebSocket: vi.fn() +})); + +const PROXY_URL = "https://proxy.example.com"; + +describe(ShellExecService.name, () => { + describe("toProxyWebSocketUrl", () => { + it("maps https to wss", () => { + expect(toProxyWebSocketUrl("https://proxy.example.com")).toBe("wss://proxy.example.com"); + }); + + it("maps http to ws", () => { + expect(toProxyWebSocketUrl("http://localhost:3000")).toBe("ws://localhost:3000"); + }); + }); + + describe("isValidProviderHost", () => { + it("accepts an https host with a domain name", () => { + expect(isValidProviderHost("https://provider.example.com:8443")).toBe(true); + }); + + it("rejects a non-https (http) host", () => { + expect(isValidProviderHost("http://provider.example.com")).toBe(false); + }); + + it("rejects an IPv4 host", () => { + expect(isValidProviderHost("https://203.0.113.10:8443")).toBe(false); + }); + + it("rejects an IPv6 host", () => { + expect(isValidProviderHost("https://[2001:db8::1]:8443")).toBe(false); + }); + + it("rejects a .local host", () => { + expect(isValidProviderHost("https://provider.local:8443")).toBe(false); + }); + + it("rejects a malformed URL", () => { + expect(isValidProviderHost("not a url")).toBe(false); + }); + }); + + describe("isStrictBase64", () => { + it("accepts a valid base64 string", () => { + expect(isStrictBase64(Buffer.from("hello").toString("base64"))).toBe(true); + }); + + it("rejects prose with spaces", () => { + expect(isStrictBase64("Received error from provider websocket")).toBe(false); + }); + + it("rejects a string whose length is not a multiple of four", () => { + expect(isStrictBase64("abc")).toBe(false); + }); + + it("rejects an empty string", () => { + expect(isStrictBase64("")).toBe(false); + }); + }); + + describe("buildShellUrl (argv mapping)", () => { + it("maps a single-token argv to cmd0", () => { + const url = buildShellUrl(createShellExecInput({ command: ["ls"] })); + + expect(url).toContain("/lease/1234/1/1/shell"); + expect(url).toContain("stdin=0"); + expect(url).toContain("tty=0"); + expect(url).toContain("podIndex=0"); + expect(url).toContain("service=test-service"); + expect(url).toContain("cmd0=ls"); + expect(url).not.toContain("cmd1="); + }); + + it("maps a multi-token argv to cmd0..cmdN", () => { + const url = buildShellUrl(createShellExecInput({ command: ["echo", "hello", "world"] })); + + expect(url).toContain("cmd0=echo"); + expect(url).toContain("cmd1=hello"); + expect(url).toContain("cmd2=world"); + expect(url).not.toContain("cmd3="); + }); + + it("URL-encodes each argv token independently, preserving whitespace inside a token", () => { + const url = buildShellUrl(createShellExecInput({ command: ["sh", "-c", "echo SECRET=v > /run/secrets/.env"] })); + + expect(url).toContain("cmd0=sh"); + expect(url).toContain("cmd1=-c"); + expect(url).toContain("cmd2=echo%20SECRET%3Dv%20%3E%20%2Frun%2Fsecrets%2F.env"); + expect(url).not.toContain("cmd3="); + }); + + it("encodes reserved characters in a token", () => { + const url = buildShellUrl(createShellExecInput({ command: ["echo", "a&b?c#d"] })); + + expect(url).toContain("cmd0=echo"); + expect(url).toContain("cmd1=a%26b%3Fc%23d"); + }); + + it("encodes the service name", () => { + const url = buildShellUrl(createShellExecInput({ service: "my service" })); + + expect(url).toContain("service=my%20service"); + }); + + it("uses the provided gseq and oseq in the path", () => { + const url = buildShellUrl(createShellExecInput({ gseq: 3, oseq: 5 })); + + expect(url).toContain("/lease/1234/3/5/shell"); + }); + + it("removes a trailing slash from the provider base URL", () => { + const url = buildShellUrl(createShellExecInput({ providerBaseUrl: "https://provider.example.com/" })); + + expect(url).toMatch(/^https:\/\/provider\.example\.com\/lease/); + }); + + it("URL-encodes the dseq path segment", () => { + const url = buildShellUrl(createShellExecInput({ dseq: "foo/bar" })); + + expect(url).toContain("/lease/foo%2Fbar/1/1/shell"); + }); + }); + + describe("buildShellUrl (stdin flag)", () => { + it("emits stdin=0 when no stdin is provided", () => { + const url = buildShellUrl(createShellExecInput()); + + expect(url).toContain("stdin=0"); + expect(url).not.toContain("stdin=1"); + }); + + it("emits stdin=1 when stdin is provided", () => { + const url = buildShellUrl(createShellExecInput({ stdin: `SECRET=${faker.string.alphanumeric(16)}` })); + + expect(url).toContain("stdin=1"); + expect(url).not.toContain("stdin=0"); + }); + + it("emits stdin=0 for an empty stdin string", () => { + const url = buildShellUrl(createShellExecInput({ stdin: "" })); + + expect(url).toContain("stdin=0"); + expect(url).not.toContain("stdin=1"); + }); + + it("never places the stdin payload in the URL", () => { + const secret = `SUPER_SECRET_VALUE_${faker.string.alphanumeric(16)}`; + const url = buildShellUrl(createShellExecInput({ command: ["sh", "-c", "cat > /run/secrets/.env"], stdin: secret })); + + expect(url).not.toContain(secret); + expect(url).not.toContain(encodeURIComponent(secret)); + }); + }); + + describe("parseExitCode", () => { + it("parses a JSON exit_code body", () => { + expect(parseExitCode(Buffer.from('{"exit_code":42}'))).toBe(42); + }); + + it("parses a JSON exit_code of 0", () => { + expect(parseExitCode(Buffer.from('{"exit_code":0}'))).toBe(0); + }); + + it("maps a JSON null exit_code to 0", () => { + expect(parseExitCode(Buffer.from('{"exit_code":null}'))).toBe(0); + }); + + it("parses a 4-byte little-endian int32 payload", () => { + expect(parseExitCode(Buffer.from([42, 0, 0, 0]))).toBe(42); + }); + + it("parses a 4-byte LE int32 whose first byte is 0x7B ('{')", () => { + expect(parseExitCode(Buffer.from([123, 0, 0, 0]))).toBe(123); + }); + + it("returns 0 for an empty payload", () => { + expect(parseExitCode(Buffer.from([]))).toBe(0); + }); + }); + + describe("execute - provider host pre-check", () => { + it("returns an Err without opening a socket when the host is not https", async () => { + const { service } = setup(); + + const result = await service.execute(createShellExecInput({ providerBaseUrl: "http://provider.example.com" })); + + expect(result.ok).toBe(false); + expect(result.val).toContain("Invalid provider host"); + expect(vi.mocked(WebSocket)).not.toHaveBeenCalled(); + }); + + it("returns an Err when the host is an IP address", async () => { + const { service } = setup(); + + const result = await service.execute(createShellExecInput({ providerBaseUrl: "https://203.0.113.10:8443" })); + + expect(result.ok).toBe(false); + expect(result.val).toContain("Invalid provider host"); + }); + }); + + describe("execute - proxy routing (D1)", () => { + it("connects to the PROVIDER_PROXY_URL as wss and sends the envelope with the provider url and no data", async () => { + const { service, mockWs, getConstructorArgs } = setup(); + const input = createShellExecInput(); + + const promise = service.execute(input); + mockWs._trigger("open"); + + const [socketUrl, socketOptions] = getConstructorArgs(); + expect(socketUrl).toBe("wss://proxy.example.com"); + // No connection-level Authorization header: auth travels inside the envelope. + expect(socketOptions).toBeUndefined(); + + expect(mockWs.send).toHaveBeenCalledTimes(1); + const envelope = JSON.parse(mockWs.send.mock.calls[0][0]); + expect(envelope.type).toBe("websocket"); + expect(envelope.url).toBe(buildShellUrl(input)); + expect(envelope.url.startsWith("https://provider.example.com/lease/")).toBe(true); + expect(envelope.auth).toEqual({ type: "jwt", token: input.jwtToken }); + expect(envelope.providerAddress).toBe(input.providerAddress); + expect(envelope.isBase64).toBe(true); + expect("data" in envelope).toBe(false); + + // finish so the promise settles + mockWs._trigger("message", exitFrameJson(0)); + await promise; + }); + }); + + describe("execute - stdin injection (Task 2b)", () => { + const STDIN_MARKER = 104; + + it("sends a 104 stdin data frame plus a 104 EOF frame, keeping the secret out of the URL", async () => { + const { service, mockWs, getConstructorArgs } = setup(); + const secret = `SECRET=${faker.string.alphanumeric(16)}\nAPI_KEY=${faker.string.alphanumeric(12)}`; + const input = createShellExecInput({ + command: ["sh", "-c", "cat > /run/secrets/.env && chmod 600 /run/secrets/.env"], + stdin: secret + }); + + const promise = service.execute(input); + mockWs._trigger("open"); + + // Three frames: connect envelope, stdin data, stdin EOF. + expect(mockWs.send).toHaveBeenCalledTimes(3); + + const connect = JSON.parse(mockWs.send.mock.calls[0][0]); + const dataFrameMsg = JSON.parse(mockWs.send.mock.calls[1][0]); + const eofFrameMsg = JSON.parse(mockWs.send.mock.calls[2][0]); + + // Connect frame opens stdin (stdin=1) and carries no data. + expect(connect.url).toContain("stdin=1"); + expect("data" in connect).toBe(false); + + // The secret must not appear in the socket URL or the connect-frame url. + const [socketUrl] = getConstructorArgs(); + expect(socketUrl).toBe("wss://proxy.example.com"); + expect(String(socketUrl)).not.toContain(secret); + expect(connect.url).not.toContain(secret); + expect(connect.url).not.toContain(encodeURIComponent(secret)); + + // Data frame: full envelope + base64 of [104, ...secretBytes]. + expect(dataFrameMsg.type).toBe("websocket"); + expect(dataFrameMsg.url).toBe(connect.url); + expect(dataFrameMsg.auth).toEqual({ type: "jwt", token: input.jwtToken }); + expect(dataFrameMsg.providerAddress).toBe(input.providerAddress); + expect(dataFrameMsg.isBase64).toBe(true); + const dataBytes = Buffer.from(dataFrameMsg.data, "base64"); + expect(dataBytes[0]).toBe(STDIN_MARKER); + expect(dataBytes.subarray(1).toString("utf-8")).toBe(secret); + expect([...dataBytes]).toEqual([STDIN_MARKER, ...Buffer.from(secret, "utf-8")]); + + // The base64 payload itself must not be the plaintext secret (sanity). + expect(dataFrameMsg.data).not.toContain(secret); + + // EOF frame: full envelope + base64 of a bare [104] marker. + expect(eofFrameMsg.url).toBe(connect.url); + const eofBytes = Buffer.from(eofFrameMsg.data, "base64"); + expect([...eofBytes]).toEqual([STDIN_MARKER]); + + // A normal exit frame still yields the correct exit code. + mockWs._trigger("message", exitFrameJson(0)); + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).exitCode).toBe(0); + }); + + it("does not send any 104 frame when stdin is omitted", async () => { + const { service, mockWs } = setup(); + const input = createShellExecInput(); + + const promise = service.execute(input); + mockWs._trigger("open"); + + // Only the connect envelope is sent. + expect(mockWs.send).toHaveBeenCalledTimes(1); + + const connect = JSON.parse(mockWs.send.mock.calls[0][0]); + expect(connect.url).toContain("stdin=0"); + expect("data" in connect).toBe(false); + + // No sent frame decodes to a 104 stdin marker. + const has104Frame = mockWs.send.mock.calls.some(([raw]) => { + const msg = JSON.parse(raw); + if (typeof msg.data !== "string") return false; + return Buffer.from(msg.data, "base64")[0] === STDIN_MARKER; + }); + expect(has104Frame).toBe(false); + + mockWs._trigger("message", exitFrameJson(0)); + const result = await promise; + expect(result.ok).toBe(true); + }); + }); + + describe("execute - receive / marker handling", () => { + it("routes marker 100 to stdout with the marker byte stripped", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", dataFrame(100, "hello")); + mockWs._trigger("message", exitFrameJson(0)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).stdout).toBe("hello"); + expect((result.val as ShellExecOutput).stderr).toBe(""); + expect((result.val as ShellExecOutput).exitCode).toBe(0); + }); + + it("routes marker 101 to stderr with the marker byte stripped", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", dataFrame(101, "boom")); + mockWs._trigger("message", exitFrameJson(1)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).stderr).toBe("boom"); + expect((result.val as ShellExecOutput).stdout).toBe(""); + expect((result.val as ShellExecOutput).exitCode).toBe(1); + }); + + it("reads the exit code from a 102 JSON result frame", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", exitFrameJson(7)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).exitCode).toBe(7); + }); + + it("reads the exit code from a 102 4-byte LE int32 result frame", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", bytesFrame([102, 9, 0, 0, 0])); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).exitCode).toBe(9); + }); + + it("treats a 103 failure frame as a provider error (mapped 502), not output", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", dataFrame(103, "container terminated")); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toContain("Provider error"); + expect(result.val).toContain("container terminated"); + }); + + it("decodes a base64-string data payload", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + const base64 = Buffer.from([100, ...Buffer.from("hi", "utf-8")]).toString("base64"); + mockWs._trigger("open"); + mockWs._trigger("message", Buffer.from(JSON.stringify({ type: "websocket", message: { data: base64 } }))); + mockWs._trigger("message", exitFrameJson(0)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).stdout).toBe("hi"); + }); + }); + + describe("execute - error / robustness handling", () => { + it("resolves an error-key frame as a mapped provider error without decoding it", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger( + "message", + Buffer.from( + JSON.stringify({ + type: "websocket", + message: "Message doesn't match expected schema", + error: "Invalid message format", + errors: [{ path: ["url"], message: "URL must use https protocol" }] + }) + ) + ); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toContain("Provider error: Invalid message format"); + expect(result.val).toContain("url: URL must use https protocol"); + }); + + it("drops a non-base64 string payload instead of leaking it as output", async () => { + const { service, mockWs, logger } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", Buffer.from(JSON.stringify({ type: "websocket", message: { data: "this is not base64 !!!" } }))); + mockWs._trigger("message", exitFrameJson(0)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).stdout).toBe(""); + expect((result.val as ShellExecOutput).stderr).toBe(""); + expect(logger.warn).toHaveBeenCalled(); + }); + + it("ignores pong keepalive frames without corrupting output", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", dataFrame(100, "Hello ")); + mockWs._trigger("message", Buffer.from(JSON.stringify({ type: "pong" }))); + mockWs._trigger("message", dataFrame(100, "World")); + mockWs._trigger("message", exitFrameJson(0)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).stdout).toBe("Hello World"); + }); + + it("returns an Err with the connection message on a socket error", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("error", new Error("ECONNREFUSED")); + mockWs._trigger("close"); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toContain("WebSocket connection failed: ECONNREFUSED"); + }); + + it("returns an Err when the socket closes before an exit code arrives", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("close"); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toBe("Connection closed without exit code"); + }); + }); + + describe("execute - truncation (M5)", () => { + it("sets truncated but still reports the correct exit code when output exceeds the cap", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", dataFrame(100, "A".repeat(MAX_OUTPUT_SIZE + 1))); + // keeps reading past the cap so the exit frame is still processed + mockWs._trigger("message", exitFrameJson(3)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).truncated).toBe(true); + expect((result.val as ShellExecOutput).exitCode).toBe(3); + }); + + it("does not truncate output that is exactly at the cap", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + const data = "A".repeat(MAX_OUTPUT_SIZE); + mockWs._trigger("open"); + mockWs._trigger("message", dataFrame(100, data)); + mockWs._trigger("message", exitFrameJson(0)); + + const result = await promise; + expect(result.ok).toBe(true); + expect((result.val as ShellExecOutput).truncated).toBe(false); + expect((result.val as ShellExecOutput).stdout.length).toBe(MAX_OUTPUT_SIZE); + }); + }); + + describe("execute - auth expiry (M4)", () => { + it("maps a 4001 close frame to an auth-expired error, not a generic provider error", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("message", Buffer.from(JSON.stringify({ type: "websocket", message: "", closed: true, code: 4001, reason: "token expired" }))); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toContain("Auth expired"); + expect(result.val).not.toContain("Provider error"); + }); + + it("maps a 4003 ws close event to an auth-expired error", async () => { + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput()); + + mockWs._trigger("open"); + mockWs._trigger("close", 4003, Buffer.from("unauthorized")); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toContain("Auth expired"); + }); + }); + + describe("execute - timeout", () => { + it("resolves with a timeout error when the command runs past the timeout", async () => { + vi.useFakeTimers(); + const { service, mockWs } = setup(); + const promise = service.execute(createShellExecInput({ timeout: 5 })); + + mockWs._trigger("open"); + vi.advanceTimersByTime(5000); + + const result = await promise; + expect(result.ok).toBe(false); + expect(result.val).toBe("Command timed out"); + expect(mockWs.close).toHaveBeenCalled(); + + vi.useRealTimers(); + }); + }); + + function createMockWebSocket() { + const handlers: Record void>> = {}; + return { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + (handlers[event] ??= []).push(handler); + }), + send: vi.fn(), + close: vi.fn(), + _trigger(event: string, ...args: unknown[]) { + (handlers[event] || []).forEach(h => h(...args)); + } + }; + } + + function setup(overrides: { proxyUrl?: string } = {}) { + const mockWs = createMockWebSocket(); + let constructorArgs: unknown[] = []; + vi.mocked(WebSocket).mockImplementation(function (this: unknown, ...args: unknown[]) { + constructorArgs = args; + return mockWs as unknown as WebSocket; + } as unknown as typeof WebSocket); + + const logger = mock(); + const config = { PROVIDER_PROXY_URL: overrides.proxyUrl ?? PROXY_URL } as DeploymentConfig; + const service = new ShellExecService(config, logger); + + return { service, mockWs, logger, getConstructorArgs: () => constructorArgs }; + } + + function createShellExecInput(overrides?: Partial): ShellExecInput { + return { + providerBaseUrl: overrides?.providerBaseUrl ?? "https://provider.example.com", + providerAddress: overrides?.providerAddress ?? faker.string.alphanumeric(44), + dseq: overrides?.dseq ?? "1234", + gseq: overrides?.gseq ?? 1, + oseq: overrides?.oseq ?? 1, + service: overrides?.service ?? "test-service", + command: overrides?.command ?? ["echo", "Hello"], + stdin: overrides?.stdin, + timeout: overrides?.timeout ?? 60, + jwtToken: overrides?.jwtToken ?? faker.string.alphanumeric(100) + }; + } + + // Provider -> proxy -> client frame shape: { type:"websocket", message:{ type:"Buffer", data:number[] } } + function bytesFrame(data: number[]): Buffer { + return Buffer.from(JSON.stringify({ type: "websocket", message: { type: "Buffer", data } })); + } + + function dataFrame(marker: number, text: string): Buffer { + return bytesFrame([marker, ...Buffer.from(text, "utf-8")]); + } + + function exitFrameJson(code: number): Buffer { + return bytesFrame([102, ...Buffer.from(JSON.stringify({ exit_code: code }), "utf-8")]); + } +}); diff --git a/apps/api/src/deployment/services/shell-exec/shell-exec.service.ts b/apps/api/src/deployment/services/shell-exec/shell-exec.service.ts new file mode 100644 index 0000000000..53720f9855 --- /dev/null +++ b/apps/api/src/deployment/services/shell-exec/shell-exec.service.ts @@ -0,0 +1,370 @@ +import { isIP } from "node:net"; +import { Err, Ok, Result } from "ts-results"; +import { inject, singleton } from "tsyringe"; +import { WebSocket } from "ws"; + +import { LoggerService } from "@src/core"; +import { DEPLOYMENT_CONFIG, type DeploymentConfig } from "@src/deployment/config/config.provider"; + +/** + * Cap on the combined stdout+stderr the service buffers per exec, in UTF-8 + * bytes. Kept small (64 KiB) so many concurrent execs cannot exhaust memory — + * a synchronous buffered response is only safe while the buffer is bounded + * (per maintainer review on #3079/#3097). Exported so tests assert against the + * same constant instead of a hardcoded literal. + */ +export const MAX_OUTPUT_SIZE = 64 * 1024; + +/** + * LeaseShell protocol markers. The provider prefixes every shell frame with a + * single marker byte, mirroring `LeaseShellCode` in the web client + * (`apps/deploy-web/src/types/shell.ts`). + */ +const LeaseShellCode = { + Stdout: 100, + Stderr: 101, + Result: 102, + Failure: 103, + Stdin: 104 +} as const; + +export type ShellExecInput = { + providerBaseUrl: string; + providerAddress: string; + dseq: string; + gseq: number; + oseq: number; + service: string; + command: string[]; + timeout: number; + jwtToken: string; + /** + * Optional data piped to the command's standard input. Kept out of the URL/argv + * (which the provider-proxy logs) and streamed as `104` LeaseShellCodeStdin + * frames instead, so secrets never touch the logged request URL. + */ + stdin?: string; +}; + +export type ShellExecOutput = { + stdout: string; + stderr: string; + exitCode: number; + truncated: boolean; +}; + +type ProviderBufferMessage = { + type?: string; + data?: number[] | string; +}; + +type ReceivedMessage = { + type?: string; + message?: string | number[] | ProviderBufferMessage | null; + closed?: boolean; + code?: number; + reason?: string | number[] | { data?: number[] } | null; + error?: unknown; + errors?: Array<{ path?: Array; message?: string }>; +}; + +type WebSocketOutgoingMessage = { + type: "websocket"; + url: string; + auth: { type: "jwt"; token: string }; + providerAddress: string; + isBase64: boolean; +}; + +// A post-connect frame carries a base64 `data` payload alongside the full +// envelope. The provider-proxy re-validates the envelope on every frame, so the +// url/auth/providerAddress fields must be resent, not just `data`. +type WebSocketDataMessage = WebSocketOutgoingMessage & { data: string }; + +@singleton() +export class ShellExecService { + private readonly proxyWsUrl: string; + + constructor( + @inject(DEPLOYMENT_CONFIG) config: DeploymentConfig, + private readonly logger: LoggerService + ) { + this.proxyWsUrl = toProxyWebSocketUrl(config.PROVIDER_PROXY_URL); + this.logger.setContext(ShellExecService.name); + } + + async execute(input: ShellExecInput): Promise> { + if (!isValidProviderHost(input.providerBaseUrl)) { + return Err(`Invalid provider host: ${input.providerBaseUrl}`); + } + + const providerUrl = buildShellUrl(input); + const auth = { type: "jwt" as const, token: input.jwtToken }; + + return new Promise(resolve => { + let settled = false; + const settle = (result: Result) => { + if (settled) return; + settled = true; + clearTimeout(timeoutId); + resolve(result); + }; + + let stdout = ""; + let stderr = ""; + let outputBytes = 0; + let exitCode: number | undefined; + let truncated = false; + + // Connect to the provider-proxy (NOT the provider): no Authorization header, + // the provider credentials travel inside the envelope instead. + const ws = new WebSocket(this.proxyWsUrl); + + const timeoutId = setTimeout(() => { + settle(Err("Command timed out")); + ws.close(); + }, input.timeout * 1000); + + ws.on("open", () => { + // Connect frame. Omit `data`: with `stdin=0` the command runs one-shot + // from the URL query params, and the proxy's `!message.data` guard still + // opens and links the provider socket. + const message: WebSocketOutgoingMessage = { + type: "websocket", + url: providerUrl, + auth, + providerAddress: input.providerAddress, + isBase64: true + }; + ws.send(JSON.stringify(message)); + + // Task 2b: stream any stdin as `104` LeaseShellCodeStdin frames instead + // of embedding secrets in the (logged) URL/argv. The provider-proxy + // queues data frames until the provider socket is verified, so no fixed + // settle delay is needed here. Each frame resends the full envelope + // (url/auth/providerAddress) with a base64 `data` payload. + if (input.stdin && input.stdin.length > 0) { + const sendData = (payload: Buffer) => { + const dataMessage: WebSocketDataMessage = { ...message, data: payload.toString("base64") }; + ws.send(JSON.stringify(dataMessage)); + }; + // Marker byte 104 prefixes the raw UTF-8 stdin bytes. + sendData(Buffer.concat([Buffer.from([LeaseShellCode.Stdin]), Buffer.from(input.stdin, "utf-8")])); + // EOF: a bare `104` marker with no payload closes stdin so the command + // (e.g. `cat`) sees end-of-input and can exit. + sendData(Buffer.from([LeaseShellCode.Stdin])); + } + }); + + const appendOutput = (marker: number, text: string) => { + if (truncated) return; + // Enforce the cap in UTF-8 bytes (not UTF-16 code units) so multi-byte + // output is accounted for by its true wire size. + const chunkBytes = Buffer.byteLength(text, "utf8"); + if (outputBytes + chunkBytes <= MAX_OUTPUT_SIZE) { + outputBytes += chunkBytes; + if (marker === LeaseShellCode.Stdout) { + stdout += text; + } else { + stderr += text; + } + } else { + // M5: stop appending once over the cap, but keep reading so the + // exit/close frame still arrives and `exitCode` stays accurate. + truncated = true; + } + }; + + const handleClose = (code?: number, reason?: string) => { + if (isAuthExpiryClose(code, reason)) { + settle(Err(`Auth expired: provider closed connection${code !== undefined ? ` (code ${code})` : ""}`)); + return; + } + if (exitCode !== undefined) { + settle(Ok({ stdout, stderr, exitCode, truncated })); + } else { + settle(Err("Connection closed without exit code")); + } + }; + + ws.on("message", (raw: Buffer) => { + let message: ReceivedMessage; + try { + message = JSON.parse(raw.toString()) as ReceivedMessage; + } catch { + // Non-JSON frame: ignore. + return; + } + + // 1. Filter keepalives and unknown frame types. + if (message.type === "pong") return; + if (message.type !== undefined && message.type !== "websocket") return; + + // 2. Error BEFORE decode (M1). The proxy reports failures as a + // `type:"websocket"` frame carrying an `error` key — never decode it. + if (message.error !== undefined && message.error !== null && message.error !== "") { + settle(Err(`Provider error: ${flattenError(message)}`)); + ws.close(); + return; + } + + if (message.closed) { + handleClose(message.code, reasonToString(message.reason)); + ws.close(); + return; + } + + // 3. Decode payload. `message.message.data` is either a byte list (Node + // Buffer -> JSON) or a base64 string. + const inner = message.message; + const rawData = Array.isArray(inner) ? inner : inner && typeof inner === "object" ? (inner as ProviderBufferMessage).data : undefined; + + let bytes: Buffer; + if (Array.isArray(rawData)) { + bytes = Buffer.from(rawData); + } else if (typeof rawData === "string") { + if (!isStrictBase64(rawData)) { + // Never dispatch a non-base64 prose frame as fake output. + this.logger.warn({ event: "SHELL_EXEC_INVALID_BASE64", length: rawData.length }); + return; + } + bytes = Buffer.from(rawData, "base64"); + } else { + return; + } + + if (bytes.length === 0) return; + + // 4. Marker dispatch (M2). marker = data[0]; payload = data.slice(1). + const marker = bytes[0]; + const payload = bytes.subarray(1); + + switch (marker) { + case LeaseShellCode.Stdout: + case LeaseShellCode.Stderr: + appendOutput(marker, payload.toString("utf-8")); + return; + case LeaseShellCode.Result: + // M3: exit code is either JSON `{"exit_code":N}` or a 4-byte LE int32. + exitCode = parseExitCode(payload); + settle(Ok({ stdout, stderr, exitCode, truncated })); + ws.close(); + return; + case LeaseShellCode.Failure: { + // M2: distinct provider-error end state, not output. + const detail = payload.toString("utf-8").trim(); + settle(Err(`Provider error: shell failure${detail ? `: ${detail}` : ""}`)); + ws.close(); + return; + } + default: + // Unknown marker (e.g. stdin/resize echoes) — ignore. + return; + } + }); + + ws.on("error", err => { + settle(Err(`WebSocket connection failed: ${err.message}`)); + ws.close(); + }); + + ws.on("close", (code?: number, reason?: Buffer) => { + handleClose(code, reason ? reason.toString() : undefined); + }); + }); + } +} + +/** Maps an http(s) proxy URL to its ws(s) equivalent. */ +export function toProxyWebSocketUrl(url: string): string { + return url.replace(/^http/, "ws"); +} + +/** + * The provider-proxy silently rejects a lease-shell URL that is not https or + * that points at an IP / `.local` host. Pre-check so we can surface a clear + * error instead of a mystery timeout. + */ +export function isValidProviderHost(hostUri: string): boolean { + try { + const parsed = new URL(hostUri); + const hostname = parsed.hostname.startsWith("[") ? parsed.hostname.slice(1, -1) : parsed.hostname; + return parsed.protocol === "https:" && !hostname.endsWith(".local") && !isIP(hostname); + } catch { + return false; + } +} + +/** + * Builds the provider lease-shell URL consumed by the proxy as the envelope + * `url`. Each argv token maps to one `cmdN` query param; the provider execs the + * tokens as argv with no shell re-interpretation of its own. + */ +export function buildShellUrl(input: Pick): string { + const cmdParts = input.command.map((arg, i) => `&cmd${i}=${encodeURIComponent(arg)}`).join(""); + const baseUrl = input.providerBaseUrl.replace(/\/$/, ""); + // `stdin=1` opens the provider stdin channel so `104` frames are accepted; the + // payload itself is NEVER placed in the URL (it would be logged by the proxy). + const stdinFlag = input.stdin && input.stdin.length > 0 ? 1 : 0; + return `${baseUrl}/lease/${encodeURIComponent(input.dseq)}/${input.gseq}/${input.oseq}/shell?stdin=${stdinFlag}&tty=0&podIndex=0&service=${encodeURIComponent(input.service)}${cmdParts}`; +} + +/** + * Parses the exit-code payload of a `102` result frame. Providers use one of two + * encodings (version-dependent): a JSON body `{"exit_code":N}` or a raw 4-byte + * little-endian int32. `null` and missing codes map to 0. + */ +export function parseExitCode(payload: Buffer): number { + const text = payload.toString("utf-8").trim(); + + if (text.startsWith("{")) { + try { + const parsed = JSON.parse(text) as { exit_code?: unknown }; + const code = parsed?.exit_code; + // Successfully parsed JSON body: a null/absent/non-numeric code maps to 0. + return typeof code === "number" && Number.isFinite(code) ? code : 0; + } catch { + // Malformed JSON that only happens to start with "{" (e.g. a 4-byte LE + // int32 whose first byte is 0x7B) — fall through to the binary encoding. + } + } + + if (payload.length >= 4) return payload.readInt32LE(0); + return 0; +} + +/** Strict base64 validation so prose frames are dropped, not decoded as output. */ +export function isStrictBase64(value: string): boolean { + if (value.length === 0 || value.length % 4 !== 0) return false; + return /^[A-Za-z0-9+/]*={0,2}$/.test(value); +} + +function isAuthExpiryClose(code?: number, reason?: string): boolean { + if (code === 4001 || code === 4003) return true; + if (reason && /expired|unauthorized/i.test(reason)) return true; + return false; +} + +function reasonToString(reason: ReceivedMessage["reason"]): string | undefined { + if (reason === undefined || reason === null) return undefined; + if (typeof reason === "string") return reason; + if (Array.isArray(reason)) return Buffer.from(reason).toString("utf-8"); + if (Array.isArray(reason.data)) return Buffer.from(reason.data).toString("utf-8"); + return undefined; +} + +/** Flattens a proxy error frame, including any Zod-style `errors[].path: message`. */ +function flattenError(message: ReceivedMessage): string { + const base = typeof message.error === "string" ? message.error : JSON.stringify(message.error); + if (Array.isArray(message.errors) && message.errors.length > 0) { + const flat = message.errors + .map(issue => { + const path = Array.isArray(issue?.path) ? issue.path.join(".") : ""; + return path ? `${path}: ${issue?.message ?? ""}` : `${issue?.message ?? ""}`; + }) + .filter(Boolean) + .join("; "); + if (flat) return `${base} (${flat})`; + } + return base; +} diff --git a/apps/api/src/routers/open-api-handlers.ts b/apps/api/src/routers/open-api-handlers.ts index e07926f3ad..dd55fa8fa1 100644 --- a/apps/api/src/routers/open-api-handlers.ts +++ b/apps/api/src/routers/open-api-handlers.ts @@ -33,6 +33,7 @@ import { import { deploymentSettingRouter } from "@src/deployment/routes/deployment-setting/deployment-setting.router"; import { deploymentsRouter } from "@src/deployment/routes/deployments/deployments.router"; import { leasesRouter } from "@src/deployment/routes/leases/leases.router"; +import { shellExecRouter } from "@src/deployment/routes/shell-exec/shell-exec.router"; import { gpuRouter } from "@src/gpu"; import { pricingRouter } from "@src/pricing"; import { proposalsRouter } from "@src/proposal"; @@ -77,6 +78,7 @@ export const openApiHonoHandlers: OpenApiHonoHandler[] = [ deploymentSettingRouter, deploymentsRouter, leasesRouter, + shellExecRouter, apiKeysRouter, bidsRouter, certificateRouter, diff --git a/apps/api/src/routes/deployment/index.ts b/apps/api/src/routes/deployment/index.ts index 57b5140b09..c55f70a1bf 100644 --- a/apps/api/src/routes/deployment/index.ts +++ b/apps/api/src/routes/deployment/index.ts @@ -2,5 +2,6 @@ import { bidsRouter } from "@src/bid/routes/bids/bids.router"; import { certificateRouter } from "@src/certificate/routes/certificate.router"; import { deploymentsRouter } from "@src/deployment/routes/deployments/deployments.router"; import { leasesRouter } from "@src/deployment/routes/leases/leases.router"; +import { shellExecRouter } from "@src/deployment/routes/shell-exec/shell-exec.router"; -export default [deploymentsRouter, bidsRouter, certificateRouter, leasesRouter]; +export default [deploymentsRouter, bidsRouter, certificateRouter, leasesRouter, shellExecRouter]; diff --git a/package-lock.json b/package-lock.json index cb9f77f9e6..48e59317f1 100644 --- a/package-lock.json +++ b/package-lock.json @@ -104,6 +104,7 @@ "ts-results": "^3.3.0", "tsyringe": "^4.10.0", "unleash-client": "^6.6.0", + "ws": "^8.18.2", "zod": "3.*" }, "devDependencies": { @@ -124,6 +125,7 @@ "@types/pg": "^8.11.6", "@types/semver": "^7.5.2", "@types/tar": "^6.1.13", + "@types/ws": "^8.5.4", "@typescript-eslint/eslint-plugin": "^7.12.0", "@vitest/coverage-v8": "^4.1.5", "drizzle-kit": "^0.22.7",