fix: stream ssh downloads over tcp-pool

This commit is contained in:
Codex
2026-06-11 16:26:35 +00:00
parent 3d0faf557e
commit 6bce211b11
6 changed files with 473 additions and 154 deletions
+135 -148
View File
@@ -1,10 +1,18 @@
import { createHash, randomBytes } from "node:crypto";
import { mkdir, readFile, writeFile } from "node:fs/promises";
import { once } from "node:events";
import { createWriteStream } from "node:fs";
import { mkdir, readFile, rm } from "node:fs/promises";
import path from "node:path";
import type { ParsedSshInvocation, ParsedSshRoute, SshCaptureResult } from "./ssh";
export interface SshRemoteCommandStreamHandlers {
onStdout(chunk: Buffer): Promise<void> | void;
onStderr?(chunk: Buffer): Promise<void> | void;
}
export interface SshRemoteCommandExecutor {
runRemoteCommand(remoteCommand: string, input?: string): Promise<SshCaptureResult>;
streamRemoteCommand?(remoteCommand: string, handlers: SshRemoteCommandStreamHandlers, input?: string): Promise<SshCaptureResult>;
}
export interface SshFileTransferCommandBuilders {
@@ -14,7 +22,6 @@ export interface SshFileTransferCommandBuilders {
type SshFileTransferOperation =
| "stat"
| "read-b64-block"
| "write-b64-argv"
| "write-b64-stdin"
| "write-b64-begin"
@@ -25,7 +32,6 @@ interface SshFileTransferCliOptions {
action: "upload" | "download";
localPath: string;
remotePath: string;
readBlockBytes: number;
}
interface SshFileTransferStat {
@@ -43,6 +49,17 @@ interface SshFileTransferWriteResult {
chunks: number;
}
interface SshFileTransferDownloadResult {
remote: SshFileTransferStat;
local: SshFileTransferStat;
strategy: "tcp-pool-stdout-stream";
transport: "host.ssh.tcp-pool";
chunks: number;
elapsedMs: number;
throughputBytesPerSecond: number;
remainingBytes: number;
}
class SshFileTransferError extends Error {
constructor(message: string, public readonly details: Record<string, unknown> = {}) {
super(message);
@@ -50,12 +67,9 @@ class SshFileTransferError extends Error {
}
}
const fileTransferReadBlockBytes = 1_048_576;
const fileTransferWriteB64ArgvLimit = 48_000;
const fileTransferWriteRawChunkBytes = 1_048_576;
const fileTransferWriteB64ChunkChars = 1_398_104;
const fileTransferReadBlockMaxAttempts = 12;
const fileTransferReadEmptyRetryDelayMs = 250;
const fileTransferProgressEveryChunks = 16;
export function isSshFileTransferOperation(args: string[]): boolean {
@@ -97,15 +111,10 @@ export async function runSshFileTransferOperation(
return 0;
}
const read = await readRemoteFileVerified(invocation, executor, builders, options.remotePath, options.readBlockBytes);
await mkdir(path.dirname(localPath), { recursive: true });
await writeFile(localPath, read.content);
const local = await readFile(localPath);
const localStat = { bytes: local.length, sha256: sha256HexBuffer(local) };
assertTransferStat("download final local verification", localPath, read.remote, localStat);
const read = await downloadRemoteFileVerified(invocation, executor, builders, options.remotePath, localPath);
const verification = buildTransferVerification(
{ side: "remote", path: options.remotePath, ...read.remote },
{ side: "local", path: localPath, ...localStat },
{ side: "local", path: localPath, ...read.local },
);
process.stdout.write(`${JSON.stringify({
ok: true,
@@ -119,9 +128,12 @@ export async function runSshFileTransferOperation(
verified: true,
verification,
transfer: {
strategy: "chunked-read",
strategy: read.strategy,
transport: read.transport,
chunks: read.chunks,
chunkBytes: options.readBlockBytes,
elapsedMs: read.elapsedMs,
throughputBytesPerSecond: read.throughputBytesPerSecond,
remainingBytes: read.remainingBytes,
},
}, null, 2)}\n`);
return 0;
@@ -131,7 +143,6 @@ function parseSshFileTransferCliOptions(args: string[]): SshFileTransferCliOptio
const action = args[0];
if (action !== "upload" && action !== "download") throw new Error("ssh file transfer requires upload or download");
const positionals: string[] = [];
let readBlockBytes = fileTransferReadBlockBytes;
for (let index = 1; index < args.length; index += 1) {
const arg = args[index] ?? "";
if (arg === "--") {
@@ -139,14 +150,10 @@ function parseSshFileTransferCliOptions(args: string[]): SshFileTransferCliOptio
break;
}
if (arg === "--chunk-bytes" || arg === "--block-bytes") {
const value = args[index + 1];
if (value === undefined) throw new Error(`ssh ${action} ${arg} requires a value`);
readBlockBytes = boundedTransferChunkBytes(value, `ssh ${action} ${arg}`);
index += 1;
continue;
throw new Error(`unsupported ssh ${action} option: ${arg}; downloads use host.ssh.tcp-pool stdout streaming and no longer support the legacy base64 block reader`);
}
if (arg === "--help" || arg === "-h" || arg === "help") {
throw new Error(`ssh ${action} usage: trans <route> ${action} ${action === "upload" ? "<local-file> <remote-file>" : "<remote-file> <local-file>"} [--chunk-bytes N]`);
throw new Error(`ssh ${action} usage: trans <route> ${action} ${action === "upload" ? "<local-file> <remote-file>" : "<remote-file> <local-file>"}`);
}
if (arg.startsWith("-")) throw new Error(`unsupported ssh ${action} option: ${arg}`);
positionals.push(arg);
@@ -157,14 +164,8 @@ function parseSshFileTransferCliOptions(args: string[]): SshFileTransferCliOptio
const [first, second] = positionals;
if (!first || !second) throw new Error(`ssh ${action} paths must be non-empty`);
return action === "upload"
? { action, localPath: first, remotePath: second, readBlockBytes }
: { action, remotePath: first, localPath: second, readBlockBytes };
}
function boundedTransferChunkBytes(raw: string, option: string): number {
const value = Number(raw);
if (!Number.isInteger(value) || value <= 0) throw new Error(`${option} must be a positive integer`);
return Math.min(4 * 1024 * 1024, Math.max(1024, value));
? { action, localPath: first, remotePath: second }
: { action, remotePath: first, localPath: second };
}
async function writeRemoteFileVerified(
@@ -203,150 +204,117 @@ async function writeRemoteFileVerified(
return { strategy: "chunked-stdin", chunks };
}
async function readRemoteFileVerified(
async function downloadRemoteFileVerified(
invocation: ParsedSshInvocation,
executor: SshRemoteCommandExecutor,
builders: SshFileTransferCommandBuilders,
remotePath: string,
readBlockBytes: number,
): Promise<{ remote: SshFileTransferStat; content: Buffer; chunks: number }> {
localPath: string,
): Promise<SshFileTransferDownloadResult> {
if (executor.streamRemoteCommand === undefined) {
throw new SshFileTransferError("ssh download requires a tcp-pool stdout stream sink", {
route: invocation.route.raw,
providerId: invocation.providerId,
remotePath,
localPath,
requiredTransport: "host.ssh.tcp-pool",
blocker: "streamRemoteCommand unavailable",
});
}
if (invocation.route.plane === "win") {
throw new SshFileTransferError("ssh download requires tcp-pool stdout streaming; win routes are not supported by the safe binary stream path", {
route: invocation.route.raw,
providerId: invocation.providerId,
remotePath,
localPath,
requiredTransport: "host.ssh.tcp-pool",
blocker: "win route has no verified binary stdout stream path",
});
}
const remote = await statRemoteFile(invocation, executor, builders, remotePath);
const chunks: Buffer[] = [];
await mkdir(path.dirname(localPath), { recursive: true });
const remoteCommand = buildStreamDownloadRemoteCommand(invocation.route, builders, remotePath);
const startedAtMs = Date.now();
const hash = createHash("sha256");
const output = createWriteStream(localPath, { flags: "w", mode: 0o600 });
let actualBytes = 0;
let chunkCount = 0;
for (let blockIndex = 0; blockIndex * readBlockBytes < remote.bytes; blockIndex += 1) {
const offsetBytes = blockIndex * readBlockBytes;
const expectedChunkBytes = Math.min(readBlockBytes, remote.bytes - offsetBytes);
const chunk = await readRemoteBase64BlockWithRetry(invocation, executor, builders, remotePath, readBlockBytes, blockIndex, expectedChunkBytes, remote.bytes, actualBytes);
chunks.push(chunk);
actualBytes += chunk.length;
chunkCount += 1;
emitDownloadProgress(invocation, remotePath, blockIndex, chunkCount, actualBytes, remote.bytes, chunk.length);
}
const content = Buffer.concat(chunks);
const actual = { bytes: content.length, sha256: sha256HexBuffer(content) };
assertTransferStat("download remote read verification", remotePath, remote, actual);
return { remote, content, chunks: chunkCount };
}
async function readRemoteBase64BlockWithRetry(
invocation: ParsedSshInvocation,
executor: SshRemoteCommandExecutor,
builders: SshFileTransferCommandBuilders,
remotePath: string,
readBlockBytes: number,
blockIndex: number,
expectedChunkBytes: number,
expectedBytes: number,
actualBytes: number,
): Promise<Buffer> {
const attemptErrors: Array<Record<string, unknown>> = [];
for (let attempt = 1; attempt <= fileTransferReadBlockMaxAttempts; attempt += 1) {
try {
const read = await checkedFileTransfer(invocation, executor, builders, "read-b64-block", [remotePath, String(blockIndex), String(readBlockBytes)]);
const encoded = read.stdout.replace(/\s+/gu, "");
const chunk = decodeRemoteBase64Block(encoded);
if (chunk !== null && chunk.length === expectedChunkBytes) return chunk;
const reason = chunk === null ? "invalid-base64" : chunk.length === 0 ? "empty-read" : "short-read";
attemptErrors.push({
attempt,
reason,
exitCode: read.exitCode,
stdoutBytes: read.stdout.length,
encodedBytes: encoded.length,
decodedBytes: chunk?.length ?? null,
expectedChunkBytes,
stderrTail: read.stderr.slice(-500),
try {
const result = await executor.streamRemoteCommand(remoteCommand, {
onStdout: async (chunk) => {
await writeStreamChunk(output, chunk);
actualBytes += chunk.length;
chunkCount += 1;
hash.update(chunk);
emitDownloadProgress(invocation, remotePath, chunkCount, actualBytes, remote.bytes, chunk.length, startedAtMs);
},
});
await closeWriteStream(output);
if (result.exitCode !== 0) {
throw new SshFileTransferError("remote ssh download stream failed", {
route: invocation.route.raw,
providerId: invocation.providerId,
remotePath,
localPath,
requiredTransport: "host.ssh.tcp-pool",
exitCode: result.exitCode,
stdout: transferTextSnapshot(result.stdout, { headChars: 120, tailChars: 500 }),
stderr: transferTextSnapshot(result.stderr, { headChars: 120, tailChars: 1000 }),
});
if (attempt < fileTransferReadBlockMaxAttempts) {
emitReadRetryProgress(invocation, remotePath, blockIndex, attempt, expectedBytes, actualBytes, expectedChunkBytes, chunk?.length ?? null, reason);
await delayMs(fileTransferReadEmptyRetryDelayMs);
}
} catch (error) {
attemptErrors.push({
attempt,
error: error instanceof Error ? error.message : String(error),
expectedChunkBytes,
});
if (attempt < fileTransferReadBlockMaxAttempts) {
emitReadRetryProgress(invocation, remotePath, blockIndex, attempt, expectedBytes, actualBytes, expectedChunkBytes, null, "remote-error");
await delayMs(fileTransferReadEmptyRetryDelayMs);
}
}
const elapsedMs = Math.max(1, Date.now() - startedAtMs);
const local = { bytes: actualBytes, sha256: hash.digest("hex") };
assertTransferStat("download tcp-pool stdout stream verification", localPath, remote, local);
emitDownloadProgress(invocation, remotePath, Math.max(1, chunkCount), actualBytes, remote.bytes, 0, startedAtMs, true);
return {
remote,
local,
strategy: "tcp-pool-stdout-stream",
transport: "host.ssh.tcp-pool",
chunks: chunkCount,
elapsedMs,
throughputBytesPerSecond: Math.round((actualBytes * 1000) / elapsedMs),
remainingBytes: Math.max(0, remote.bytes - actualBytes),
};
} catch (error) {
output.destroy();
await rm(localPath, { force: true }).catch(() => undefined);
throw error;
}
throw new SshFileTransferError("remote download returned an invalid block before EOF after retries", {
route: invocation.route.raw,
remotePath,
blockIndex,
attempts: fileTransferReadBlockMaxAttempts,
expectedBytes,
actualBytes,
expectedChunkBytes,
attemptErrors,
});
}
function decodeRemoteBase64Block(encoded: string): Buffer | null {
if (encoded.length === 0) return Buffer.alloc(0);
if (!/^[A-Za-z0-9+/]*={0,2}$/u.test(encoded) || encoded.length % 4 !== 0) return null;
return Buffer.from(encoded, "base64");
}
function emitDownloadProgress(
invocation: ParsedSshInvocation,
remotePath: string,
blockIndex: number,
chunkCount: number,
actualBytes: number,
expectedBytes: number,
lastChunkBytes: number,
startedAtMs: number,
force = false,
): void {
if (chunkCount !== 1 && actualBytes < expectedBytes && chunkCount % fileTransferProgressEveryChunks !== 0) return;
if (!force && chunkCount !== 1 && actualBytes < expectedBytes && chunkCount % fileTransferProgressEveryChunks !== 0) return;
const elapsedMs = Math.max(1, Date.now() - startedAtMs);
process.stderr.write(`${JSON.stringify({
event: "unidesk.ssh.download.progress",
at: new Date().toISOString(),
route: invocation.route.raw,
providerId: invocation.providerId,
remotePath,
blockIndex,
strategy: "tcp-pool-stdout-stream",
transport: "host.ssh.tcp-pool",
chunks: chunkCount,
bytes: actualBytes,
totalBytes: expectedBytes,
actualBytes,
expectedBytes,
lastChunkBytes,
remainingBytes: Math.max(0, expectedBytes - actualBytes),
elapsedMs,
throughputBytesPerSecond: Math.round((actualBytes * 1000) / elapsedMs),
})}\n`);
}
function emitReadRetryProgress(
invocation: ParsedSshInvocation,
remotePath: string,
blockIndex: number,
attempt: number,
expectedBytes: number,
actualBytes: number,
expectedChunkBytes: number,
actualChunkBytes: number | null,
reason: string,
): void {
const event = reason === "empty-read" ? "unidesk.ssh.download.empty-read-retry" : "unidesk.ssh.download.short-read-retry";
process.stderr.write(`${JSON.stringify({
event,
route: invocation.route.raw,
providerId: invocation.providerId,
remotePath,
blockIndex,
attempt,
maxAttempts: fileTransferReadBlockMaxAttempts,
reason,
actualBytes,
expectedBytes,
actualChunkBytes,
expectedChunkBytes,
})}\n`);
}
function delayMs(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
async function statRemoteFile(
invocation: ParsedSshInvocation,
executor: SshRemoteCommandExecutor,
@@ -357,6 +325,31 @@ async function statRemoteFile(
return parseFileTransferStat(stat.stdout, stat.stderr, remotePath);
}
function buildStreamDownloadRemoteCommand(
route: ParsedSshRoute,
builders: SshFileTransferCommandBuilders,
remotePath: string,
): string {
return builders.buildRouteCommand(route, ["sh", "-c", "cat -- \"$1\"", "unidesk-download", remotePath]);
}
async function writeStreamChunk(stream: ReturnType<typeof createWriteStream>, chunk: Buffer): Promise<void> {
if (chunk.length === 0) return;
await new Promise<void>((resolve, reject) => {
stream.write(chunk, (error) => {
if (error) reject(error);
else resolve();
});
});
}
async function closeWriteStream(stream: ReturnType<typeof createWriteStream>): Promise<void> {
if (stream.closed) return;
const closed = once(stream, "close");
stream.end();
await closed;
}
async function checkedFileTransfer(
invocation: ParsedSshInvocation,
executor: SshRemoteCommandExecutor,
@@ -486,11 +479,6 @@ function posixFileTransferScript(): string {
" stat)",
" target=$1; bytes=$(wc -c < \"$target\" | tr -d '[:space:]'); digest=$(sha256_file \"$target\"); printf '%s %s\\n' \"$bytes\" \"$digest\"",
" ;;",
" read-b64-block)",
" target=$1; block_index=$2; block_size=$3",
" case \"$block_index:$block_size\" in *[!0-9:]*|:*) printf 'invalid read block args\\n' >&2; exit 2;; esac",
" dd if=\"$target\" bs=\"$block_size\" skip=\"$block_index\" count=1 2>/dev/null | base64 | tr -d '\\n'",
" ;;",
" write-b64-argv)",
" target=$1; expected_bytes=$2; expected_sha256=$3; shift 3",
" ensure_parent \"$target\"; base=${target##*/}; dir=.; case \"$target\" in */*) dir=${target%/*};; esac",
@@ -563,7 +551,6 @@ function windowsFileTransferScript(basePath: string | null, operation: SshFileTr
"$target = Resolve-UnideskPath $targetArg;",
"switch ($operation) {",
" 'stat' { if (-not (Test-Path -LiteralPath $target -PathType Leaf)) { Fail ('file not found: ' + $target) 1 }; $bytes = ([System.IO.FileInfo]$target).Length; $digest = Get-Sha256 $target; [Console]::Out.WriteLine(([string]$bytes) + ' ' + $digest); break }",
" 'read-b64-block' { $blockIndex = [Int64]$arg1; $blockSize = [Int32]$arg2; if ($blockIndex -lt 0 -or $blockSize -le 0) { Fail 'invalid read block args' 2 }; $fs = [System.IO.File]::Open($target, [System.IO.FileMode]::Open, [System.IO.FileAccess]::Read, [System.IO.FileShare]::ReadWrite); try { [void]$fs.Seek($blockIndex * $blockSize, [System.IO.SeekOrigin]::Begin); $buffer = New-Object byte[] $blockSize; $read = $fs.Read($buffer, 0, $blockSize); if ($read -gt 0) { [Console]::Out.Write([Convert]::ToBase64String($buffer, 0, $read)) } } finally { $fs.Dispose() }; break }",
" 'write-b64-argv' { Decode-ToTarget $target ([string]::Concat($argvChunks)) ([Int64]$arg1) $arg2; break }",
" 'write-b64-stdin' { Decode-ToTarget $target ([Console]::In.ReadToEnd()) ([Int64]$arg1) $arg2; break }",
" 'write-b64-begin' { Ensure-Parent $target; Set-TmpPaths $target $arg1; [System.IO.File]::WriteAllText($script:tmpB64, '', [System.Text.Encoding]::ASCII); break }",