forward_server/index.ts

import {createServer} from 'node:http';
import type {Server as HttpServer} from 'node:http';
import type {Socket as NetSocket} from 'node:net';
import {Server} from 'socket.io';
import type {ControlTransport, ForwardServerOptions} from './types.ts';
import {resolve_options} from './utils/relay_config.ts';
import {create_session_store} from './runtime/session_store.ts';
import type {RelaySessionStore} from './runtime/session_store.ts';
import {install_relay_routes} from './runtime/relay_routes.ts';
import {install_control_routes, report_status} from './runtime/init.ts';
import {mount_http} from './runtime/http_mount.ts';

export interface ForwardServerHandle {
    readonly io: Server;
    readonly http_server: HttpServer;
    readonly store: RelaySessionStore;
    readonly forward_id: number;
    readonly path: string;
    readonly namespace: string;
    readonly owns_http_server: boolean;
    bind_control(transport: ControlTransport): () => void;
    stop(): Promise<void>;
}

let current: ForwardServerHandle | null = null;
let initialization: Promise<ForwardServerHandle> | null = null;
let startup_options: ReturnType<typeof resolve_options> | null = null;
let stopping: Promise<void> | null = null;

export async function start_forward_server(input: ForwardServerOptions): Promise<ForwardServerHandle> {
    const options = resolve_options(input);
    if (stopping) await stopping;
    if (startup_options) {
        const same = startup_options.forward_id === options.forward_id && startup_options.key === options.key
            && startup_options.http_server === options.http_server && startup_options.path === options.path
            && startup_options.namespace === options.namespace && startup_options.port === options.port
            && startup_options.host === options.host
            && JSON.stringify(startup_options.allowed_origins) === JSON.stringify(options.allowed_origins)
            && startup_options.allow_no_origin === options.allow_no_origin
            && JSON.stringify(startup_options.limits) === JSON.stringify(options.limits);
        if (!same) throw new Error('Relay already initialized with different options; stop it before reconfiguration');
    }
    if (current) return current;
    if (initialization) return initialization;
    startup_options = options;
    initialization = create_relay(options);
    try {
        current = await initialization;
        return current;
    } catch (error) {
        startup_options = null;
        throw error;
    } finally { initialization = null; }
}

async function create_relay(options: ReturnType<typeof resolve_options>): Promise<ForwardServerHandle> {
    const owns_http_server = !options.http_server;
    const host = options.http_server ?? createServer((_request, response) => {
        response.statusCode = 404;
        response.end();
    });
    // Socket.IO owns only this private server in borrowed mode. Its close()
    // therefore cannot close the Main HTTP server or unrelated Socket.IO services.
    const transport_server = owns_http_server ? host : createServer();
    const owned_connections = new Set<NetSocket>();
    const track_connection = (socket: NetSocket) => {
        owned_connections.add(socket);
        socket.once('close', () => owned_connections.delete(socket));
    };
    if (owns_http_server) host.on('connection', track_connection);
    let control: ControlTransport | undefined;
    let dispose_control = () => {};
    let unmount = () => {};
    const io = new Server(transport_server, {
        path: options.path,
        serveClient: false,
        maxHttpBufferSize: options.limits.max_payload_bytes,
        connectTimeout: options.limits.auth_timeout_ms,
        cors: {origin: options.allowed_origins, credentials: false},
        allowRequest: (request, callback) => {
            const origin = request.headers.origin;
            const permitted = origin === undefined ? options.allow_no_origin : options.allowed_origins.includes(origin);
            callback(null, permitted && io.engine.clientsCount < options.limits.max_connections);
        },
    });
    const store = create_session_store({
        forward_id: options.forward_id,
        key: options.key,
        limits: options.limits,
        on_status: status => {
            try { report_status(control, status); }
            catch { options.logger?.log('Unable to report relay session status'); }
        },
    });
    store.set_control_available(false);
    const namespace = io.of(options.namespace);
    if (options.namespace !== '/') io.of('/').use((_socket, next) => next(new Error('UNAUTHORIZED')));
    const dispose_routes = install_relay_routes(namespace, store);
    let stop_promise: Promise<void> | null = null;
    const handle: ForwardServerHandle = {
        io, http_server: host, store, forward_id: options.forward_id,
        path: options.path, namespace: options.namespace, owns_http_server,
        bind_control(transport) {
            if (stop_promise) throw new Error('Relay is stopped');
            dispose_control();
            control = transport;
            const dispose = install_control_routes(transport, store, options.logger);
            dispose_control = dispose;
            return () => {
                if (dispose_control !== dispose) return;
                dispose();
                dispose_control = () => {};
                control = undefined;
            };
        },
        stop() {
            if (stop_promise) return stop_promise;
            stop_promise = (async () => {
                host.off('close', host_closed);
                store.close();
                dispose_routes();
                dispose_control();
                control = undefined;
                unmount();
                let deadline: ReturnType<typeof setTimeout> | undefined;
                try {
                    await Promise.race([
                        io.close(),
                        new Promise<never>((_resolve, reject) => {
                            deadline = setTimeout(() => {
                                for (const socket of owned_connections) socket.destroy();
                                reject(new Error('Relay shutdown timed out'));
                            }, options.limits.close_timeout_ms);
                        }),
                    ]);
                } finally {
                    clearTimeout(deadline);
                    host.off('connection', track_connection);
                    for (const socket of owned_connections) socket.destroy();
                    owned_connections.clear();
                }
            })().finally(() => {
                if (current === handle) current = null;
                startup_options = null;
                stopping = null;
            });
            stopping = stop_promise;
            return stop_promise;
        },
    };
    const host_closed = () => { void handle.stop().catch(() => options.logger?.log('Relay shutdown failed')); };
    try {
        if (options.control) handle.bind_control(options.control);
        if (owns_http_server) {
            await new Promise<void>((resolve, reject) => {
                const on_error = (error: Error) => { host.off('listening', on_listening); reject(error); };
                const on_listening = () => { host.off('error', on_error); resolve(); };
                host.once('error', on_error);
                host.once('listening', on_listening);
                host.listen(options.port, options.host);
            });
        } else unmount = mount_http(host, options.path, transport_server);
        host.once('close', host_closed);
        return handle;
    } catch (error) {
        await handle.stop();
        throw error;
    }
}

export async function stop_forward_server(): Promise<void> {
    if (initialization) {
        try { await initialization; } catch { return; }
    }
    if (current) await current.stop();
    else if (stopping) await stopping;
}

export type * from './types.ts';