vault-link/frontend/sync-client/src/services/websocket-manager.ts

import type { Logger } from "../tracing/logger";
import type { Settings } from "../persistence/settings";
import type { WebSocketServerMessage } from "./types/WebSocketServerMessage";
import type { WebSocketClientMessage } from "./types/WebSocketClientMessage";
import type { CursorPositionFromClient } from "./types/CursorPositionFromClient";
import type { ClientCursors } from "./types/ClientCursors";
import { WEBSOCKET_RECEIVE_TIMEOUT_MS } from "../consts";
import { EventListeners } from "../utils/data-structures/event-listeners";
import { awaitAll } from "../utils/await-all";

export class WebSocketManager {
    public readonly onWebSocketStatusChanged = new EventListeners<
        (isConnected: boolean) => unknown
    >();

    public readonly onRemoteVaultUpdateReceived = new EventListeners<
        () => Promise<void>
    >();

    public readonly onRemoteCursorsUpdateReceived = new EventListeners<
        (cursors: ClientCursors[]) => Promise<void>
    >();

    private isStopped = true;
    private receiveTimeoutId: ReturnType<typeof setTimeout> | undefined;
    private reconnectTimeoutId: ReturnType<typeof setTimeout> | undefined;

    private readonly outstandingPromises = new Set<Promise<void>>();

    private webSocket: WebSocket | undefined;
    private readonly webSocketFactoryImplementation: typeof globalThis.WebSocket;

    public constructor(
        private readonly deviceId: string,
        private readonly logger: Logger,
        private readonly settings: Settings,
        webSocketImplementation?: typeof globalThis.WebSocket
    ) {
        if (webSocketImplementation) {
            this.webSocketFactoryImplementation = webSocketImplementation;
        } else if (typeof globalThis.WebSocket !== "undefined") {
            this.webSocketFactoryImplementation = globalThis.WebSocket;
        } else {
            this.webSocketFactoryImplementation =
                // eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-unsafe-type-assertion -- Node versions without WebSocket use ws.
                require("ws") as typeof globalThis.WebSocket;
        }
    }

    public get isWebSocketConnected(): boolean {
        return (
            this.webSocket !== undefined &&
            this.webSocket.readyState ===
                this.webSocketFactoryImplementation.OPEN
        );
    }

    public start(): void {
        if (!this.isStopped) return;
        this.isStopped = false;
        this.initializeWebSocket();
    }

    public async stop(): Promise<void> {
        this.isStopped = true;
        clearTimeout(this.reconnectTimeoutId);
        this.reconnectTimeoutId = undefined;
        this.retireSocket();
        await this.waitUntilFinished();
    }

    public async waitUntilFinished(): Promise<void> {
        await awaitAll([...this.outstandingPromises]);
    }

    public updateLocalCursors(cursorPositions: CursorPositionFromClient): void {
        if (!this.isWebSocketConnected || !this.webSocket) {
            // A missing cursor update is fine, we can just skip it if needed
            this.logger.warn(
                "WebSocket is not connected, cannot send cursor positions"
            );
            return;
        }

        const message: WebSocketClientMessage = {
            type: "cursorPositions",
            ...cursorPositions
        };

        try {
            this.webSocket.send(JSON.stringify(message));
            this.logger.debug(
                `Sent cursor positions: ${JSON.stringify(cursorPositions)}`
            );
        } catch (error) {
            this.logger.warn(
                `Failed to send cursor positions: ${String(error)}`
            );
        }
    }

    private expectResponse(socket: WebSocket): void {
        clearTimeout(this.receiveTimeoutId);
        // Cursor heartbeats are acknowledged every CURSOR_HEARTBEAT_INTERVAL_MS,
        // including when there are no open editors. An OPEN TCP socket alone
        // proves nothing: silence past the heartbeat deadline means a
        // half-open connection.
        this.receiveTimeoutId = setTimeout(() => {
            if (this.webSocket !== socket || this.isStopped) return;
            this.logger.warn("WebSocket receive deadline expired");
            this.reconnect(socket);
        }, WEBSOCKET_RECEIVE_TIMEOUT_MS);
        // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition -- Browser timers have no unref method.
        this.receiveTimeoutId.unref?.();
    }

    private retireSocket(): void {
        clearTimeout(this.receiveTimeoutId);
        const socket = this.webSocket;
        if (!socket) return;
        this.webSocket = undefined;
        socket.onopen = null;
        socket.onclose = null;
        socket.onmessage = null;
        socket.onerror = (): void => {
            /* Retired transports may still emit errors. */
        };
        this.onWebSocketStatusChanged.trigger(false);
        try {
            socket.close(1000, "WebSocket retired");
        } catch (error) {
            this.logger.warn(`Failed to close WebSocket: ${String(error)}`);
        }
    }

    private reconnect(socket: WebSocket): void {
        if (this.webSocket !== socket) return;
        this.retireSocket();
        this.scheduleReconnect();
    }

    private scheduleReconnect(): void {
        if (this.isStopped) return;
        clearTimeout(this.reconnectTimeoutId);
        this.reconnectTimeoutId = setTimeout(() => {
            this.reconnectTimeoutId = undefined;
            this.initializeWebSocket();
        }, this.settings.getSettings().webSocketRetryIntervalMs);
    }

    private initializeWebSocket(): void {
        if (this.isStopped) return;
        try {
            this.openWebSocket();
        } catch (error) {
            this.logger.warn(`Cannot initialize WebSocket: ${String(error)}`);
            this.webSocket = undefined;
            clearTimeout(this.receiveTimeoutId);
            this.onWebSocketStatusChanged.trigger(false);
            this.scheduleReconnect();
        }
    }

    private openWebSocket(): void {
        this.retireSocket();

        const { remoteUri, vaultName } = this.settings.getSettings();
        const wsUri = new URL(remoteUri);
        wsUri.protocol = wsUri.protocol === "https:" ? "wss:" : "ws:";
        wsUri.pathname = `${wsUri.pathname.replace(/\/$/u, "")}/vaults/${encodeURIComponent(vaultName)}/ws`;
        wsUri.search = "";
        wsUri.hash = "";

        this.logger.info(`Connecting to WebSocket at ${wsUri.toString()}`);

        const socket = new this.webSocketFactoryImplementation(wsUri);
        this.webSocket = socket;
        this.expectResponse(socket);

        socket.onerror = (): void => {
            this.logger.warn("WebSocket transport error");
            this.reconnect(socket);
        };

        socket.onopen = (): void => {
            if (this.webSocket !== socket || this.isStopped) return;
            try {
                const handshake: WebSocketClientMessage = {
                    type: "handshake",
                    token: this.settings.getSettings().token,
                    deviceId: this.deviceId
                };
                socket.send(JSON.stringify(handshake));
                this.logger.info("WebSocket connection opened");
                this.onWebSocketStatusChanged.trigger(true);
                // An observer may stop or replace the connection during delivery.
                if (this.webSocket !== socket)
                    this.onWebSocketStatusChanged.trigger(false);
            } catch (error) {
                this.logger.warn(
                    `WebSocket handshake failed: ${String(error)}`
                );
                this.reconnect(socket);
            }
        };

        socket.onmessage = (event): void => {
            if (this.webSocket !== socket || this.isStopped) return;
            try {
                // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
                const message = JSON.parse(
                    event.data
                ) as WebSocketServerMessage;

                if (["vaultChanged", "cursorPositions"].includes(message.type))
                    this.expectResponse(socket);
                // Track the message handling promise
                const messageHandlingPromise = this.handleWebSocketMessage(
                    message
                )
                    .catch((error: unknown) => {
                        this.logger.error(
                            `Error handling WebSocket message: ${String(error)}`
                        );
                    })
                    .finally(() => {
                        this.outstandingPromises.delete(messageHandlingPromise);
                    });

                this.outstandingPromises.add(messageHandlingPromise);
            } catch (error) {
                this.logger.error(
                    `Error parsing WebSocket message: ${String(error)}`
                );
            }
        };

        socket.onclose = (event): void => {
            this.logger.warn(
                `WebSocket closed with code ${event.code} (${event.reason == "" ? "unknown reason" : event.reason})`
            );
            this.reconnect(socket);
        };
    }

    private async handleWebSocketMessage(
        message: WebSocketServerMessage
    ): Promise<void> {
        if (message.type === "vaultChanged") {
            await this.onRemoteVaultUpdateReceived.triggerAsync();

            // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
        } else if (message.type === "cursorPositions") {
            this.logger.debug(
                `Received cursor positions for ${JSON.stringify(message.clients)}`
            );

            await this.onRemoteCursorsUpdateReceived.triggerAsync(
                message.clients
            );
        } else {
            this.logger.warn(
                `Received unknown message type: ${JSON.stringify(message)}`
            );
        }
    }
}