@nestjs/websockets/web-sockets-controller.js

Nest - modern, fast, powerful node.js web framework (@websockets)

import { from as fromPromise, isObservable, of, } from 'rxjs';
import { distinctUntilChanged, mergeAll } from 'rxjs/operators';
import { GATEWAY_OPTIONS, MESSAGE_METADATA, PORT_METADATA, } from './constants.js';
import { InvalidSocketPortException } from './errors/invalid-socket-port.exception.js';
import { GatewayMetadataExplorer, } from './gateway-metadata-explorer.js';
import { compareElementAt } from './utils/compare-element.util.js';
import { isFunction } from '@nestjs/common/internal';
import { Logger } from '@nestjs/common';
import { ContextIdFactory, MetadataScanner, } from '@nestjs/core';
import { ExecutionContextHost, REQUEST_CONTEXT_ID, STATIC_CONTEXT, } from '@nestjs/core/internal';
export class WebSocketsController {
    socketServerProvider;
    config;
    contextCreator;
    container;
    injector;
    exceptionFiltersContext;
    graphInspector;
    appOptions;
    logger = new Logger(WebSocketsController.name, {
        timestamp: true,
    });
    metadataExplorer = new GatewayMetadataExplorer(new MetadataScanner());
    exceptionFiltersCache = new WeakMap();
    pendingDisconnects = new WeakMap();
    constructor(socketServerProvider, config, contextCreator, container, injector, exceptionFiltersContext, graphInspector, appOptions = {}) {
        this.socketServerProvider = socketServerProvider;
        this.config = config;
        this.contextCreator = contextCreator;
        this.container = container;
        this.injector = injector;
        this.exceptionFiltersContext = exceptionFiltersContext;
        this.graphInspector = graphInspector;
        this.appOptions = appOptions;
    }
    connectGatewayToServer(instanceOrWrapper, metatypeOrModuleKey, moduleKey, instanceWrapperId) {
        const isInstanceWrapper = typeof metatypeOrModuleKey === 'string' &&
            'instance' in instanceOrWrapper;
        const instance = isInstanceWrapper
            ? instanceOrWrapper.instance
            : instanceOrWrapper;
        const metatype = isInstanceWrapper
            ? instanceOrWrapper.metatype
            : metatypeOrModuleKey;
        const targetModuleKey = isInstanceWrapper
            ? metatypeOrModuleKey
            : moduleKey;
        const targetInstanceWrapperId = isInstanceWrapper
            ? instanceOrWrapper.id
            : instanceWrapperId;
        const instanceWrapper = isInstanceWrapper
            ? instanceOrWrapper
            : {
                instance,
                metatype,
                id: targetInstanceWrapperId,
                isDependencyTreeStatic: () => true,
                isDependencyTreeDurable: () => false,
                getEnhancersMetadata: () => [],
            };
        const gatewayMetatype = metatype ?? instance.constructor;
        const options = Reflect.getMetadata(GATEWAY_OPTIONS, gatewayMetatype) || {};
        const port = Reflect.getMetadata(PORT_METADATA, gatewayMetatype) || 0;
        if (!Number.isInteger(port)) {
            throw new InvalidSocketPortException(port, gatewayMetatype);
        }
        this.subscribeToServerEvents(instanceWrapper, options, port, targetModuleKey, targetInstanceWrapperId);
    }
    subscribeToServerEvents(instanceOrWrapper, options, port, moduleKey, instanceWrapperId) {
        const instanceWrapper = 'instance' in instanceOrWrapper
            ? instanceOrWrapper
            : {
                instance: instanceOrWrapper,
                metatype: instanceOrWrapper.constructor,
                id: instanceWrapperId,
                isDependencyTreeStatic: () => true,
                isDependencyTreeDurable: () => false,
                getEnhancersMetadata: () => [],
            };
        const { instance } = instanceWrapper;
        const nativeMessageHandlers = this.metadataExplorer.explore(instance);
        const isStatic = instanceWrapper.isDependencyTreeStatic();
        const hasGlobalScopedEnhancers = this.getGlobalScopedEnhancers().length > 0;
        const moduleRef = this.container.getModuleByKey(moduleKey);
        const messageHandlers = nativeMessageHandlers.map(({ callback, isAckHandledManually, message, methodName }) => ({
            message,
            methodName,
            callback: isStatic && !hasGlobalScopedEnhancers
                ? this.contextCreator.create(instance, callback, moduleKey, methodName, STATIC_CONTEXT)
                : this.createRequestScopedHandler(instanceWrapper, moduleRef, moduleKey, methodName),
            isAckHandledManually,
        }));
        this.inspectEntrypointDefinitions(instance, port, messageHandlers, instanceWrapperId);
        if (this.appOptions.preview) {
            return;
        }
        const observableServer = this.socketServerProvider.scanForSocketServer(options, port);
        this.assignServerToProperties(instance, observableServer.server);
        this.subscribeEvents(instanceWrapper, messageHandlers, observableServer, isStatic
            ? this.createStaticEventHandler(instance, moduleKey, 'handleConnection')
            : this.createRequestScopedEventHandler(instanceWrapper, moduleRef, moduleKey, 'handleConnection', observableServer.server), isStatic && !hasGlobalScopedEnhancers
            ? this.createStaticEventHandler(instance, moduleKey, 'handleDisconnect')
            : this.createRequestScopedEventHandler(instanceWrapper, moduleRef, moduleKey, 'handleDisconnect', observableServer.server));
    }
    subscribeEvents(instanceWrapper, subscribersMap, observableServer, connectionHandler, disconnectHandler) {
        const { instance } = instanceWrapper;
        const { init, disconnect, connection, server } = observableServer;
        const adapter = this.config.getIoAdapter();
        this.subscribeInitEvent(instance, init);
        this.subscribeConnectionEvent(connectionHandler, connection);
        this.subscribeDisconnectEvent(disconnectHandler, disconnect);
        const handler = this.getConnectionHandler(this, instance, subscribersMap, disconnect, connection);
        adapter.bindClientConnect(server, handler);
        this.printSubscriptionLogs(instance, subscribersMap);
    }
    getConnectionHandler(context, instance, subscribersMap, disconnect, connection) {
        const adapter = this.config.getIoAdapter();
        return (...args) => {
            const [client] = args;
            connection.next(args);
            context.subscribeMessages(subscribersMap, client, instance);
            const disconnectHook = adapter.bindClientDisconnect;
            disconnectHook &&
                disconnectHook.call(adapter, client, (reason) => disconnect.next({ client, reason }));
        };
    }
    subscribeInitEvent(instance, event) {
        if (instance.afterInit) {
            event.subscribe(instance.afterInit.bind(instance));
        }
    }
    subscribeConnectionEvent(handlerOrGateway, event) {
        const handler = typeof handlerOrGateway === 'function'
            ? handlerOrGateway
            : handlerOrGateway?.handleConnection?.bind(handlerOrGateway);
        if (handler) {
            event
                .pipe(distinctUntilChanged((prev, curr) => compareElementAt(prev, curr, 0)))
                .subscribe((args) => handler(...args));
        }
    }
    subscribeDisconnectEvent(handlerOrGateway, event) {
        const handler = typeof handlerOrGateway === 'function'
            ? handlerOrGateway
            : handlerOrGateway?.handleDisconnect?.bind(handlerOrGateway);
        if (handler) {
            event
                .pipe(distinctUntilChanged((prev, curr) => {
                const prevClient = prev?.client || prev;
                const currClient = curr?.client || curr;
                return prevClient === currClient;
            }))
                .subscribe((data) => {
                if (data && typeof data === 'object' && 'client' in data) {
                    handler(data.client, data.reason);
                }
                else {
                    // Backward compatibility: if it's just the client
                    handler(data);
                }
            });
        }
    }
    subscribeMessages(subscribersMap, client, instance) {
        const adapter = this.config.getIoAdapter();
        const handlers = subscribersMap.map(({ callback, message, isAckHandledManually }) => ({
            message,
            callback: callback.bind(instance, client),
            isAckHandledManually,
        }));
        adapter.bindMessageHandlers(client, handlers, data => fromPromise(this.pickResult(data)).pipe(mergeAll()));
    }
    async pickResult(deferredResult) {
        const result = await deferredResult;
        if (isObservable(result)) {
            return result;
        }
        if (result instanceof Promise) {
            return fromPromise(result);
        }
        return of(result);
    }
    inspectEntrypointDefinitions(instance, port, messageHandlers, instanceWrapperId) {
        messageHandlers.forEach(handler => {
            this.graphInspector.insertEntrypointDefinition({
                type: 'websocket',
                methodName: handler.methodName,
                className: instance.constructor?.name,
                classNodeId: instanceWrapperId,
                metadata: {
                    port,
                    key: handler.message,
                    message: handler.message,
                },
            }, instanceWrapperId);
        });
    }
    createRequestScopedHandler(instanceWrapper, moduleRef, moduleKey, methodName) {
        const { instance } = instanceWrapper;
        const collection = moduleRef.providers;
        const isStatic = instanceWrapper.isDependencyTreeStatic();
        const isTreeDurable = isStatic
            ? this.areGlobalScopedEnhancersDurable()
            : instanceWrapper.isDependencyTreeDurable();
        const scopedEnhancers = [
            ...(isStatic ? instanceWrapper.getEnhancersMetadata() || [] : []),
            ...this.getGlobalScopedEnhancers(),
        ];
        const pattern = Reflect.getMetadata(MESSAGE_METADATA, instance[methodName]) ?? 'unknown';
        return async (...args) => {
            const [client] = args;
            try {
                const contextId = this.getContextId(client, isTreeDurable);
                const contextInstance = isStatic
                    ? instance
                    : await this.injector.loadPerContext(instance, moduleRef, collection, contextId);
                // Resolve global enhancers without adding them to the gateway's dependency
                // tree, so a singleton gateway keeps its state and lifecycle hooks.
                await this.injector.loadEnhancersPerContext(instanceWrapper, contextId, instanceWrapper, scopedEnhancers);
                return this.contextCreator.create(contextInstance, contextInstance[methodName], moduleKey, methodName, contextId, instanceWrapper.id)(...args);
            }
            catch (err) {
                let exceptionFilter = this.exceptionFiltersCache.get(instance[methodName]);
                if (!exceptionFilter) {
                    exceptionFilter = this.exceptionFiltersContext.create(instance, instance[methodName], moduleKey);
                    this.exceptionFiltersCache.set(instance[methodName], exceptionFilter);
                }
                // The context failed to resolve before `WsProxy` could append the
                // pattern, so add it here the same way.
                const host = new ExecutionContextHost([...args, pattern]);
                host.setType('ws');
                exceptionFilter.handle(err, host);
            }
        };
    }
    createRequestScopedEventHandler(instanceWrapper, moduleRef, moduleKey, methodName, server) {
        const { instance } = instanceWrapper;
        const collection = moduleRef.providers;
        const isStatic = instanceWrapper.isDependencyTreeStatic();
        const isTreeDurable = isStatic
            ? this.areGlobalScopedEnhancersDurable()
            : instanceWrapper.isDependencyTreeDurable();
        const targetCallback = instance[methodName];
        return async (...args) => {
            const [client] = args;
            let contextId;
            try {
                contextId = this.getContextId(client, isTreeDurable);
                if (methodName === 'handleDisconnect') {
                    this.pendingDisconnects.set(contextId, (this.pendingDisconnects.get(contextId) ?? 0) + 1);
                }
                const contextInstance = isStatic
                    ? instance
                    : await this.injector.loadPerContext(instance, moduleRef, collection, contextId);
                if (!isStatic) {
                    this.assignServerToProperties(contextInstance, server);
                }
                const scopedMethod = contextInstance[methodName];
                return await scopedMethod?.apply(contextInstance, args);
            }
            catch (err) {
                if (!targetCallback) {
                    // The gateway failed to resolve and has no hook whose exception
                    // filters could handle the error. Nothing awaits this handler, so
                    // rethrowing would end as an unhandled rejection.
                    this.logger.error(`Failed to resolve ${instance.constructor.name} on client ${methodName === 'handleConnection' ? 'connection' : 'disconnection'}`, err?.stack);
                    return;
                }
                this.handleLifecycleHookError(instance, moduleKey, methodName, args, err);
            }
            finally {
                if (methodName === 'handleDisconnect' && contextId) {
                    this.cleanupRequestScopedContext(instanceWrapper, contextId, client);
                }
            }
        };
    }
    // A default-scoped gateway is resolved per connection only for its global
    // scoped enhancers, so their trees decide whether the context is durable
    // (mirrors InstanceWrapper#isDependencyTreeDurable).
    areGlobalScopedEnhancersDurable() {
        const scopedEnhancers = this.getGlobalScopedEnhancers().filter(wrapper => !wrapper.isDependencyTreeStatic());
        return (scopedEnhancers.length > 0 &&
            scopedEnhancers.every(wrapper => wrapper.isDependencyTreeDurable()));
    }
    /**
     * Runs a lifecycle hook of a default-scoped gateway and passes the errors it
     * throws or rejects with to the gateway's exception filters, as for message
     * handlers. The hook still runs synchronously.
     */
    createStaticEventHandler(instance, moduleKey, methodName) {
        const hook = instance[methodName];
        if (!hook) {
            return undefined;
        }
        return (...args) => {
            const handleError = (err) => this.handleLifecycleHookError(instance, moduleKey, methodName, args, err);
            try {
                const result = hook.apply(instance, args);
                return isFunction(result?.then)
                    ? result.catch(handleError)
                    : result;
            }
            catch (err) {
                handleError(err);
            }
        };
    }
    handleLifecycleHookError(instance, moduleKey, methodName, args, err) {
        // Only called for gateways that implement the hook
        const targetCallback = instance[methodName];
        let exceptionFilter = this.exceptionFiltersCache.get(targetCallback);
        if (!exceptionFilter) {
            exceptionFilter = this.exceptionFiltersContext.create(instance, targetCallback, moduleKey);
            this.exceptionFiltersCache.set(targetCallback, exceptionFilter);
        }
        const [client, disconnectReason] = args;
        // Client first, data second and the pattern last, as for message handlers.
        // With the raw hook arguments, the filter would send the last one (the
        // socket itself with socket.io, the upgrade request with ws) back to the
        // client as the pattern. The disconnect reason is safe to pass as data,
        // as nothing is sent to a client that has already disconnected.
        const data = methodName === 'handleDisconnect' ? disconnectReason : undefined;
        const host = new ExecutionContextHost([client, data, methodName]);
        host.setType('ws');
        exceptionFilter.handle(err, host);
    }
    getGlobalScopedEnhancers() {
        return [
            ...this.config.getGlobalRequestGuards(),
            ...this.config.getGlobalRequestPipes(),
            ...this.config.getGlobalRequestInterceptors(),
        ];
    }
    getContextId(request, isTreeDurable) {
        const contextId = ContextIdFactory.getByRequest(request);
        if (!request[REQUEST_CONTEXT_ID]) {
            Object.defineProperty(request, REQUEST_CONTEXT_ID, {
                value: contextId,
                enumerable: false,
                writable: false,
                configurable: true,
            });
            const requestProviderValue = isTreeDurable
                ? contextId.payload
                : Object.assign(request, contextId.payload);
            this.container.registerRequestProvider(requestProviderValue, contextId);
        }
        return contextId;
    }
    cleanupRequestScopedContext(_instanceWrapper, contextId, request) {
        const remaining = (this.pendingDisconnects.get(contextId) ?? 1) - 1;
        if (remaining > 0) {
            this.pendingDisconnects.set(contextId, remaining);
            return;
        }
        this.pendingDisconnects.delete(contextId);
        // Other gateways can still be executing their asynchronous disconnect hooks.
        // A finished cycle must also never remove a newer context on the same client.
        if (request[REQUEST_CONTEXT_ID] === contextId) {
            Reflect.deleteProperty(request, REQUEST_CONTEXT_ID);
        }
    }
    assignServerToProperties(instance, server) {
        for (const propertyKey of this.metadataExplorer.scanForServerHooks(instance)) {
            Reflect.set(instance, propertyKey, server);
        }
    }
    printSubscriptionLogs(instance, subscribersMap) {
        const gatewayClassName = instance?.constructor?.name;
        if (!gatewayClassName) {
            return;
        }
        subscribersMap.forEach(({ message }) => this.logger.log(`${gatewayClassName} subscribed to the "${message}" message`));
    }
}