promigence 0.1.0-beta.12

package/dist/sdk/pty.js

@@ -0,0 +1,496 @@
+import { createHash, randomBytes } from "node:crypto";
+import { CALLER_HEADER } from "./caller.js";
+import { requestVia } from "./proxy.js";
+export const PTY_INPUT_CHUNK = 64 * 1024;
+export const PTY_INPUT_KEEP = 256 * 1024;
+export const PTY_RESUME_HEADER = "promigence-resume";
+const WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
+export const PTY_COMMAND_HEADER = "promigence-pty-command";
+export const PTY_SECRETS_HEADER = "promigence-pty-secrets";
+export function wsAccept(key) {
+    return createHash("sha1").update(key + WS_GUID).digest("base64");
+}
+export function encodeFrame(opcode, payload) {
+    const len = payload.length;
+    const head = [0x80 | opcode];
+    if (len < 126)
+        head.push(0x80 | len);
+    else if (len < 65536)
+        head.push(0x80 | 126, len >> 8, len & 0xff);
+    else {
+        head.push(0x80 | 127);
+        const big = Buffer.alloc(8);
+        big.writeBigUInt64BE(BigInt(len));
+        head.push(...big);
+    }
+    const mask = randomBytes(4);
+    const out = Buffer.alloc(head.length + 4 + len);
+    Buffer.from(head).copy(out, 0);
+    mask.copy(out, head.length);
+    const at = head.length + 4;
+    for (let i = 0; i < len; i++)
+        out[at + i] = payload[i] ^ mask[i % 4];
+    return out;
+}
+export class FrameReader {
+    buf = Buffer.alloc(0);
+    partial;
+    push(chunk) {
+        this.buf = this.buf.length ? Buffer.concat([this.buf, chunk]) : Buffer.from(chunk);
+        const out = [];
+        for (;;) {
+            if (this.buf.length < 2)
+                break;
+            const b0 = this.buf[0];
+            const b1 = this.buf[1];
+            let len = b1 & 0x7f;
+            let at = 2;
+            if (len === 126) {
+                if (this.buf.length < 4)
+                    break;
+                len = this.buf.readUInt16BE(2);
+                at = 4;
+            }
+            else if (len === 127) {
+                if (this.buf.length < 10)
+                    break;
+                len = Number(this.buf.readBigUInt64BE(2));
+                at = 10;
+            }
+            const masked = (b1 & 0x80) !== 0;
+            if (this.buf.length < at + (masked ? 4 : 0) + len)
+                break;
+            let payload = this.buf.subarray(at + (masked ? 4 : 0), at + (masked ? 4 : 0) + len);
+            if (masked) {
+                const key = this.buf.subarray(at, at + 4);
+                payload = Buffer.from(payload.map((b, i) => b ^ key[i % 4]));
+            }
+            else {
+                payload = Buffer.from(payload);
+            }
+            this.buf = this.buf.subarray(at + (masked ? 4 : 0) + len);
+            const fin = (b0 & 0x80) !== 0;
+            const opcode = b0 & 0x0f;
+            if (opcode >= 0x8) {
+                out.push({ opcode, payload });
+                continue;
+            }
+            if (opcode !== 0)
+                this.partial = { opcode, parts: [payload] };
+            else
+                this.partial?.parts.push(payload);
+            if (fin && this.partial) {
+                out.push({ opcode: this.partial.opcode, payload: Buffer.concat(this.partial.parts) });
+                this.partial = undefined;
+            }
+        }
+        return out;
+    }
+}
+export class PtySession {
+    socket;
+    pid;
+    cols;
+    rows;
+    closed = false;
+    done;
+    pending;
+    resumeToken;
+    inputSent = 0;
+    inputBase = 0;
+    inputLog = Buffer.alloc(0);
+    constructor(socket, open, done) {
+        this.socket = socket;
+        this.pid = open.pid;
+        this.cols = open.cols;
+        this.rows = open.rows;
+        this.resumeToken = open.resume;
+        this.done = done;
+        void done.then(() => { this.#ended = true; });
+    }
+    #ended = false;
+    get ended() {
+        return this.#ended;
+    }
+    holdInput() {
+        this.pending ??= [];
+    }
+    rebind(socket, open) {
+        this.socket = socket;
+        this.resumeToken = open.resume;
+        this.inputSent = 0;
+        this.inputBase = 0;
+        this.inputLog = Buffer.alloc(0);
+        this.flush(open);
+    }
+    resumed(socket, open) {
+        this.socket = socket;
+        this.resumeToken = open.resume ?? this.resumeToken;
+        const took = open.input ?? this.inputSent;
+        if (took < this.inputSent) {
+            const from = Math.max(took, this.inputBase) - this.inputBase;
+            const again = this.inputLog.subarray(from);
+            this.inputSent = Math.max(took, this.inputBase);
+            for (let at = 0; at < again.length; at += PTY_INPUT_CHUNK) {
+                this.writeNow(0x2, again.subarray(at, at + PTY_INPUT_CHUNK));
+            }
+        }
+        this.flush(open);
+    }
+    flush(open) {
+        const queued = this.pending ?? [];
+        this.pending = undefined;
+        if (open.cols !== this.cols || open.rows !== this.rows) {
+            this.resize({ cols: this.cols, rows: this.rows });
+        }
+        for (const f of queued)
+            this.writeNow(f.opcode, f.payload);
+    }
+    writeNow(opcode, payload) {
+        if (this.closed || this.socket.destroyed)
+            return;
+        this.socket.write(encodeFrame(opcode, payload));
+        if (opcode === 0x2) {
+            this.inputSent += payload.length;
+            this.inputLog = Buffer.concat([this.inputLog, payload]);
+            if (this.inputLog.length > PTY_INPUT_KEEP) {
+                const drop = this.inputLog.length - PTY_INPUT_KEEP;
+                this.inputLog = this.inputLog.subarray(drop);
+                this.inputBase += drop;
+            }
+        }
+    }
+    get hungUp() {
+        return this.closed;
+    }
+    send(opcode, payload) {
+        if (this.closed)
+            return;
+        if (this.pending) {
+            if (this.pending.reduce((n, f) => n + f.payload.length, 0) + payload.length <= 1 << 20) {
+                this.pending.push({ opcode, payload: Buffer.from(payload) });
+            }
+            return;
+        }
+        this.writeNow(opcode, payload);
+    }
+    write(data) {
+        const bytes = typeof data === "string" ? Buffer.from(data, "utf8") : data;
+        for (let at = 0; at < bytes.length || at === 0; at += PTY_INPUT_CHUNK) {
+            this.send(0x2, bytes.subarray(at, at + PTY_INPUT_CHUNK));
+            if (bytes.length === 0)
+                break;
+        }
+    }
+    resize(size) {
+        this.cols = size.cols;
+        this.rows = size.rows;
+        this.send(0x1, Buffer.from(JSON.stringify({ resize: { cols: size.cols, rows: size.rows } })));
+    }
+    signal(sig) {
+        this.send(0x1, Buffer.from(JSON.stringify({ signal: sig })));
+    }
+    eof() {
+        this.send(0x1, Buffer.from(JSON.stringify({ eof: true })));
+    }
+    wait() {
+        return this.done;
+    }
+    disconnect() {
+        if (this.closed)
+            return;
+        this.pending = undefined;
+        this.send(0x8, Buffer.from([0x03, 0xe8]));
+        this.closed = true;
+        this.socket.end();
+    }
+}
+export const PTY_RECONNECT_CODE = 1012;
+export const PTY_RECONNECT_FOR_MS = 30_000;
+export const PTY_RECONNECT_FIRST_WAIT_MS = 250;
+export const PTY_RECONNECT_MAX_WAIT_MS = 4_000;
+export const PTY_RESTART_EXIT = {
+    exitCode: null,
+    signal: null,
+    error: {
+        code: "service_restart",
+        message: "the service restarted and the terminal could not reconnect within 30 s; open it again",
+    },
+};
+export const PTY_RESTART_CLOSED_EXIT = {
+    exitCode: null,
+    signal: null,
+    error: { code: "service_restart", message: "the service restarted and closed the terminal; open it again" },
+};
+const CLOSED_EXIT = {
+    exitCode: null,
+    signal: null,
+    error: { code: "closed", message: "the terminal connection closed without an exit record" },
+};
+class LegRefused {
+    status;
+    error;
+    restart;
+    constructor(status, error, restart = false) {
+        this.status = status;
+        this.error = error;
+        this.restart = restart;
+    }
+}
+function openLeg(opts, url, resume) {
+    const key = randomBytes(16).toString("base64");
+    const req = requestVia(url, {
+        method: "GET",
+        agent: false,
+        headers: {
+            connection: "Upgrade",
+            upgrade: "websocket",
+            "sec-websocket-version": "13",
+            "sec-websocket-key": key,
+            authorization: `Bearer ${opts.apiKey}`,
+            "user-agent": opts.userAgent,
+            ...(opts.caller ? { [CALLER_HEADER]: opts.caller } : {}),
+            ...(resume ? { [PTY_RESUME_HEADER]: resume } : {}),
+            ...(opts.command && !resume
+                ? { [PTY_COMMAND_HEADER]: Buffer.from(opts.command, "utf8").toString("base64") }
+                : {}),
+            ...(opts.secrets?.length && !resume
+                ? {
+                    [PTY_SECRETS_HEADER]: Buffer.from(JSON.stringify(opts.secrets), "utf8").toString("base64"),
+                }
+                : {}),
+        },
+    });
+    let resolveEnded;
+    const ended = new Promise((r) => { resolveEnded = r; });
+    const opened = new Promise((resolve, reject) => {
+        let settled = false;
+        const fail = (status, body, retryAfter, restart = false) => {
+            if (settled)
+                return;
+            settled = true;
+            reject(new LegRefused(status, opts.refuse(status, body, retryAfter), restart));
+        };
+        const timer = setTimeout(() => {
+            req.destroy();
+            fail(504, JSON.stringify({ error: { code: "timeout", message: `terminal handshake to ${url.host} took over ${opts.timeoutMs ?? 30_000} ms`, fix: "Retry; if it repeats, check the sandbox is ready (`promigence run report <run_id>`)." } }));
+        }, opts.timeoutMs ?? 30_000);
+        req.on("error", (e) => { clearTimeout(timer); fail(0, JSON.stringify({ error: { code: "unreachable", message: `cannot reach ${url.host}: ${e.code ?? e.message}` } })); });
+        req.on("response", (res) => {
+            clearTimeout(timer);
+            const parts = [];
+            res.on("data", (d) => parts.push(d));
+            const retryAfter = res.headers["retry-after"];
+            res.on("end", () => fail(res.statusCode ?? 0, Buffer.concat(parts).toString("utf8"), typeof retryAfter === "string" ? retryAfter : undefined));
+        });
+        req.on("upgrade", (res, socket, head) => {
+            clearTimeout(timer);
+            if (res.headers["sec-websocket-accept"] !== wsAccept(key)) {
+                socket.destroy();
+                fail(502, JSON.stringify({ error: { code: "internal", message: "the terminal server answered the WebSocket handshake with the wrong Sec-WebSocket-Accept" } }));
+                return;
+            }
+            socket.setNoDelay(true);
+            const reader = new FrameReader();
+            let exit;
+            let closeCode;
+            let wasOpen = false;
+            const onFrames = (chunk) => {
+                for (const f of reader.push(chunk)) {
+                    if (f.opcode === 0x2) {
+                        opts.onData?.(new Uint8Array(f.payload));
+                    }
+                    else if (f.opcode === 0x1) {
+                        let v = {};
+                        try {
+                            v = JSON.parse(f.payload.toString("utf8"));
+                        }
+                        catch {
+                            continue;
+                        }
+                        if (v["event"] === "open" && !wasOpen) {
+                            wasOpen = true;
+                            settled = true;
+                            resolve({
+                                socket,
+                                pid: typeof v["req_id"] === "number" ? v["req_id"] : 0,
+                                cols: Number(v["cols"] ?? 80),
+                                rows: Number(v["rows"] ?? 24),
+                                ...(typeof v["resume"] === "string" ? { resume: v["resume"] } : {}),
+                                ...(v["reattached"] === true ? { reattached: true } : {}),
+                                ...(typeof v["input"] === "number" ? { input: v["input"] } : {}),
+                            });
+                        }
+                        else if (v["event"] === "exit") {
+                            exit = {
+                                exitCode: typeof v["exit_code"] === "number" ? v["exit_code"] : null,
+                                signal: typeof v["signal"] === "number" ? v["signal"] : null,
+                                ...(typeof v["wall_ms"] === "number" ? { wallMs: v["wall_ms"] } : {}),
+                            };
+                        }
+                        else if (v["event"] === "error" || v["event"] === "closed") {
+                            exit = {
+                                exitCode: null, signal: null,
+                                error: { code: String(v["code"] ?? v["event"]), message: String(v["msg"] ?? v["message"] ?? "the terminal ended") },
+                            };
+                        }
+                    }
+                    else if (f.opcode === 0x9) {
+                        socket.write(encodeFrame(0xa, f.payload));
+                    }
+                    else if (f.opcode === 0x8) {
+                        if (closeCode === undefined)
+                            closeCode = f.payload.length >= 2 ? f.payload.readUInt16BE(0) : 1005;
+                        if (!socket.destroyed)
+                            socket.end(encodeFrame(0x8, f.payload.subarray(0, 2)));
+                    }
+                }
+            };
+            socket.on("data", onFrames);
+            socket.on("close", () => {
+                if (!wasOpen) {
+                    const restart = closeCode === PTY_RECONNECT_CODE && exit === undefined;
+                    const e = restart ? PTY_RESTART_CLOSED_EXIT : (exit ?? CLOSED_EXIT);
+                    fail(restart ? 503 : 502, JSON.stringify({ error: { code: e.error?.code ?? "internal", message: e.error?.message ?? "the terminal ended before it opened" } }), undefined, restart);
+                    return;
+                }
+                resolveEnded({ code: closeCode, exit });
+            });
+            socket.on("error", () => { });
+            if (head.length)
+                onFrames(head);
+        });
+        req.end();
+    });
+    return { opened, ended };
+}
+function sizedUrl(url, size) {
+    const u = new URL(url);
+    u.searchParams.set("cols", String(Math.floor(size.cols)));
+    u.searchParams.set("rows", String(Math.floor(size.rows)));
+    return u;
+}
+function resumeUrl(url, size, received) {
+    const u = sizedUrl(url, size);
+    u.searchParams.set("received", String(received));
+    return u;
+}
+function worthRetrying(status) {
+    return status === 0 || status === 429 || status >= 500;
+}
+export function connectPty(opts) {
+    const reconnect = opts.reconnect ?? opts.command === undefined;
+    const resumeOn = opts.resume ?? true;
+    let received = 0;
+    const legOpts = {
+        ...opts,
+        onData: (d) => {
+            received += d.length;
+            opts.onData?.(d);
+        },
+    };
+    let resolveDone;
+    const done = new Promise((r) => { resolveDone = r; });
+    let finished = false;
+    const finish = (e) => {
+        if (finished)
+            return;
+        finished = true;
+        opts.onExit?.(e);
+        resolveDone(e);
+    };
+    const sleep = (ms) => new Promise((r) => setTimeout(r, ms));
+    const follow = (session, leg) => {
+        void leg.ended.then(async ({ code, exit }) => {
+            if (exit !== undefined || session.hungUp) {
+                finish(exit ?? CLOSED_EXIT);
+                return;
+            }
+            const restart = code === PTY_RECONNECT_CODE;
+            const dropped = code === undefined || code === 1006;
+            let resumable = resumeOn && session.resumeToken !== undefined && (restart || dropped);
+            if (!restart && !resumable) {
+                finish(CLOSED_EXIT);
+                return;
+            }
+            if (!resumable && !reconnect) {
+                finish(PTY_RESTART_CLOSED_EXIT);
+                return;
+            }
+            session.holdInput();
+            const t0 = Date.now();
+            for (let attempt = 0;; attempt++) {
+                const wait = Math.min(PTY_RECONNECT_FIRST_WAIT_MS * 2 ** attempt, PTY_RECONNECT_MAX_WAIT_MS) * (0.5 + Math.random() / 2);
+                if (Date.now() - t0 + wait > (opts.reconnectForMs ?? PTY_RECONNECT_FOR_MS))
+                    break;
+                await sleep(wait);
+                if (session.hungUp) {
+                    finish(CLOSED_EXIT);
+                    return;
+                }
+                const next = resumable
+                    ? openLeg(legOpts, resumeUrl(opts.url, session, received), session.resumeToken)
+                    : openLeg(legOpts, sizedUrl(opts.url, session));
+                try {
+                    const o = await next.opened;
+                    if (session.hungUp) {
+                        o.socket.destroy();
+                        finish(CLOSED_EXIT);
+                        return;
+                    }
+                    if (resumable) {
+                        session.resumed(o.socket, o);
+                    }
+                    else {
+                        received = 0;
+                        session.rebind(o.socket, o);
+                    }
+                    follow(session, next);
+                    opts.onReconnected?.({ resumed: resumable });
+                    return;
+                }
+                catch (e) {
+                    const refused = e instanceof LegRefused ? e : undefined;
+                    if (resumable && refused?.status === 404) {
+                        resumable = false;
+                        if (!restart) {
+                            finish(CLOSED_EXIT);
+                            return;
+                        }
+                        if (!reconnect) {
+                            finish(PTY_RESTART_CLOSED_EXIT);
+                            return;
+                        }
+                        continue;
+                    }
+                    if (refused && !worthRetrying(refused.status)) {
+                        const err = refused.error;
+                        finish({ exitCode: null, signal: null, error: { code: typeof err.code === "string" ? err.code : "internal", message: err.message } });
+                        return;
+                    }
+                }
+            }
+            finish(restart ? PTY_RESTART_EXIT : CLOSED_EXIT);
+        });
+    };
+    return (async () => {
+        const t0 = Date.now();
+        for (let attempt = 0;; attempt++) {
+            const first = openLeg(legOpts, new URL(opts.url));
+            try {
+                const o = await first.opened;
+                const session = new PtySession(o.socket, o, done);
+                follow(session, first);
+                return session;
+            }
+            catch (e) {
+                const wait = Math.min(PTY_RECONNECT_FIRST_WAIT_MS * 2 ** attempt, PTY_RECONNECT_MAX_WAIT_MS) * (0.5 + Math.random() / 2);
+                if (e instanceof LegRefused && e.restart && reconnect && Date.now() - t0 + wait <= (opts.reconnectForMs ?? PTY_RECONNECT_FOR_MS)) {
+                    await sleep(wait);
+                    continue;
+                }
+                throw e instanceof LegRefused ? e.error : e;
+            }
+        }
+    })();
+}