126 lines
3.6 KiB
TypeScript
126 lines
3.6 KiB
TypeScript
import { randomUUID } from "node:crypto";
|
|
import { WebSocketServer, type WebSocket } from "ws";
|
|
import { formatComputer, parseJsonFrame, parseLinkMessage, type ResponseMessage } from "./protocol.js";
|
|
|
|
export type ComputerConnection = {
|
|
computerId: number;
|
|
label: string | null;
|
|
ws: WebSocket;
|
|
connectedAt: number;
|
|
lastSeenAt: number;
|
|
};
|
|
|
|
type PendingRequest = {
|
|
computerId: number;
|
|
resolve: (message: ResponseMessage) => void;
|
|
};
|
|
|
|
export class LinkRegistry {
|
|
private readonly computers = new Map<number, ComputerConnection>();
|
|
private readonly pending = new Map<string, PendingRequest>();
|
|
|
|
register(connection: ComputerConnection): void {
|
|
const existing = this.computers.get(connection.computerId);
|
|
if (existing && existing.ws !== connection.ws) {
|
|
existing.ws.close();
|
|
}
|
|
this.computers.set(connection.computerId, connection);
|
|
}
|
|
|
|
unregister(ws: WebSocket): void {
|
|
for (const [computerId, connection] of this.computers) {
|
|
if (connection.ws === ws) {
|
|
this.computers.delete(computerId);
|
|
}
|
|
}
|
|
}
|
|
|
|
count(): number {
|
|
return this.computers.size;
|
|
}
|
|
|
|
snapshot(): ComputerConnection[] {
|
|
return [...this.computers.values()];
|
|
}
|
|
|
|
handleFrame(ws: WebSocket, data: unknown): void {
|
|
const message = parseLinkMessage(parseJsonFrame(data));
|
|
if (!message) {
|
|
return;
|
|
}
|
|
|
|
if (message.type === "hello") {
|
|
const now = Date.now();
|
|
this.register({
|
|
computerId: message.computerId,
|
|
label: message.computerLabel,
|
|
ws,
|
|
connectedAt: now,
|
|
lastSeenAt: now,
|
|
});
|
|
ws.send(JSON.stringify({ type: "hello-ok" }));
|
|
return;
|
|
}
|
|
|
|
const pending = this.pending.get(message.id);
|
|
if (!pending) {
|
|
return;
|
|
}
|
|
|
|
const connection = this.computers.get(pending.computerId);
|
|
if (connection) {
|
|
connection.lastSeenAt = Date.now();
|
|
}
|
|
this.pending.delete(message.id);
|
|
pending.resolve(message);
|
|
}
|
|
|
|
async probeComputers(timeoutMs: number): Promise<string> {
|
|
const computers = this.snapshot();
|
|
if (computers.length === 0) {
|
|
return "No computers connected.";
|
|
}
|
|
|
|
const lines = await Promise.all(computers.map((computer) => this.probeComputer(computer, timeoutMs)));
|
|
return lines.join("\n");
|
|
}
|
|
|
|
private async probeComputer(computer: ComputerConnection, timeoutMs: number): Promise<string> {
|
|
const id = randomUUID();
|
|
const response = new Promise<ResponseMessage>((resolve) => {
|
|
this.pending.set(id, { computerId: computer.computerId, resolve });
|
|
computer.ws.send(JSON.stringify({ type: "request", id, method: "ping" }));
|
|
});
|
|
|
|
const timeout = new Promise<null>((resolve) => {
|
|
setTimeout(() => {
|
|
this.pending.delete(id);
|
|
resolve(null);
|
|
}, timeoutMs);
|
|
});
|
|
|
|
const result = await Promise.race([response, timeout]);
|
|
if (!result) {
|
|
return `timeout from ${formatComputer(computer.computerId, computer.label)}`;
|
|
}
|
|
|
|
if (result.ok && typeof result.result === "string") {
|
|
return result.result;
|
|
}
|
|
|
|
return `error from ${formatComputer(computer.computerId, computer.label)}: ${result.error ?? "unknown error"}`;
|
|
}
|
|
}
|
|
|
|
export function startLinkServer(options: { host: string; port: number; registry: LinkRegistry }): WebSocketServer {
|
|
const server = new WebSocketServer({ host: options.host, port: options.port });
|
|
|
|
server.on("connection", (ws) => {
|
|
ws.on("message", (data) => options.registry.handleFrame(ws, data));
|
|
ws.on("close", () => options.registry.unregister(ws));
|
|
ws.on("error", () => options.registry.unregister(ws));
|
|
});
|
|
|
|
return server;
|
|
}
|