feat(core): Add support for WebSockets as an alternative to Server-Sent Events (#5443)

Co-authored-by: Matthijs Knigge <matthijs@volcano.nl>

packages/cli/src/push/abstract.push.ts

@@ -0,0 +1,49 @@
+import { LoggerProxy as Logger } from 'n8n-workflow';
+import type { IPushDataType } from '@/Interfaces';
+
+export abstract class AbstractPush<T> {
+	protected connections: Record<string, T> = {};
+
+	protected abstract close(connection: T): void;
+	protected abstract sendToOne(connection: T, data: string): void;
+
+	protected add(sessionId: string, connection: T): void {
+		const { connections } = this;
+		Logger.debug('Add editor-UI session', { sessionId });
+
+		const existingConnection = connections[sessionId];
+		if (existingConnection) {
+			// Make sure to remove existing connection with the same id
+			this.close(existingConnection);
+		}
+
+		connections[sessionId] = connection;
+	}
+
+	protected remove(sessionId?: string): void {
+		if (sessionId !== undefined) {
+			Logger.debug('Remove editor-UI session', { sessionId });
+			delete this.connections[sessionId];
+		}
+	}
+
+	send<D>(type: IPushDataType, data: D, sessionId: string | undefined = undefined) {
+		const { connections } = this;
+		if (sessionId !== undefined && connections[sessionId] === undefined) {
+			Logger.error(`The session "${sessionId}" is not registered.`, { sessionId });
+			return;
+		}
+
+		Logger.debug(`Send data of type "${type}" to editor-UI`, { dataType: type, sessionId });
+
+		const sendData = JSON.stringify({ type, data });
+
+		if (sessionId === undefined) {
+			// Send to all connected clients
+			Object.values(connections).forEach((connection) => this.sendToOne(connection, sendData));
+		} else {
+			// Send only to a specific client
+			this.sendToOne(connections[sessionId], sendData);
+		}
+	}
+}

packages/cli/src/push/index.ts

@@ -0,0 +1,105 @@
+import { ServerResponse } from 'http';
+import type { Server } from 'http';
+import type { Socket } from 'net';
+import type { Application, RequestHandler } from 'express';
+import { Server as WSServer } from 'ws';
+import { parse as parseUrl } from 'url';
+import config from '@/config';
+import { resolveJwt } from '@/auth/jwt';
+import { AUTH_COOKIE_NAME } from '@/constants';
+import { SSEPush } from './sse.push';
+import { WebSocketPush } from './websocket.push';
+import type { Push, PushResponse, SSEPushRequest, WebSocketPushRequest } from './types';
+export type { Push } from './types';
+
+const useWebSockets = config.getEnv('push.backend') === 'websocket';
+
+let pushInstance: Push;
+export const getPushInstance = () => {
+	if (!pushInstance) pushInstance = useWebSockets ? new WebSocketPush() : new SSEPush();
+	return pushInstance;
+};
+
+export const setupPushServer = (restEndpoint: string, server: Server, app: Application) => {
+	if (useWebSockets) {
+		const wsServer = new WSServer({ noServer: true });
+		server.on('upgrade', (request: WebSocketPushRequest, socket: Socket, head) => {
+			if (parseUrl(request.url).pathname === `/${restEndpoint}/push`) {
+				wsServer.handleUpgrade(request, socket, head, (ws) => {
+					request.ws = ws;
+
+					const response = new ServerResponse(request);
+					response.writeHead = (statusCode) => {
+						if (statusCode > 200) ws.close();
+						return response;
+					};
+
+					// @ts-ignore
+					// eslint-disable-next-line @typescript-eslint/no-unsafe-call
+					app.handle(request, response);
+				});
+			}
+		});
+	}
+};
+
+export const setupPushHandler = (
+	restEndpoint: string,
+	app: Application,
+	isUserManagementEnabled: boolean,
+) => {
+	const push = getPushInstance();
+	const endpoint = `/${restEndpoint}/push`;
+
+	const pushValidationMiddleware: RequestHandler = async (
+		req: SSEPushRequest | WebSocketPushRequest,
+		res,
+		next,
+	) => {
+		const ws = req.ws;
+
+		const { sessionId } = req.query;
+		if (sessionId === undefined) {
+			if (ws) {
+				ws.send('The query parameter "sessionId" is missing!');
+				ws.close(400);
+			} else {
+				next(new Error('The query parameter "sessionId" is missing!'));
+			}
+			return;
+		}
+
+		// Handle authentication
+		if (isUserManagementEnabled) {
+			try {
+				// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access
+				const authCookie: string = req.cookies?.[AUTH_COOKIE_NAME] ?? '';
+				await resolveJwt(authCookie);
+			} catch (error) {
+				if (ws) {
+					ws.send(`Unauthorized: ${(error as Error).message}`);
+					ws.close(401);
+				} else {
+					res.status(401).send('Unauthorized');
+				}
+				return;
+			}
+		}
+
+		next();
+	};
+
+	app.use(
+		endpoint,
+		pushValidationMiddleware,
+		(req: SSEPushRequest | WebSocketPushRequest, res: PushResponse) => {
+			if (req.ws) {
+				(push as WebSocketPush).add(req.query.sessionId, req.ws);
+			} else if (!useWebSockets) {
+				(push as SSEPush).add(req.query.sessionId, { req, res });
+			} else {
+				res.status(401).send('Unauthorized');
+			}
+		},
+	);
+};

packages/cli/src/push/sse.push.ts

@@ -0,0 +1,32 @@
+import SSEChannel from 'sse-channel';
+import { AbstractPush } from './abstract.push';
+import type { PushRequest, PushResponse } from './types';
+
+type Connection = { req: PushRequest; res: PushResponse };
+
+export class SSEPush extends AbstractPush<Connection> {
+	readonly channel = new SSEChannel();
+
+	readonly connections: Record<string, Connection> = {};
+
+	constructor() {
+		super();
+		this.channel.on('disconnect', (channel, { req }) => {
+			this.remove(req?.query?.sessionId);
+		});
+	}
+
+	add(sessionId: string, connection: Connection) {
+		super.add(sessionId, connection);
+		this.channel.addClient(connection.req, connection.res);
+	}
+
+	protected close({ res }: Connection): void {
+		res.end();
+		this.channel.removeClient(res);
+	}
+
+	protected sendToOne(connection: Connection, data: string): void {
+		this.channel.send(data, [connection.res]);
+	}
+}

packages/cli/src/push/types.ts

@@ -0,0 +1,15 @@
+import type { Request, Response } from 'express';
+import type { WebSocket } from 'ws';
+import type { SSEPush } from './sse.push';
+import type { WebSocketPush } from './websocket.push';
+
+// TODO: move all push related types here
+
+export type Push = SSEPush | WebSocketPush;
+
+export type PushRequest = Request<{}, {}, {}, { sessionId: string }>;
+
+export type SSEPushRequest = PushRequest & { ws: undefined };
+export type WebSocketPushRequest = PushRequest & { ws: WebSocket };
+
+export type PushResponse = Response & { req: PushRequest };

packages/cli/src/push/websocket.push.ts

@@ -0,0 +1,19 @@
+import type WebSocket from 'ws';
+import { AbstractPush } from './abstract.push';
+
+export class WebSocketPush extends AbstractPush<WebSocket> {
+	add(sessionId: string, connection: WebSocket) {
+		super.add(sessionId, connection);
+
+		// Makes sure to remove the session if the connection is closed
+		connection.once('close', () => this.remove(sessionId));
+	}
+
+	protected close(connection: WebSocket): void {
+		connection.close();
+	}
+
+	protected sendToOne(connection: WebSocket, data: string): void {
+		connection.send(data);
+	}
+}