@dbx-tools/lakebox 0.0.0-stage → 0.9.132

package/src/rest.ts

@@ -0,0 +1,264 @@
+import { errorUtils, log, net } from "@dbx-tools/shared-core";
+import { z } from "zod";
+
+const LAKEBOX_ROOT = "/api/2.0/lakebox";
+const LAKEBOXES_PATH = `${LAKEBOX_ROOT}/sandboxes`;
+const SSH_KEYS_PATH = `${LAKEBOX_ROOT}/ssh-keys`;
+const ORG_ID_HEADER = "X-Databricks-Org-Id";
+const logger = log.logger("lakebox:rest");
+
+export const LakeboxEntrySchema = z
+  .object({
+    sandboxId: z.string().describe("Lakebox identifier; also the SSH gateway username."),
+    status: z.string().describe("Upstream lifecycle status, such as Running or Stopped."),
+    gatewayHost: z.string().optional().describe("SSH gateway hostname."),
+    name: z.string().optional().describe("Display label."),
+    createTime: z.string().optional().describe("Creation time."),
+    lastStartTime: z.string().optional().describe("Last start time."),
+    idleTimeout: z.string().optional().describe("Platform idle timeout, such as 600s."),
+    noAutostop: z.boolean().optional().describe("Whether idle auto-stop is disabled."),
+  })
+  .passthrough()
+  .describe("Lakebox record returned by the Databricks Lakebox REST API.");
+/** Lakebox record returned by the Databricks Lakebox REST API. */
+export type LakeboxEntry = z.infer<typeof LakeboxEntrySchema>;
+
+const LakeboxListSchema = z
+  .object({
+    sandboxes: z.array(LakeboxEntrySchema).default([]),
+    nextPageToken: z.string().optional(),
+  })
+  .passthrough();
+
+export const LakeboxSshKeySchema = z
+  .object({
+    keyHash: z.string().describe("Registered public key hash."),
+    name: z.string().optional().describe("Key label."),
+    gatewayHost: z.string().optional().describe("Gateway hostname for this key."),
+  })
+  .passthrough()
+  .describe("Registered SSH key returned by the Databricks Lakebox REST API.");
+/** Registered SSH key returned by the Databricks Lakebox REST API. */
+export type LakeboxSshKey = z.infer<typeof LakeboxSshKeySchema>;
+
+const LakeboxSshKeyListSchema = z
+  .object({ sshKeys: z.array(LakeboxSshKeySchema).default([]) })
+  .passthrough();
+
+/**
+ * Workspace credentials for Lakebox REST calls.
+ *
+ * The Databricks SDK and AppKit `WorkspaceClient.config` satisfy this shape, as
+ * does a small adapter over `@dbx-tools/auth`. Lakebox lifecycle calls must use
+ * the app service principal, never request-scoped OBO credentials.
+ */
+export interface LakeboxAuth {
+  /** Workspace URL. */
+  readonly host?: string;
+  /** Numeric workspace identifier, when known. */
+  readonly workspaceId?: string;
+  /** Resolves lazy configuration before `host` is read. */
+  ensureResolved?(): Promise<void>;
+  /** Adds authentication headers to one request. */
+  authenticate(headers: Headers): Promise<void>;
+}
+
+/** Error raised for a non-success Lakebox REST response. */
+export class LakeboxApiError extends Error {
+  constructor(
+    message: string,
+    readonly status: number,
+    readonly errorCode?: string,
+    readonly requestId?: string,
+  ) {
+    super(message);
+    this.name = "LakeboxApiError";
+  }
+}
+
+type RequestMethod = "GET" | "POST" | "PATCH" | "DELETE";
+
+/**
+ * REST client for Databricks Lakeboxes and registered SSH keys.
+ *
+ * Responses are validated before they reach callers. Diagnostics never include
+ * tokens or request headers.
+ */
+export class LakeboxClient {
+  /** Creates a client for one workspace credential. */
+  constructor(private readonly auth: LakeboxAuth) {}
+
+  /** Returns the workspace identifier, or the workspace origin when the identifier is unknown. */
+  async workspaceScope(): Promise<string> {
+    await this.auth.ensureResolved?.();
+    return (
+      process.env.DATABRICKS_WORKSPACE_ID?.trim() ||
+      this.auth.workspaceId?.trim() ||
+      net.urlBuilder(this.auth.host ?? "")?.origin ||
+      "workspace"
+    );
+  }
+
+  /** Lists every Lakebox visible to the caller, following pagination. */
+  async list(signal?: AbortSignal): Promise<LakeboxEntry[]> {
+    const entries: LakeboxEntry[] = [];
+    let pageToken: string | undefined;
+    do {
+      const query = new URLSearchParams({ page_size: "100" });
+      if (pageToken) query.set("page_token", pageToken);
+      const response = LakeboxListSchema.parse(
+        await this.request("GET", `${LAKEBOXES_PATH}?${query}`, undefined, signal),
+      );
+      entries.push(...response.sandboxes);
+      pageToken = response.nextPageToken || undefined;
+    } while (pageToken);
+    return entries;
+  }
+
+  /** Returns one Lakebox by ID. */
+  async get(id: string, signal?: AbortSignal): Promise<LakeboxEntry> {
+    return LakeboxEntrySchema.parse(await this.request("GET", lakeboxPath(id), undefined, signal));
+  }
+
+  /** Creates a Lakebox with a Databricks-assigned display name. */
+  async create(signal?: AbortSignal): Promise<LakeboxEntry> {
+    return LakeboxEntrySchema.parse(
+      await this.request("POST", LAKEBOXES_PATH, { sandbox: {} }, signal),
+    );
+  }
+
+  /** Sets the platform idle timeout that stops the Lakebox after inactivity. */
+  async setIdleTimeout(id: string, seconds: number, signal?: AbortSignal): Promise<LakeboxEntry> {
+    return LakeboxEntrySchema.parse(
+      await this.request(
+        "PATCH",
+        lakeboxPath(id),
+        { sandbox_id: id, idle_timeout: `${seconds}s` },
+        signal,
+      ),
+    );
+  }
+
+  /** Permanently deletes one Lakebox and its persistent home. */
+  async delete(id: string, signal?: AbortSignal): Promise<void> {
+    await this.request("DELETE", lakeboxPath(id), undefined, signal);
+  }
+
+  /** Requests a start and returns the updated record. */
+  async start(id: string, signal?: AbortSignal): Promise<LakeboxEntry> {
+    return LakeboxEntrySchema.parse(
+      await this.request("POST", `${lakeboxPath(id)}/start`, { sandbox_id: id }, signal),
+    );
+  }
+
+  /** Requests a stop and returns the updated record. */
+  async stop(id: string, signal?: AbortSignal): Promise<LakeboxEntry> {
+    return LakeboxEntrySchema.parse(
+      await this.request("POST", `${lakeboxPath(id)}/stop`, { sandbox_id: id }, signal),
+    );
+  }
+
+  /** Lists SSH public keys registered by the caller. */
+  async listKeys(signal?: AbortSignal): Promise<LakeboxSshKey[]> {
+    return LakeboxSshKeyListSchema.parse(
+      await this.request("GET", SSH_KEYS_PATH, undefined, signal),
+    ).sshKeys;
+  }
+
+  /** Registers an SSH public key and returns the Databricks key record. */
+  async registerKey(publicKey: string, name: string, signal?: AbortSignal): Promise<LakeboxSshKey> {
+    return LakeboxSshKeySchema.parse(
+      await this.request("POST", SSH_KEYS_PATH, { public_key: publicKey, name }, signal),
+    );
+  }
+
+  private async request(
+    method: RequestMethod,
+    path: string,
+    payload: unknown,
+    signal?: AbortSignal,
+  ): Promise<unknown> {
+    await this.auth.ensureResolved?.();
+    const host = this.auth.host?.replace(/\/+$/, "");
+    if (!host) throw new Error("Lakebox REST requests require a Databricks workspace host");
+    const headers = new Headers({ Accept: "application/json" });
+    const workspaceId = this.auth.workspaceId;
+    if (workspaceId && workspaceId !== "none") headers.set(ORG_ID_HEADER, workspaceId);
+    if (payload !== undefined) headers.set("Content-Type", "application/json");
+    await this.auth.authenticate(headers);
+    const startedAt = Date.now();
+    const metadata = { method, path: path.split("?")[0], hasPayload: payload !== undefined };
+    logger.debug("lakebox REST request", metadata);
+    let response: Response;
+    try {
+      response = await fetch(`${host}${path}`, {
+        method,
+        headers,
+        ...(payload === undefined ? {} : { body: JSON.stringify(payload) }),
+        ...(signal ? { signal } : {}),
+      });
+    } catch (error) {
+      logger.warn("lakebox REST request failed", {
+        ...metadata,
+        durationMs: Date.now() - startedAt,
+        error: errorUtils.toError(error),
+      });
+      throw error;
+    }
+    const text = await response.text();
+    const body = parseBody(text);
+    if (!response.ok) {
+      const record = z
+        .object({
+          message: z.string().optional(),
+          error_code: z.string().optional(),
+          request_id: z.string().optional(),
+        })
+        .passthrough()
+        .safeParse(body);
+      const details = record.success ? record.data : {};
+      const error = new LakeboxApiError(
+        details.message || `Lakebox request failed (${response.status})`,
+        response.status,
+        details.error_code,
+        details.request_id ?? response.headers.get("x-request-id") ?? undefined,
+      );
+      logger.warn("lakebox REST request failed", {
+        ...metadata,
+        durationMs: Date.now() - startedAt,
+        status: error.status,
+        errorCode: error.errorCode,
+        requestId: error.requestId,
+      });
+      throw error;
+    }
+    logger.debug("lakebox REST response", { ...metadata, durationMs: Date.now() - startedAt });
+    return body ?? {};
+  }
+}
+
+/**
+ * Returns whether an API error indicates that Lakeboxes are unavailable in the
+ * workspace or region.
+ */
+export function isLakeboxUnavailable(error: unknown): boolean {
+  if (!(error instanceof LakeboxApiError)) return false;
+  return (
+    error.status === 404 ||
+    error.status === 503 ||
+    /not available in your region/i.test(error.message)
+  );
+}
+
+function lakeboxPath(id: string): string {
+  return `${LAKEBOXES_PATH}/${encodeURIComponent(id)}`;
+}
+
+function parseBody(text: string): unknown {
+  if (!text.trim()) return undefined;
+  try {
+    return JSON.parse(text) as unknown;
+  } catch {
+    return { message: text.slice(0, 4_000) };
+  }
+}

package/src/session.ts

@@ -0,0 +1,371 @@
+import { createHash } from "node:crypto";
+import { isIP } from "node:net";
+import type { Duplex } from "node:stream";
+import { errorUtils, log } from "@dbx-tools/shared-core";
+import ssh2, { type Client as Ssh2Client } from "ssh2";
+import { z } from "zod";
+import {
+  createSshTransport,
+  SshCommand,
+  sshCommandScript,
+  type SshCommandIdentity,
+  type SshTransport,
+} from "./_ssh.ts";
+import type { Ssh2PatchMode, SshTransportMode } from "./transport.ts";
+
+export { SshCommand, sshCommandScript, type SshCommandIdentity } from "./_ssh.ts";
+
+const { Client } = ssh2;
+const logger = log.logger("lakebox:ssh");
+const SERVICE_DIRECT_SUFFIX = ".service-direct.cloud.databricks.com";
+
+const DnsAnswerSchema = z
+  .object({
+    type: z.number().int().describe("DNS record type number."),
+    data: z.string().describe("DNS record value."),
+  })
+  .passthrough()
+  .describe("One DNS-over-HTTPS answer record.");
+
+const DnsResponseSchema = z
+  .object({
+    Answer: z.array(DnsAnswerSchema).optional().describe("Resolved DNS answer records."),
+  })
+  .passthrough()
+  .describe("Google DNS-over-HTTPS JSON response.");
+
+/** Connection, key, host-verification, and command settings for one SSH machine. */
+export interface SshSessionOptions {
+  /** Databricks Lakebox gateway hostname. */
+  host: string;
+  /** SSH gateway port. */
+  port: number;
+  /** SSH username accepted by the gateway: the Lakebox ID. */
+  username: string;
+  /** In-memory private key used for the verified SSH handshake and ssh2 fallback. */
+  privateKey: string;
+  /** Private key file used by native OpenSSH. */
+  privateKeyPath: string;
+  /** Dedicated OpenSSH known-host file for this workspace, user, and gateway. */
+  knownHostsPath: string;
+  /** Connection timeout in milliseconds. */
+  connectTimeoutMs: number;
+  /** SSH transport selected for this session. */
+  sshTransport: SshTransportMode;
+  /** ssh2 missing-exec-reply patch policy. */
+  ssh2Patch: Ssh2PatchMode;
+  /** Previously trusted SHA-256 host-key fingerprint. */
+  knownHost?: string;
+  /** Persists the accepted host key for both ssh2 and native OpenSSH. */
+  rememberHost?: (fingerprint: string, key: Buffer) => Promise<void>;
+  /** Native OpenSSH executable. Defaults to `ssh` on `PATH`. */
+  sshCommand?: string;
+  /** Resolves public fallback addresses for a service-direct gateway. */
+  resolvePublicAddresses?: (host: string) => Promise<string[]>;
+}
+
+/**
+ * Verified SSH session that delegates commands and remote TCP connections to
+ * the transport selected by `src/_ssh.ts`.
+ *
+ * The initial ssh2 handshake verifies and persists the Databricks gateway host
+ * key. The selected transport then starts its pooled command and forwarding
+ * resources. Resolved transport and patch policy determine whether the verified
+ * client seeds the ssh2 pool or is replaced by native OpenSSH ControlMasters.
+ */
+export class SshSession {
+  private readonly activeCommands = new Set<SshCommand>();
+  private readonly pendingClients = new Set<Ssh2Client>();
+  private transport?: SshTransport;
+  private connecting?: Promise<void>;
+  private closing?: Promise<void>;
+  private knownHost?: string;
+  private closed = false;
+
+  /** Creates a disconnected session. Call {@link connect} or another operation to connect. */
+  constructor(private readonly options: SshSessionOptions) {
+    this.knownHost = options.knownHost;
+  }
+
+  /**
+   * Verifies the gateway, resolves its public fallback when necessary, and
+   * selects the command and forwarding transport.
+   */
+  async connect(): Promise<void> {
+    if (this.closed) throw new Error("SSH session is closed");
+    if (this.transport) return;
+    if (this.connecting) return this.connecting;
+    this.connecting = this.connectWithFallback().catch((error) => {
+      this.connecting = undefined;
+      throw error;
+    });
+    return this.connecting;
+  }
+
+  private async connectWithFallback(): Promise<void> {
+    try {
+      await this.connectHost(this.options.host);
+    } catch (error) {
+      if (!shouldResolvePublicGateway(this.options.host, error)) throw error;
+      const resolveAddresses = this.options.resolvePublicAddresses ?? resolvePublicGatewayAddresses;
+      let addresses: string[];
+      try {
+        addresses = await resolveAddresses(this.options.host);
+      } catch (resolveError) {
+        logger.warn("public Lakebox gateway resolution failed", {
+          host: this.options.host,
+          error: errorUtils.toError(resolveError),
+        });
+        throw error;
+      }
+      logger.warn("private Lakebox gateway unreachable; retrying public addresses", {
+        host: this.options.host,
+        addressCount: addresses.length,
+      });
+      let lastError: unknown = error;
+      for (const address of addresses) {
+        try {
+          await this.connectHost(address);
+          return;
+        } catch (connectError) {
+          lastError = connectError;
+        }
+      }
+      throw lastError;
+    }
+  }
+
+  private async connectHost(connectHost: string): Promise<void> {
+    if (this.closed) throw new Error("SSH session is closed");
+    const client = await this.createVerifiedClient(connectHost, true);
+    if (this.closed) {
+      client.destroy();
+      throw new Error("SSH session is closed");
+    }
+    try {
+      const transport = await createSshTransport({
+        host: this.options.host,
+        connectHost,
+        port: this.options.port,
+        username: this.options.username,
+        privateKeyPath: this.options.privateKeyPath,
+        knownHostsPath: this.options.knownHostsPath,
+        connectTimeoutMs: this.options.connectTimeoutMs,
+        ssh2Client: client,
+        createSsh2CommandClient: () => this.createVerifiedClient(connectHost, false),
+        transport: this.options.sshTransport,
+        ssh2Patch: this.options.ssh2Patch,
+        ...(this.options.sshCommand ? { sshCommand: this.options.sshCommand } : {}),
+      });
+      if (this.closed) {
+        await transport.close();
+        throw new Error("SSH session is closed");
+      }
+      this.transport = transport;
+      this.connecting = undefined;
+      if (transport.kind === "openssh") client.destroy();
+    } catch (error) {
+      client.destroy();
+      throw error;
+    }
+  }
+
+  private createVerifiedClient(connectHost: string, rememberHost: boolean): Promise<Ssh2Client> {
+    return new Promise((resolve, reject) => {
+      if (this.closed) {
+        reject(new Error("SSH session is closed"));
+        return;
+      }
+      const client = new Client();
+      this.pendingClients.add(client);
+      let verifiedHost: { fingerprint: string; key: Buffer } | undefined;
+      let settled = false;
+      const finishSettlement = () => {
+        this.pendingClients.delete(client);
+      };
+      client.once("ready", () => {
+        const finish = async () => {
+          if (!verifiedHost) throw new Error("SSH gateway did not provide a verified host key");
+          this.knownHost = verifiedHost.fingerprint;
+          if (rememberHost) {
+            await this.options.rememberHost?.(verifiedHost.fingerprint, verifiedHost.key);
+          }
+          settled = true;
+          finishSettlement();
+          resolve(client);
+        };
+        void finish().catch((error) => {
+          client.destroy();
+          if (settled) return;
+          settled = true;
+          finishSettlement();
+          reject(error);
+        });
+      });
+      client.once("error", (error) => {
+        if (settled) return;
+        settled = true;
+        finishSettlement();
+        reject(error);
+      });
+      client.once("close", () => {
+        if (settled) return;
+        settled = true;
+        finishSettlement();
+        reject(new Error("SSH gateway closed before authentication completed"));
+      });
+      try {
+        client.connect({
+          host: connectHost,
+          port: this.options.port,
+          username: this.options.username,
+          privateKey: this.options.privateKey,
+          readyTimeout: this.options.connectTimeoutMs,
+          keepaliveInterval: 15_000,
+          keepaliveCountMax: 3,
+          hostVerifier: (key: Buffer) => {
+            const fingerprint = createHash("sha256").update(key).digest("base64");
+            if (this.knownHost && fingerprint !== this.knownHost) return false;
+            verifiedHost = { fingerprint, key: Buffer.from(key) };
+            return true;
+          },
+        });
+      } catch (error) {
+        settled = true;
+        finishSettlement();
+        client.destroy();
+        reject(error);
+      }
+    });
+  }
+
+  /**
+   * Starts one command as a Lakebox user and returns its live output and
+   * completion state immediately.
+   */
+  execute(
+    command: string,
+    args: string[],
+    options: {
+      user: { name: string; home: string };
+      cwd?: string;
+      env?: Record<string, string>;
+      timeoutMs: number;
+      signal?: AbortSignal;
+    },
+  ): SshCommand {
+    const { user, ...rest } = options;
+    return this.startCommand(command, args, rest, { kind: "user", ...user });
+  }
+
+  /**
+   * @internal Starts one fixed package-owned bootstrap command as root.
+   *
+   * Never route REST, agent-tool, service, or environment input here.
+   */
+  executePrivileged(
+    command: string,
+    args: string[],
+    options: { timeoutMs: number; signal?: AbortSignal },
+  ): SshCommand {
+    return this.startCommand(command, args, options, { kind: "root" });
+  }
+
+  private startCommand(
+    command: string,
+    args: string[],
+    options: {
+      cwd?: string;
+      env?: Record<string, string>;
+      timeoutMs: number;
+      signal?: AbortSignal;
+    },
+    identity: SshCommandIdentity,
+  ): SshCommand {
+    const script = sshCommandScript(command, args, options.cwd, options.env, identity);
+    const sshCommand = new SshCommand(command, args, options, async (execution) => {
+      await this.connect();
+      if (execution.signal.aborted) return 130;
+      const transport = this.transport;
+      if (!transport) throw new Error("SSH transport did not initialize");
+      return transport.run(script, execution);
+    });
+    this.activeCommands.add(sshCommand);
+    const remove = () => this.activeCommands.delete(sshCommand);
+    void sshCommand.result.then(remove, remove);
+    return sshCommand;
+  }
+
+  /** Opens a remote TCP connection through native SOCKS5 or ssh2 direct forwarding. */
+  async forwardOut(remoteHost: string, remotePort: number): Promise<Duplex> {
+    await this.connect();
+    const transport = this.transport;
+    if (!transport) throw new Error("SSH transport did not initialize");
+    return transport.connect(remoteHost, remotePort);
+  }
+
+  /** Cancels active commands, closes the selected transport, and clears connection state. */
+  close(): Promise<void> {
+    if (this.closing) return this.closing;
+    this.closed = true;
+    this.closing = this.closeResources();
+    return this.closing;
+  }
+
+  private async closeResources(): Promise<void> {
+    const cancellations = [...this.activeCommands].map((command) => command.cancel());
+    this.activeCommands.clear();
+    for (const client of this.pendingClients) client.destroy();
+    this.pendingClients.clear();
+    const transport = this.transport;
+    const connecting = this.connecting;
+    this.transport = undefined;
+    this.connecting = undefined;
+    const results = await Promise.allSettled([
+      ...cancellations,
+      ...(transport ? [transport.close()] : []),
+      ...(connecting ? [connecting] : []),
+    ]);
+    for (const result of results) {
+      if (result.status === "rejected") {
+        logger.warn("failed to close SSH session resource", {
+          error: errorUtils.toError(result.reason),
+        });
+      }
+    }
+  }
+}
+
+/**
+ * Resolves public IPv4 addresses for a Databricks service-direct gateway.
+ *
+ * Returns an empty list for other hostnames. DNS-over-HTTPS requests time out
+ * after five seconds and reject invalid responses.
+ */
+export async function resolvePublicGatewayAddresses(
+  host: string,
+  fetcher: typeof fetch = fetch,
+): Promise<string[]> {
+  if (!host.toLowerCase().endsWith(SERVICE_DIRECT_SUFFIX)) return [];
+  const response = await fetcher(
+    `https://dns.google/resolve?name=${encodeURIComponent(host)}&type=A`,
+    { headers: { Accept: "application/dns-json" }, signal: AbortSignal.timeout(5_000) },
+  );
+  if (!response.ok) throw new Error(`Public DNS resolution failed (${response.status})`);
+  const payload = DnsResponseSchema.parse(await response.json());
+  return [
+    ...new Set(
+      (payload.Answer ?? [])
+        .filter((answer) => answer.type === 1 && isIP(answer.data) === 4)
+        .map((answer) => answer.data),
+    ),
+  ];
+}
+
+function shouldResolvePublicGateway(host: string, error: unknown): boolean {
+  if (!host.toLowerCase().endsWith(SERVICE_DIRECT_SUFFIX)) return false;
+  const code =
+    (error as { code?: unknown; cause?: { code?: unknown } })?.code ??
+    (error as { cause?: { code?: unknown } })?.cause?.code;
+  return code === "EHOSTUNREACH" || code === "ENETUNREACH" || code === "ETIMEDOUT";
+}

package/src/transport.ts

@@ -0,0 +1,25 @@
+import { z } from "zod";
+
+export const SshTransportModeSchema = z
+  .string()
+  .trim()
+  .toLowerCase()
+  .pipe(z.enum(["auto", "openssh", "ssh2"]))
+  .default("auto")
+  .describe(
+    "SSH transport: auto prefers native OpenSSH with SOCKS5 forwarding and otherwise selects ssh2; explicit modes never switch.",
+  );
+/** SSH transport selected for a Lakebox session. */
+export type SshTransportMode = z.output<typeof SshTransportModeSchema>;
+
+export const Ssh2PatchModeSchema = z
+  .string()
+  .trim()
+  .toLowerCase()
+  .pipe(z.enum(["auto", "true", "false"]))
+  .default("auto")
+  .describe(
+    "ssh2 missing-exec-reply workaround: auto patches only in development, true requires it, false never patches.",
+  );
+/** Missing exec-reply patch policy applied when ssh2 is selected. */
+export type Ssh2PatchMode = z.output<typeof Ssh2PatchModeSchema>;