@prismnetwork/agent-sdk 0.5.0 → 0.6.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@prismnetwork/agent-sdk",
-  "version": "0.5.0",
+  "version": "0.6.0",
   "description": "Headless GPU leasing and renter-encrypted storage on Prism Network for wallet-holding agents.",
   "type": "module",
   "main": "prism.mjs",
@@ -24,6 +24,7 @@
   },
   "files": [
     "prism.mjs",
+    "relay.mjs",
     "toolset.d.mts",
     "prism.d.mts",
     "toolset.mjs",

package/prism.d.mts

@@ -3,15 +3,33 @@ export declare const USDG: string;
 export declare const DEFAULT_IMAGE: string;
 export declare const TRUST_CLASSES: readonly ["open", "isolated", "attested", "confidential"];
 
+/// `mode` says which of the two shapes arrived. Brokered capacity fills in
+/// `ssh_host` and `ssh_port`; a node that accepts nothing inbound fills in the
+/// gateway fields instead and is reached through a relay.
 export interface LeaseAccess {
-  mode?: string;
+  mode?: "direct_ssh" | "gateway" | string;
   ssh_host?: string;
   ssh_port?: number;
   ssh_user?: string;
+  gateway_host?: string;
+  relay_port?: number;
+  /// The root the relay's certificate chains to, in PEM. It is served under a
+  /// private CA, so this is what the client pins.
+  gateway_ca?: string;
+  token?: string;
+  jupyter_path?: string;
+  jupyter_token?: string;
   expires_at?: string;
   [key: string]: unknown;
 }
 
+/// A local address that forwards to the workspace until it is closed.
+export interface RelayForwarder {
+  host: string;
+  port: number;
+  close(): Promise<void>;
+}
+
 export interface LeaseHandle {
   leaseId: number;
   access: LeaseAccess;
@@ -74,6 +92,9 @@ export declare class PrismAgent {
     command: string,
     options?: { timeoutMs?: number; connectRetries?: number; connectDelayMs?: number; stdin?: string | null },
   ): Promise<RunResult>;
+  /// Only for a lease reached through the gateway. Use it for anything that is
+  /// not a one-shot command: scp, a notebook client, an interactive shell.
+  forward(lease: LeaseHandle, options?: { service?: "ssh" | "jupyter" }): Promise<RelayForwarder>;
   endLease(lease: LeaseHandle): void;
 }
 

package/prism.mjs

@@ -14,6 +14,7 @@ import {
   stringToBytes,
 } from "viem";
 import { privateKeyToAccount } from "viem/accounts";
+import { openRelayForwarder } from "./relay.mjs";
 import { PrismVault } from "./vault.mjs";
 import { PrismWorkspace } from "./workspace.mjs";
 
@@ -419,28 +420,66 @@ export class PrismAgent {
   // can lag a few minutes after the box reports ready. `stdin` feeds the command
   // its input, which keeps anything sensitive out of the remote process table.
   async run(lease, command, { timeoutMs = 120_000, connectRetries = 24, connectDelayMs = 10_000, stdin = null } = {}) {
-    if (!lease?.access?.ssh_host || !lease.access.ssh_port || !lease.keyPath) {
+    if (typeof command !== "string" || command.length === 0) throw new PrismError(400, "command_required");
+    if (!lease?.keyPath) {
       throw new PrismError(400, "invalid_lease_handle", {
         mode: lease?.access?.mode ?? null,
         lease_id: lease?.leaseId ?? null,
-        hint: "gateway-mode access has no ssh endpoint",
+        hint: "the lease handle carries no ssh key",
       });
     }
-    if (typeof command !== "string" || command.length === 0) throw new PrismError(400, "command_required");
-    const target = {
-      host: lease.access.ssh_host,
-      port: lease.access.ssh_port,
-      user: lease.access.ssh_user ?? "root",
-      keyPath: lease.keyPath,
-    };
-    let last;
-    for (let attempt = 0; attempt <= connectRetries; attempt++) {
-      const res = await this.#ssh(target, command, timeoutMs, stdin);
-      if (!isSshWarmup(res)) return res;
-      last = res;
-      if (attempt < connectRetries) await sleep(connectDelayMs);
+
+    // A physical node accepts nothing inbound, so its session arrives through
+    // the gateway. Opening the renter's half of that tunnel gives a local port
+    // that behaves like any other host, which is why the retry loop below does
+    // not care which kind of capacity it is talking to.
+    const forwarder =
+      lease.access?.mode === "gateway" ? await openRelayForwarder(lease.access) : null;
+    try {
+      const target = forwarder
+        ? {
+            host: forwarder.host,
+            port: forwarder.port,
+            user: lease.access.ssh_user ?? "workspace",
+            keyPath: lease.keyPath,
+          }
+        : {
+            host: lease.access?.ssh_host,
+            port: lease.access?.ssh_port,
+            user: lease.access?.ssh_user ?? "root",
+            keyPath: lease.keyPath,
+          };
+      if (!target.host || !target.port) {
+        throw new PrismError(400, "invalid_lease_handle", {
+          mode: lease.access?.mode ?? null,
+          lease_id: lease.leaseId ?? null,
+          hint: "the access grant names no reachable endpoint",
+        });
+      }
+      let last;
+      for (let attempt = 0; attempt <= connectRetries; attempt++) {
+        const res = await this.#ssh(target, command, timeoutMs, stdin);
+        if (!isSshWarmup(res)) return res;
+        last = res;
+        if (attempt < connectRetries) await sleep(connectDelayMs);
+      }
+      return last;
+    } finally {
+      if (forwarder) await forwarder.close();
+    }
+  }
+
+  /// A local address that forwards to the workspace for as long as it is open.
+  /// Use it for anything that is not a one-shot command: `scp`, port forwards,
+  /// an interactive shell, a notebook client. The caller closes it.
+  async forward(lease, { service = "ssh" } = {}) {
+    if (lease?.access?.mode !== "gateway") {
+      throw new PrismError(400, "forward_not_supported", {
+        mode: lease?.access?.mode ?? null,
+        hint: "this lease is reachable directly and needs no relay",
+      });
     }
-    return last;
+    return openRelayForwarder(lease.access, { service });
   }
 
   // Releases local key material. The on-chain lease settles at the end of its duration.

package/relay.mjs

@@ -0,0 +1,176 @@
+// Reaching a workspace that has no public address.
+//
+// Capacity brokered from a cloud gives the renter an SSH endpoint on the host.
+// A physical node has none: it dials out to the gateway and accepts nothing
+// inbound, so the renter's session is carried back through that tunnel. This
+// opens the renter's half of it and presents the result as a local port, which
+// is what lets `ssh`, `scp`, or anything else speak to a machine that cannot be
+// addressed.
+//
+// The relay wants one JSON frame naming the grant and the service, answers with
+// one saying whether it paired, and from then on the connection is the workspace
+// socket. Frames are a big-endian u32 length followed by the payload.
+import { createServer } from "node:net";
+import { connect as tlsConnect } from "node:tls";
+
+const MAX_FRAME_BYTES = 16 * 1_024;
+const HANDSHAKE_TIMEOUT_MS = 20_000;
+
+export class RelayError extends Error {
+  constructor(code, detail) {
+    super(code);
+    this.name = "RelayError";
+    this.code = code;
+    this.detail = detail ?? null;
+  }
+}
+
+function frame(value) {
+  const payload = Buffer.from(JSON.stringify(value));
+  if (payload.length > MAX_FRAME_BYTES) throw new RelayError("relay_frame_too_large");
+  const header = Buffer.alloc(4);
+  header.writeUInt32BE(payload.length, 0);
+  return Buffer.concat([header, payload]);
+}
+
+// Resolves with the first frame and whatever bytes arrived behind it. Those
+// trailing bytes are already workspace traffic, so losing them corrupts the
+// session before it starts.
+function readFrame(socket) {
+  return new Promise((resolve, reject) => {
+    let buffer = Buffer.alloc(0);
+    const timer = setTimeout(() => {
+      cleanup();
+      reject(new RelayError("relay_handshake_timeout"));
+    }, HANDSHAKE_TIMEOUT_MS);
+
+    const onData = (chunk) => {
+      buffer = Buffer.concat([buffer, chunk]);
+      if (buffer.length < 4) return;
+      const length = buffer.readUInt32BE(0);
+      if (length > MAX_FRAME_BYTES) {
+        cleanup();
+        reject(new RelayError("relay_frame_too_large"));
+        return;
+      }
+      if (buffer.length < 4 + length) return;
+      cleanup();
+      try {
+        resolve({
+          message: JSON.parse(buffer.subarray(4, 4 + length).toString("utf8")),
+          rest: buffer.subarray(4 + length),
+        });
+      } catch (err) {
+        reject(new RelayError("relay_frame_malformed", err?.message ?? String(err)));
+      }
+    };
+    const onError = (err) => {
+      cleanup();
+      reject(new RelayError("relay_disconnected", err?.message ?? String(err)));
+    };
+    const onEnd = () => {
+      cleanup();
+      reject(new RelayError("relay_closed_early"));
+    };
+    function cleanup() {
+      clearTimeout(timer);
+      socket.off("data", onData);
+      socket.off("error", onError);
+      socket.off("end", onEnd);
+    }
+
+    socket.on("data", onData);
+    socket.on("error", onError);
+    socket.on("end", onEnd);
+  });
+}
+
+function dial(access) {
+  return new Promise((resolve, reject) => {
+    const socket = tlsConnect(
+      {
+        host: access.gateway_host,
+        port: access.relay_port,
+        servername: access.gateway_host,
+        // The relay runs under a private CA, so the public trust store says
+        // nothing about it. Pinning the root the control plane handed back is
+        // the whole reason it is in the grant.
+        ca: access.gateway_ca ? [access.gateway_ca] : undefined,
+      },
+      () => resolve(socket),
+    );
+    socket.once("error", (err) => reject(new RelayError("relay_connect_failed", err?.message ?? String(err))));
+  });
+}
+
+// Checked before anything is opened, so a grant that can never work says so at
+// once instead of leaving the caller with a port that resets every connection.
+function assertUsable(access) {
+  if (!access?.gateway_host || !access?.relay_port || !access?.token) {
+    throw new RelayError("relay_access_incomplete");
+  }
+  if (!access.gateway_ca) {
+    throw new RelayError("relay_ca_missing", "the access grant carries no gateway root to verify against");
+  }
+}
+
+// One relay connection, paired and ready to carry traffic.
+export async function openRelayConnection(access, service = "ssh") {
+  assertUsable(access);
+  const socket = await dial(access);
+  socket.write(frame({ token: access.token, service }));
+  const { message, rest } = await readFrame(socket);
+  if (!message?.ready) {
+    socket.destroy();
+    throw new RelayError("relay_refused", message?.error ?? null);
+  }
+  return { socket, rest };
+}
+
+/// A local port that forwards to the workspace, one relay connection per
+/// inbound connection. `ssh` gets an address it can use and never learns the
+/// session is tunnelled.
+export async function openRelayForwarder(access, { service = "ssh" } = {}) {
+  assertUsable(access);
+  const server = createServer();
+  const sockets = new Set();
+
+  server.on("connection", (local) => {
+    sockets.add(local);
+    local.on("close", () => sockets.delete(local));
+    local.on("error", () => local.destroy());
+    openRelayConnection(access, service)
+      .then(({ socket, rest }) => {
+        if (local.destroyed) {
+          socket.destroy();
+          return;
+        }
+        sockets.add(socket);
+        socket.on("close", () => sockets.delete(socket));
+        socket.on("error", () => {
+          socket.destroy();
+          local.destroy();
+        });
+        if (rest.length > 0) local.write(rest);
+        local.pipe(socket);
+        socket.pipe(local);
+      })
+      .catch(() => local.destroy());
+  });
+
+  await new Promise((resolve, reject) => {
+    server.once("error", reject);
+    server.listen(0, "127.0.0.1", resolve);
+  });
+
+  const { port } = server.address();
+  return {
+    host: "127.0.0.1",
+    port,
+    async close() {
+      for (const socket of sockets) socket.destroy();
+      sockets.clear();
+      await new Promise((resolve) => server.close(resolve));
+    },
+  };
+}