miniSCP/src/bun/transfers/engine.ts

import { pipeline } from "node:stream/promises";
import type { Readable } from "node:stream";
import type { EntryType, TransferSnapshot, TransferState } from "../../shared/types";
import type { FileSystemProvider } from "../fs/provider";
import type { SessionRegistry } from "../fs/registry";
import { walkTree } from "./walk";

const PROGRESS_INTERVAL_MS = 100;
// workers per job; the shared slot budget is the real limit
const MAX_WORKERS_PER_JOB = 8;

interface Job {
	snapshot: TransferSnapshot;
	type: EntryType;
	cancelled: boolean;
	activeStreams: Set<Readable>;
}

/** Counting semaphore with live resize (concurrency setting changes mid-flight). */
class Semaphore {
	private inUse = 0;
	private waiters: (() => void)[] = [];

	constructor(private capacity: number) {}

	async acquire(): Promise<void> {
		if (this.inUse < this.capacity) {
			this.inUse++;
			return;
		}
		// wake() reserves the slot on our behalf before resolving
		await new Promise<void>((resolve) => this.waiters.push(resolve));
	}

	release(): void {
		this.inUse--;
		this.wake();
	}

	resize(capacity: number): void {
		this.capacity = capacity;
		this.wake();
	}

	private wake(): void {
		while (this.waiters.length && this.inUse < this.capacity) {
			this.inUse++;
			this.waiters.shift()!();
		}
	}
}

const TERMINAL: TransferState[] = ["completed", "failed", "cancelled"];

async function ensureDir(fs: FileSystemProvider, path: string): Promise<void> {
	try {
		const st = await fs.stat(path);
		if (st.type === "dir") return;
		throw new Error(`"${path}" exists and is not a directory`);
	} catch {
		await fs.mkdir(path);
	}
}

/**
 * FIFO queue of copy jobs. Because both providers speak node streams, one
 * pipeline() uniformly covers local->local, upload, download and
 * remote->remote (piped through this process).
 */
export class TransferEngine {
	onUpdate: (snapshot: TransferSnapshot) => void = () => {};

	private jobs = new Map<string, Job>();
	private queue: Job[] = [];
	private running = 0;
	private counter = 0;
	private lastEmit = new Map<string, number>();
	// serial by default; the transfers setting raises it. The same number limits
	// active jobs AND total concurrent file streams (folder jobs parallelize
	// internally out of this shared budget).
	private concurrency = 1;
	private slots = new Semaphore(1);

	constructor(private registry: SessionRegistry) {}

	setConcurrency(n: number): void {
		this.concurrency = Math.max(1, Math.min(8, Math.round(n) || 1));
		this.slots.resize(this.concurrency);
		this.pump(); // a raised limit starts queued jobs immediately
	}

	enqueue(params: {
		srcSessionId: string;
		items: { path: string; type: EntryType }[];
		dstSessionId: string;
		dstDir: string;
	}): string[] {
		const src = this.registry.get(params.srcSessionId);
		const jobIds: string[] = [];
		for (const item of params.items) {
			const id = `job-${++this.counter}`;
			const job: Job = {
				type: item.type,
				cancelled: false,
				activeStreams: new Set(),
				snapshot: {
					id,
					name: src.basename(item.path),
					srcSessionId: params.srcSessionId,
					srcPath: item.path,
					dstSessionId: params.dstSessionId,
					dstDir: params.dstDir,
					state: "queued",
					totalFiles: 0,
					totalBytes: 0,
					doneFiles: 0,
					doneBytes: 0,
					currentFile: "",
					currentFileBytes: 0,
					currentFileTotal: 0,
				},
			};
			this.jobs.set(id, job);
			this.queue.push(job);
			jobIds.push(id);
			this.emit(job, true);
		}
		this.pump();
		return jobIds;
	}

	cancel(jobId: string): void {
		const job = this.jobs.get(jobId);
		if (!job || TERMINAL.includes(job.snapshot.state)) return;
		job.cancelled = true;
		if (job.snapshot.state === "queued") {
			this.queue = this.queue.filter((j) => j !== job);
			this.finish(job, "cancelled");
		} else {
			this.destroyStreams(job, new Error("transfer cancelled"));
		}
	}

	dismiss(jobId: string): void {
		const job = this.jobs.get(jobId);
		if (job && TERMINAL.includes(job.snapshot.state)) this.jobs.delete(jobId);
	}

	/** Fail every non-terminal job that touches a dead session. */
	failSession(sessionId: string, reason: string): void {
		for (const job of this.jobs.values()) {
			const { srcSessionId, dstSessionId, state } = job.snapshot;
			if (TERMINAL.includes(state)) continue;
			if (srcSessionId !== sessionId && dstSessionId !== sessionId) continue;
			job.cancelled = true;
			job.snapshot.error = reason;
			if (state === "queued") {
				this.queue = this.queue.filter((j) => j !== job);
				this.finish(job, "failed");
			} else {
				this.destroyStreams(job, new Error(reason));
			}
		}
	}

	private pump(): void {
		while (this.running < this.concurrency && this.queue.length) {
			const job = this.queue.shift()!;
			this.running++;
			this.run(job)
				.catch(() => {})
				.finally(() => {
					this.running--;
					this.pump();
				});
		}
	}

	private async run(job: Job): Promise<void> {
		const snap = job.snapshot;
		try {
			const src = this.registry.get(snap.srcSessionId);
			const dst = this.registry.get(snap.dstSessionId);
			// same FILESYSTEM (never compare sessionIds — two panels on one server
			// hold two sessions, and that comparison once truncated files to zero)
			const sameFs = src.fsKey === dst.fsKey;

			snap.state = "scanning";
			this.emit(job, true);

			// collect the work list
			let files: { srcPath: string; dstPath: string; size: number; name: string }[];
			if (job.type === "dir") {
				const dstRoot = dst.join(snap.dstDir, snap.name);
				await this.guardAgainstSelfCopy(snap, dstRoot, sameFs, src, dst, true);
				const tree = await walkTree(src, snap.srcPath);
				for (const segments of tree.dirs) {
					if (job.cancelled) throw new Error("transfer cancelled");
					await ensureDir(dst, dst.join(dstRoot, ...segments));
				}
				files = tree.files.map((f) => ({
					srcPath: f.entry.path,
					dstPath: dst.join(dstRoot, ...f.dirSegments, f.entry.name),
					size: f.entry.size,
					name: f.entry.name,
				}));
			} else {
				const dstPath = dst.join(snap.dstDir, snap.name);
				await this.guardAgainstSelfCopy(snap, dstPath, sameFs, src, dst, false);
				const st = await src.stat(snap.srcPath);
				files = [{ srcPath: snap.srcPath, dstPath, size: st.size, name: snap.name }];
			}

			// belt & suspenders: an identical path must never reach the write
			// stream (createWriteStream truncates the destination immediately)
			if (sameFs) {
				for (const file of files) {
					if (file.srcPath === file.dstPath) {
						throw new Error(`"${file.name}": source and destination are the same file`);
					}
				}
			}

			snap.totalFiles = files.length;
			snap.totalBytes = files.reduce((sum, f) => sum + f.size, 0);
			snap.state = "running";
			this.emit(job, true);

			// copy with a small worker pool drawing from the SHARED stream budget:
			// one folder job can use several slots, several jobs share them fairly
			let nextIndex = 0;
			let firstError: unknown = null;

			const copyOne = async (file: (typeof files)[number]): Promise<void> => {
				snap.currentFile = file.name;
				snap.currentFileBytes = 0;
				snap.currentFileTotal = file.size;
				this.emit(job, true);

				const readStream = src.createReadStream(file.srcPath);
				job.activeStreams.add(readStream);
				readStream.on("data", (chunk: Buffer | string) => {
					const n = typeof chunk === "string" ? Buffer.byteLength(chunk) : chunk.length;
					snap.doneBytes += n;
					this.emit(job);
				});
				try {
					await pipeline(readStream, dst.createWriteStream(file.dstPath));
				} finally {
					job.activeStreams.delete(readStream);
				}
				snap.doneFiles++;
				this.emit(job, true);
			};

			const worker = async (): Promise<void> => {
				while (!job.cancelled && !firstError && nextIndex < files.length) {
					const file = files[nextIndex++];
					await this.slots.acquire();
					try {
						if (job.cancelled || firstError) return;
						await copyOne(file);
					} catch (error) {
						firstError ??= error;
						this.destroyStreams(
							job,
							error instanceof Error ? error : new Error(String(error)),
						);
					} finally {
						this.slots.release();
					}
				}
			};

			const workerCount = Math.max(1, Math.min(files.length, MAX_WORKERS_PER_JOB));
			await Promise.all(Array.from({ length: workerCount }, worker));
			if (job.cancelled) throw new Error("transfer cancelled");
			if (firstError) throw firstError;

			this.finish(job, "completed");
		} catch (error) {
			if (!job.cancelled && !snap.error) {
				snap.error = error instanceof Error ? error.message : String(error);
			}
			this.finish(job, job.cancelled && !snap.error ? "cancelled" : "failed");
		}
	}

	/** Tear down every in-flight stream of a job (cancel, session death, first error). */
	private destroyStreams(job: Job, reason: Error): void {
		for (const stream of job.activeStreams) {
			stream.destroy(reason);
		}
		job.activeStreams.clear();
	}

	/**
	 * Copying onto itself truncates the source; into a descendant never ends.
	 * Keyed on filesystem identity (fsKey), not sessionId. When the keys differ
	 * but the paths collide anyway, a probe file decides whether the two
	 * connections secretly reach the same machine (NAT / alias hosts).
	 */
	private async guardAgainstSelfCopy(
		snap: TransferSnapshot,
		dstPath: string,
		sameFs: boolean,
		src: FileSystemProvider,
		dst: FileSystemProvider,
		isDir: boolean,
	): Promise<void> {
		const samePath = dstPath === snap.srcPath;
		const nestedPath = dstPath.startsWith(
			snap.srcPath.endsWith("/") ? snap.srcPath : `${snap.srcPath}/`,
		);
		if (!samePath && !nestedPath) return;

		// Paths collide. Same fsKey -> definitely the same file/folder. Different
		// fsKey -> either a coincidence (two servers both have /root/x.conf) or a
		// hidden alias of one machine; only a probe can tell (SFTP exposes no
		// device/inode identity — the technique cp/rsync use locally).
		const definitelySame =
			sameFs || (await this.probeSameFilesystem(src, dst, dst.dirname(dstPath)));
		if (!definitelySame) return;

		if (samePath) {
			throw new Error(
				isDir ? "Cannot copy a folder onto itself" : "Source and destination are the same file",
			);
		}
		throw new Error("Cannot copy a folder into itself");
	}

	/**
	 * Deterministic same-filesystem test across two connections: write a token
	 * file via dst, look for it via src. Visible + matching content = one
	 * machine behind two addresses. Fails open (returns false) when the probe
	 * itself cannot run (e.g. destination not writable) — the copy then
	 * proceeds and permissions produce their own error.
	 */
	private async probeSameFilesystem(
		src: FileSystemProvider,
		dst: FileSystemProvider,
		dir: string,
	): Promise<boolean> {
		const token = crypto.randomUUID();
		const probePath = dst.join(dir, `.bunscp-probe-${token.slice(0, 8)}`);
		try {
			await new Promise<void>((resolve, reject) => {
				const stream = dst.createWriteStream(probePath);
				stream.on("error", reject);
				stream.end(token, () => resolve());
			});
			try {
				const chunks: Buffer[] = [];
				for await (const chunk of src.createReadStream(src.join(dir, `.bunscp-probe-${token.slice(0, 8)}`))) {
					chunks.push(typeof chunk === "string" ? Buffer.from(chunk) : chunk);
				}
				const seen = Buffer.concat(chunks).toString("utf8") === token;
				if (seen) console.log(`[engine] probe: ${src.fsKey} and ${dst.fsKey} are the same machine`);
				return seen;
			} catch {
				return false; // src cannot see the token -> different filesystems
			}
		} catch (error) {
			console.log(`[engine] same-fs probe skipped: ${error instanceof Error ? error.message : error}`);
			return false;
		} finally {
			await dst.deleteFile(probePath).catch(() => {});
		}
	}

	private finish(job: Job, state: TransferState): void {
		if (TERMINAL.includes(job.snapshot.state)) return;
		job.snapshot.state = state;
		this.emit(job, true);
	}

	private emit(job: Job, force = false): void {
		const now = Date.now();
		if (!force && now - (this.lastEmit.get(job.snapshot.id) ?? 0) < PROGRESS_INTERVAL_MS) {
			return;
		}
		this.lastEmit.set(job.snapshot.id, now);
		this.onUpdate({ ...job.snapshot });
	}
}