@dbx-tools/lakebox 0.0.0-stage → 0.9.132

package/src/_ssh_ssh2.ts

@@ -0,0 +1,273 @@
+import type { Duplex } from "node:stream";
+import { log } from "@dbx-tools/shared-core";
+import type { Client as Ssh2Client, ClientChannel } from "ssh2";
+import {
+  PooledSshTransport,
+  type SshCommandExecution,
+  type SshPoolOptions,
+  type SshResourceLifecycle,
+  type SshTransportOptions,
+} from "./_ssh.ts";
+
+const logger = log.logger("lakebox:ssh2");
+const patchedSsh2Protocols = new WeakSet<Ssh2Protocol>();
+
+interface Ssh2Resource {
+  readonly client: Ssh2Client;
+  alive: boolean;
+}
+
+/** Exclusive ssh2 connection pool for commands and direct TCP forwarding. */
+export class Ssh2Transport extends PooledSshTransport<Ssh2Resource> {
+  readonly kind = "ssh2" as const;
+  private readonly connectTimeoutMs: number;
+
+  constructor(options: SshTransportOptions, poolOptions: SshPoolOptions = {}) {
+    super(ssh2Lifecycle(options), {
+      acquireTimeoutMillis: options.connectTimeoutMs,
+      ...poolOptions,
+    });
+    this.connectTimeoutMs = options.connectTimeoutMs;
+  }
+
+  protected runWithResource(
+    resource: Ssh2Resource,
+    script: string,
+    execution: SshCommandExecution,
+  ): Promise<number> {
+    return executeCommand(resource, script, execution);
+  }
+
+  protected connectWithResource(
+    resource: Ssh2Resource,
+    remoteHost: string,
+    remotePort: number,
+  ): Promise<Duplex> {
+    return new Promise((resolve, reject) => {
+      let settled = false;
+      let timeout: ReturnType<typeof setTimeout>;
+      const cleanup = () => {
+        clearTimeout(timeout);
+        resource.client.removeListener("close", clientClosed);
+        resource.client.removeListener("error", clientFailed);
+      };
+      const fail = (error: Error) => {
+        if (settled) return;
+        settled = true;
+        cleanup();
+        reject(error);
+      };
+      const clientClosed = () => {
+        fail(new Error("ssh2 connection closed while opening forwarding"));
+      };
+      const clientFailed = (error: Error) => {
+        fail(error);
+      };
+      timeout = setTimeout(
+        () => fail(new Error(`ssh2 forwarding did not open within ${this.connectTimeoutMs}ms`)),
+        this.connectTimeoutMs,
+      );
+      timeout.unref?.();
+      resource.client.once("close", clientClosed);
+      resource.client.once("error", clientFailed);
+      resource.client.forwardOut("127.0.0.1", 0, remoteHost, remotePort, (error, channel) => {
+        if (settled) {
+          channel?.destroy();
+          return;
+        }
+        settled = true;
+        cleanup();
+        if (error) reject(error);
+        else resolve(channel);
+      });
+    });
+  }
+}
+
+function ssh2Lifecycle(options: SshTransportOptions): SshResourceLifecycle<Ssh2Resource> {
+  const patchMode = options.ssh2Patch ?? "auto";
+  const requirePatch = patchMode === "true";
+  const patchCommands = requirePatch || (patchMode === "auto" && isNodeDevelopment());
+  let initialClient: Ssh2Client | undefined = options.ssh2Client;
+  const initialPatched = patchCommands ? patchSsh2ExecWithoutReply(options.ssh2Client) : false;
+  if (requirePatch && !initialPatched) {
+    options.ssh2Client.destroy();
+    throw new Error("Installed ssh2 internals do not support the required exec no-reply patch");
+  }
+  let logged = false;
+  return {
+    create: async () => {
+      const seeded = initialClient !== undefined;
+      const client = initialClient ?? (await options.createSsh2CommandClient());
+      initialClient = undefined;
+      const resource = trackClient(client);
+      const patched = seeded
+        ? initialPatched
+        : patchCommands
+          ? patchSsh2ExecWithoutReply(client)
+          : false;
+      if (requirePatch && !patched) {
+        client.destroy();
+        throw new Error("Installed ssh2 internals do not support the required exec no-reply patch");
+      }
+      if (!logged) {
+        logged = true;
+        logger.warn("using pooled ssh2 Lakebox transport", {
+          patchMode,
+          patched,
+          nodeEnv: process.env.NODE_ENV,
+          transportMode: options.transport ?? "auto",
+        });
+      }
+      return resource;
+    },
+    destroy: (resource) => closeClient(resource, options.connectTimeoutMs),
+    // A remote no-op would consume the connection's only permitted operation.
+    test: async (resource) => resource.alive,
+  };
+}
+
+function trackClient(client: Ssh2Client): Ssh2Resource {
+  const resource: Ssh2Resource = { client, alive: true };
+  const markClosed = () => {
+    resource.alive = false;
+  };
+  client.once("close", markClosed);
+  client.once("error", markClosed);
+  return resource;
+}
+
+async function closeClient(resource: Ssh2Resource, timeoutMs: number): Promise<void> {
+  if (!resource.alive) {
+    resource.client.destroy();
+    return;
+  }
+  await new Promise<void>((resolve) => {
+    let settled = false;
+    const finish = () => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timeout);
+      resource.alive = false;
+      resolve();
+    };
+    const timeout = setTimeout(
+      () => {
+        resource.client.destroy();
+        finish();
+      },
+      Math.min(timeoutMs, 1_000),
+    );
+    timeout.unref?.();
+    resource.client.once("close", finish);
+    resource.client.once("error", finish);
+    resource.client.end();
+  });
+}
+
+function executeCommand(
+  resource: Ssh2Resource,
+  script: string,
+  execution: SshCommandExecution,
+): Promise<number> {
+  return new Promise<number>((resolve, reject) => {
+    let channel: ClientChannel | undefined;
+    let exitCode: number | undefined;
+    let settled = false;
+    const cleanup = () => {
+      resource.client.removeListener("close", clientClosed);
+      resource.client.removeListener("error", clientClosed);
+    };
+    const finish = (code: number) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      resolve(code);
+    };
+    const fail = (error: Error) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      reject(error);
+    };
+    const clientClosed = () => fail(new Error("ssh2 connection closed during command execution"));
+    resource.client.once("close", clientClosed);
+    resource.client.once("error", clientClosed);
+    execution.setCancel(() => {
+      if (settled) return;
+      try {
+        channel?.signal("KILL");
+      } catch {
+        // The remote may close the channel before the signal is written.
+      }
+      finish(130);
+      channel?.close();
+    });
+    if (execution.signal.aborted) {
+      finish(130);
+      return;
+    }
+    resource.client.exec(script, (error, openedChannel) => {
+      if (error) {
+        fail(error);
+        return;
+      }
+      if (settled) {
+        openedChannel.close();
+        return;
+      }
+      channel = openedChannel;
+      channel.on("data", execution.stdout);
+      channel.stderr.on("data", execution.stderr);
+      channel.on("exit", (code: number | null) => {
+        exitCode = code ?? 1;
+      });
+      channel.once("error", fail);
+      channel.once("close", () => finish(exitCode ?? 1));
+    });
+  });
+}
+
+/** Apply the Databricks gateway exec-request workaround when internals permit it. */
+function patchSsh2ExecWithoutReply(client: Ssh2Client): boolean {
+  const internals = client as unknown as Ssh2ClientInternals;
+  const protocol = internals._protocol;
+  const channelManager = internals._chanMgr;
+  if (!protocol || typeof protocol.exec !== "function" || !channelManager) {
+    return false;
+  }
+  if (patchedSsh2Protocols.has(protocol)) return true;
+  const originalExec = protocol.exec.bind(protocol);
+  protocol.exec = (channelId, command) => {
+    originalExec(channelId, command, false);
+    const channel = Object.values(channelManager._channels).find(
+      (candidate) => candidate?.outgoing.id === channelId,
+    );
+    const callback = channel?._callbacks.shift();
+    if (!callback) {
+      throw new Error("ssh2 exec callback queue was not available for the no-reply workaround");
+    }
+    queueMicrotask(() => callback(false));
+  };
+  patchedSsh2Protocols.add(protocol);
+  return true;
+}
+
+function isNodeDevelopment(): boolean {
+  const nodeEnv = process.env.NODE_ENV?.trim().toLowerCase();
+  return nodeEnv === "development" || nodeEnv === "dev";
+}
+
+interface Ssh2ExecChannel {
+  outgoing: { id: number };
+  _callbacks: Array<(hadError: boolean | Error) => void>;
+}
+
+interface Ssh2Protocol {
+  exec(channelId: number, command: string | Buffer, wantReply?: boolean): void;
+}
+
+interface Ssh2ClientInternals {
+  _protocol?: Ssh2Protocol;
+  _chanMgr?: { _channels: Record<number, Ssh2ExecChannel | undefined> };
+}

package/src/key-store.ts

@@ -0,0 +1,177 @@
+import { createHash, randomBytes } from "node:crypto";
+import { chmod, mkdir, readFile, rename, rm, writeFile } from "node:fs/promises";
+import { homedir } from "node:os";
+import { join } from "node:path";
+import { fileLock } from "@dbx-tools/core";
+import ssh2 from "ssh2";
+import type { LakeboxClient } from "./rest.ts";
+
+const { utils } = ssh2;
+
+interface StoredKey {
+  privateKey: string;
+  privateKeyPath: string;
+  publicKey: string;
+  keyHash: string;
+}
+
+/**
+ * Stores caller-scoped SSH keys and trusted Databricks gateway host keys.
+ *
+ * Each workspace and caller identity receives a separate Ed25519 key pair. The
+ * store writes private keys, gateway fingerprints, and OpenSSH known-host files
+ * with owner-only permissions beneath the configured root directory.
+ */
+export class SshKeyStore {
+  /** Creates a key store rooted at the default dbx-tools application-data path. */
+  constructor(private readonly root = join(homedir(), ".dbx-tools", "lakebox")) {}
+
+  /**
+   * Loads or creates the caller key pair and registers its public key when the
+   * workspace does not already contain the matching key hash.
+   *
+   * @returns The local key material and the gateway host returned by Databricks,
+   * when available.
+   */
+  async ensureRegistered(
+    api: LakeboxClient,
+    owner: string,
+    signal?: AbortSignal,
+  ): Promise<StoredKey & { gatewayHost?: string }> {
+    const key = await this.loadOrCreate(await api.workspaceScope(), owner);
+    const keys = await api.listKeys(signal);
+    const existing = keys.find((entry) => entry.keyHash === key.keyHash);
+    if (existing) {
+      return { ...key, ...(existing.gatewayHost ? { gatewayHost: existing.gatewayHost } : {}) };
+    }
+    const registered = await api.registerKey(
+      key.publicKey,
+      `dbx-tools-lakebox:${key.keyHash}`,
+      signal,
+    );
+    return {
+      ...key,
+      ...(registered.gatewayHost ? { gatewayHost: registered.gatewayHost } : {}),
+    };
+  }
+
+  /** Returns the stored SHA-256 gateway fingerprint, or `undefined` when untrusted. */
+  async knownHost(
+    scope: string,
+    owner: string,
+    host: string,
+    port: number,
+  ): Promise<string | undefined> {
+    try {
+      const value = await readFile(this.hostFingerprintPath(scope, owner, host, port), "utf8");
+      return value.trim() || undefined;
+    } catch {
+      return undefined;
+    }
+  }
+
+  /** Returns the OpenSSH known-host file path for one workspace, user, and gateway. */
+  knownHostsPath(scope: string, owner: string, host: string, port: number): string {
+    return join(this.directory(scope, owner), `host-${digest(`${host}\0${port}`)}.known_hosts`);
+  }
+
+  /** Stores the accepted gateway fingerprint and OpenSSH host-key entry. */
+  async rememberHost(
+    scope: string,
+    owner: string,
+    host: string,
+    port: number,
+    fingerprint: string,
+    key: Buffer,
+  ): Promise<void> {
+    const directory = this.directory(scope, owner);
+    await mkdir(directory, { recursive: true, mode: 0o700 });
+    const fingerprintPath = this.hostFingerprintPath(scope, owner, host, port);
+    const knownHostsPath = this.knownHostsPath(scope, owner, host, port);
+    const knownHost = `[${host}]:${port} ${sshKeyAlgorithm(key)} ${key.toString("base64")}\n`;
+    await Promise.all([
+      writeFile(fingerprintPath, `${fingerprint}\n`, { mode: 0o600 }),
+      writeFile(knownHostsPath, knownHost, { mode: 0o600 }),
+    ]);
+    await Promise.all([chmod(fingerprintPath, 0o600), chmod(knownHostsPath, 0o600)]);
+  }
+
+  private async loadOrCreate(scope: string, owner: string): Promise<StoredKey> {
+    const directory = this.directory(scope, owner);
+    const privatePath = join(directory, "id_ed25519");
+    const publicPath = `${privatePath}.pub`;
+    const stored = await readStoredKey(privatePath, publicPath);
+    if (stored) return stored;
+    await mkdir(directory, { recursive: true, mode: 0o700 });
+    return fileLock.withFileLock(["lakebox-ssh-key", privatePath], async () => {
+      const rechecked = await readStoredKey(privatePath, publicPath);
+      if (rechecked) return rechecked;
+
+      const generated = utils.generateKeyPairSync("ed25519");
+      const privateKey = generated.private;
+      const publicKey = generated.public;
+      const suffix = `${process.pid}-${randomBytes(8).toString("hex")}`;
+      const stagedPrivate = `${privatePath}.${suffix}`;
+      const stagedPublic = `${publicPath}.${suffix}`;
+      try {
+        await writeFile(stagedPrivate, privateKey, { mode: 0o600 });
+        await writeFile(stagedPublic, `${publicKey.trim()}\n`, { mode: 0o600 });
+        await rename(stagedPrivate, privatePath);
+        await rename(stagedPublic, publicPath);
+        await Promise.all([chmod(privatePath, 0o600), chmod(publicPath, 0o600)]);
+        const created = await readStoredKey(privatePath, publicPath);
+        if (!created) throw new Error("SSH key files were not published atomically");
+        return created;
+      } finally {
+        await Promise.all([rm(stagedPrivate, { force: true }), rm(stagedPublic, { force: true })]);
+      }
+    });
+  }
+
+  private hostFingerprintPath(scope: string, owner: string, host: string, port: number): string {
+    return join(this.directory(scope, owner), `host-${digest(`${host}\0${port}`)}.sha256`);
+  }
+
+  private directory(scope: string, owner: string): string {
+    return join(this.root, digest(`${scope}\0${owner}`));
+  }
+}
+
+async function readStoredKey(
+  privatePath: string,
+  publicPath: string,
+): Promise<StoredKey | undefined> {
+  try {
+    const [privateKey, publicKey] = await Promise.all([
+      readFile(privatePath, "utf8"),
+      readFile(publicPath, "utf8"),
+    ]);
+    return {
+      privateKey,
+      privateKeyPath: privatePath,
+      publicKey,
+      keyHash: publicKeyHash(publicKey),
+    };
+  } catch {
+    return undefined;
+  }
+}
+
+function publicKeyHash(publicKey: string): string {
+  const normalized = publicKey.trim().split(/\s+/).slice(0, 2).join(" ");
+  return createHash("sha256").update(normalized).digest().subarray(0, 16).toString("hex");
+}
+
+function digest(value: string): string {
+  return createHash("sha256").update(value).digest("hex").slice(0, 32);
+}
+
+function sshKeyAlgorithm(key: Buffer): string {
+  if (key.length < 5) throw new Error("SSH gateway returned an invalid host key");
+  const length = key.readUInt32BE(0);
+  const end = 4 + length;
+  if (length === 0 || end > key.length) throw new Error("SSH gateway returned an invalid host key");
+  const algorithm = key.subarray(4, end).toString("ascii");
+  if (/\s/.test(algorithm)) throw new Error("SSH gateway returned an invalid host key algorithm");
+  return algorithm;
+}

package/src/lifecycle.ts

@@ -0,0 +1,88 @@
+import { asyncUtils } from "@dbx-tools/shared-core";
+import type { LakeboxStatus } from "@dbx-tools/shared-lakebox";
+import type { LakeboxClient, LakeboxEntry } from "./rest.ts";
+
+const POLL_INTERVAL_MS = 1_000;
+const LIFECYCLE_TIMEOUT_MS = 300_000;
+
+/** Maps an upstream Lakebox status to the shared lifecycle state. */
+export function normalizeLakeboxStatus(value: string): Exclude<LakeboxStatus, "missing"> {
+  switch (value.trim().toLowerCase()) {
+    case "starting":
+    case "provisioning":
+    case "pending":
+      return "starting";
+    case "running":
+      return "running";
+    case "stopping":
+      return "stopping";
+    case "stopped":
+      return "stopped";
+    case "failed":
+    case "error":
+      return "error";
+    default:
+      return "unknown";
+  }
+}
+
+/** Parses an upstream duration such as `600s` into seconds. */
+export function lakeboxIdleTimeoutSeconds(value: string | undefined): number | undefined {
+  if (!value?.endsWith("s")) return undefined;
+  const seconds = Number(value.slice(0, -1));
+  return Number.isFinite(seconds) && seconds > 0 ? seconds : undefined;
+}
+
+/**
+ * Starts a Lakebox when needed and waits until it is running.
+ *
+ * @throws When the Lakebox does not reach running within five minutes.
+ */
+export async function ensureLakeboxRunning(
+  api: LakeboxClient,
+  current: LakeboxEntry,
+  signal?: AbortSignal,
+): Promise<LakeboxEntry> {
+  let entry = current;
+  const status = normalizeLakeboxStatus(entry.status);
+  if (status === "running") return entry;
+  if (status === "starting") return waitForLakeboxStatus(api, entry, "running", signal);
+  if (status === "stopping") entry = await waitForLakeboxStatus(api, entry, "stopped", signal);
+  return waitForLakeboxStatus(api, await api.start(entry.sandboxId, signal), "running", signal);
+}
+
+/**
+ * Stops a Lakebox when needed and waits until it is stopped.
+ *
+ * @throws When the Lakebox does not reach stopped within five minutes.
+ */
+export async function ensureLakeboxStopped(
+  api: LakeboxClient,
+  current: LakeboxEntry,
+  signal?: AbortSignal,
+): Promise<LakeboxEntry> {
+  let entry = current;
+  const status = normalizeLakeboxStatus(entry.status);
+  if (status === "stopped") return entry;
+  if (status === "stopping") return waitForLakeboxStatus(api, entry, "stopped", signal);
+  if (status === "starting") entry = await waitForLakeboxStatus(api, entry, "running", signal);
+  return waitForLakeboxStatus(api, await api.stop(entry.sandboxId, signal), "stopped", signal);
+}
+
+async function waitForLakeboxStatus(
+  api: LakeboxClient,
+  current: LakeboxEntry,
+  expected: "running" | "stopped",
+  signal?: AbortSignal,
+): Promise<LakeboxEntry> {
+  const deadline = Date.now() + LIFECYCLE_TIMEOUT_MS;
+  let entry = current;
+  while (normalizeLakeboxStatus(entry.status) !== expected) {
+    if (Date.now() >= deadline) {
+      throw new Error(`Timed out waiting for Lakebox ${entry.sandboxId} to become ${expected}`);
+    }
+    await asyncUtils.sleep(POLL_INTERVAL_MS, signal);
+    entry = await api.get(entry.sandboxId, signal);
+  }
+  return entry;
+}