factorio-clusterio/test/controller/WsServer.js

import assert from "node:assert/strict";
import events from "node:events";
import jwt from "jsonwebtoken";
import * as lib from "@clusterio/lib";
import { Controller, HostRecord, WsServer } from "@clusterio/controller";
import { ControllerConfig } from "@clusterio/lib";
import packageJson from "@clusterio/controller/package.json" with { type: "json" };

import "../setup_logging.js";

class MockWebSocket extends events.EventEmitter {
	constructor() {
		super();
		this.sent = [];
		this.closed = null;
	}

	send(message) {
		this.sent.push(JSON.parse(message));
	}

	close(code, reason) {
		this.closed = { code, reason };
	}

	terminate() {
		this.emit("close", 1006, "");
	}
}

describe("controller/src/WsServer", function() {
	describe("class WsServer", function() {
		/** @type {WsServer} */
		let wsServer;
		before(function() {
			let config = new ControllerConfig("controller");
			config.set("controller.trusted_proxies", "127.0.0.0/8, ::1");
			let controller = new Controller({}, [], config);
			wsServer = controller.wsServer;
		});
		describe(".remoteAddr()", function() {
			it("should resolve X-Forward-For header for trusted proxies", function() {
				const ipv4 = wsServer.remoteAddr({
					socket: {
						remoteAddress: "127.0.0.1",
						remoteFamily: "IPv4",
					},
					headers: { "x-forwarded-for": "1.2.3.4" },
				});
				assert.equal(ipv4, "1.2.3.4");
				const ipv6 = wsServer.remoteAddr({
					socket: {
						remoteAddress: "::1",
						remoteFamily: "IPv6",
					},
					headers: { "x-forwarded-for": "2000::1" },
				});
				assert.equal(ipv6, "2000::1");
			});
			it("should pick the last address in X-Forward-For header for trusted proxies", function() {
				const ipv4 = wsServer.remoteAddr({
					socket: {
						remoteAddress: "127.0.0.1",
						remoteFamily: "IPv4",
					},
					headers: { "x-forwarded-for": "1.0.0.0, 2.0.0.0, 1.2.3.4" },
				});
				assert.equal(ipv4, "1.2.3.4");
				const ipv6 = wsServer.remoteAddr({
					socket: {
						remoteAddress: "::1",
						remoteFamily: "IPv6",
					},
					headers: { "x-forwarded-for": "2001::, 2002::, 2000::1" },
				});
				assert.equal(ipv6, "2000::1");
			});
			it("should not resolve X-Forward-For header for untrusted proxies", function() {
				const ipv4 = wsServer.remoteAddr({
					socket: {
						remoteAddress: "1.2.3.4",
						remoteFamily: "IPv4",
					},
					headers: { "x-forwarded-for": "127.0.0.1" },
				});
				assert.equal(ipv4, "1.2.3.4");
				const ipv6 = wsServer.remoteAddr({
					socket: {
						remoteAddress: "2000::1",
						remoteFamily: "IPv6",
					},
					headers: { "x-forwarded-for": "::1" },
				});
				assert.equal(ipv6, "2000::1");
			});
		});
		describe(".handleConnection()", function() {
			it("should handle socket errors before the handshake completes", function() {
				// A ws socket emits "error" for e.g. a close frame with the reserved code 1006
				const socket = new events.EventEmitter();
				socket.send = () => {};
				socket.close = () => {};
				socket.terminate = () => {};
				const req = { socket: { remoteAddress: "1.2.3.4", remoteFamily: "IPv4" }, headers: {} };
				wsServer.handleConnection(socket, req);
				try {
					assert.doesNotThrow(() => {
						socket.emit("error", new RangeError("Invalid WebSocket frame: invalid status code 1006"));
					});
				} finally {
					socket.emit("message", "not json"); // ends the handshake and its timeout
				}
			});
		});
	});
	describe("connection metrics", function() {
		const req = { socket: { remoteAddress: "10.0.0.1", remoteFamily: "IPv4" }, headers: {} };
		/** @type {WsServer} */
		let wsServer;
		before(function() {
			const config = new ControllerConfig("controller");
			config.set("controller.auth_secret", Buffer.from("TestSecretDoNotUse").toString("base64"));
			wsServer = new Controller(lib.logger, [], config).wsServer;
		});

		function counter(name) {
			return lib.defaultRegistry.collectors.find(collector => collector.metric.name === name);
		}
		function value(name, type) {
			const collector = counter(`clusterio_controller_websocket_${name}_total`);
			return type ? collector.labels(type).get() : collector.get();
		}

		[["host", lib.Address.host], ["control", lib.Address.control]].forEach(([type, addressType]) => {
			it(`should count ready, drop, resume and timeout for ${type} sessions`, function() {
				const before = Object.fromEntries(
					["ready", "drop", "resume", "session_timeout"].map(name => [name, value(name, type)])
				);
				const [connector, sessionToken] = wsServer.createSession(new lib.Address(addressType, 1));
				connector.ready(new MockWebSocket(), new lib.Address(addressType, 1), sessionToken, undefined);
				assert.equal(value("ready", type), before.ready + 1);

				connector._socket.terminate();
				assert.equal(value("drop", type), before.drop + 1);

				wsServer.resumeSession(new lib.ResumeData(sessionToken, 0), new MockWebSocket(), req);
				assert.equal(value("resume", type), before.resume + 1);

				connector._socket.terminate();
				connector._timedOut();
				assert.equal(value("session_timeout", type), before.session_timeout + 1);
				assert.equal(value("drop", type), before.drop + 2);
				assert.equal(value("ready", type), before.ready + 1);
			});
		});
		it("should count resumes of unknown sessions as failed", function() {
			const before = value("failed_resume");
			const socket = new MockWebSocket();
			wsServer.resumeSession(new lib.ResumeData("invalid", 0), socket, req);
			try {
				assert.equal(value("failed_resume"), before + 1);
				assert.equal(socket.sent[0].type, "invalidate");
			} finally {
				socket.emit("message", "not json"); // ends the handshake and its timeout
			}
		});
	});
	describe(".registerHost()", function() {
		const secret = Buffer.from("TestSecretDoNotUse");
		const req = { socket: { remoteAddress: "10.0.0.1", remoteFamily: "IPv4" }, headers: {} };
		/** @type {Controller} */
		let controller;
		/** @type {WsServer} */
		let wsServer;
		let socket;
		beforeEach(function() {
			const config = new ControllerConfig("controller");
			config.set("controller.auth_secret", secret.toString("base64"));
			controller = new Controller(lib.logger, [], config);
			wsServer = controller.wsServer;
			socket = new MockWebSocket();
		});
		afterEach(function() {
			for (const connection of wsServer.hostConnections.values()) {
				connection.connector.terminate();
			}
		});

		function addHost(id, tokenValidAfter = 0) {
			const record = new HostRecord(id, `host-${id}`, "", new Map(), false, "", "203.0.113.1", tokenValidAfter);
			controller.hosts.records.set(record);
			return record;
		}

		async function register(token, id = 5) {
			await wsServer.registerHost(new lib.RegisterHostData(token, packageJson.version, id, {}), socket, req);
		}

		function assertRejected(reason) {
			assert.deepEqual(socket.closed, { code: lib.ConnectionClosed.Unauthorized, reason });
			assert.equal(socket.closed.code, 3000);
			assert.equal(wsServer.hostConnections.size, 0);
			assert(!socket.sent.some(message => message.type === "ready"), "ready was sent");
		}

		it("should reject a validly signed token for a host that is not registered", async function() {
			await register(controller.hosts.signHostToken(5));
			assertRejected(
				"Authentication failed: host is not registered, start the host again to request to join the cluster"
			);
			assert.equal(controller.hosts.has(5), false, "host record was created");
		});
		it("should accept a registered host", async function() {
			const record = addHost(5);
			await register(controller.hosts.signHostToken(5));
			assert.equal(socket.closed, null);
			assert(socket.sent.some(message => message.type === "ready"), "ready was not sent");
			assert(wsServer.hostConnections.has(5));

			assert.equal(controller.hosts.records.size, 1);
			assert.equal(controller.hosts.get(5), record, "host record was replaced");
			assert.equal(record.connected, true);
			assert.equal(record.version, packageJson.version);
			assert.equal(record.remoteAddress, "10.0.0.1");
			assert.equal(record.name, "host-5");
			assert.equal(record.publicAddress, "203.0.113.1");

			wsServer.hostConnections.get(5).connector.terminate();
			assert.equal(wsServer.hostConnections.has(5), false);
			assert.equal(controller.hosts.get(5).connected, false);
		});
		it("should reject a token issued before tokenValidAfter", async function() {
			const token = controller.hosts.signHostToken(5);
			addHost(5, Math.floor(Date.now() / 1000) + 60);
			await register(token);
			assertRejected("Authentication failed: invalid token");
		});
		it("should accept a token issued at tokenValidAfter", async function() {
			const now = Math.floor(Date.now() / 1000);
			addHost(5, now);
			await register(jwt.sign({ aud: "host", host: 5, iat: now }, secret));
			assert.equal(socket.closed, null);
			assert(wsServer.hostConnections.has(5));
		});
		it("should reject a host revoked while its previous connection closes", async function() {
			const record = addHost(5);
			wsServer.hostConnections.set(5, {
				remoteAddress: "10.0.0.1",
				async disconnect() {
					controller.hosts.records.delete(record);
					wsServer.hostConnections.delete(5);
				},
			});
			await register(controller.hosts.signHostToken(5));
			assertRejected("Authentication failed: host was revoked");
			assert.equal(controller.hosts.has(5), false, "host record was recreated");
		});
		it("should reject hosts in recovery mode", async function() {
			addHost(5);
			controller.recoveryMode = true;
			await register(controller.hosts.signHostToken(5));
			assert.deepEqual(socket.closed, {
				code: lib.ConnectionClosed.RecoveryMode,
				reason: "Connection Refused: Controller in recovery mode",
			});
		});
	});
});