factorio-clusterio/packages/controller/src/HostConnection.ts

import type Controller from "./Controller.js";
import type WsServerConnector from "./WsServerConnector.js";

import * as lib from "@clusterio/lib";
import { logger } from "@clusterio/lib";

import BaseConnection from "./BaseConnection.js";
import InstanceRecord from "./InstanceRecord.js";
import HostRecord from "./HostRecord.js";


/**
 * Represents the connection to a host
 *
 * @extends module:controller/src/BaseConnection
 * @alias module:controller/src/HostConnection
 */
export default class HostConnection extends BaseConnection {
	private info: HostRecord;
	plugins: Map<string, string>;

	constructor(
		registerData: lib.RegisterHostData,
		connector: WsServerConnector,
		controller: Controller,
		remoteAddress: string,
	) {
		super(connector, controller);

		this.plugins = new Map(Object.entries(registerData.plugins));

		const previousHostInfo = this._controller.hosts.get(registerData.id);
		this.info = new HostRecord(
			registerData.id,
			previousHostInfo?.name ?? "",
			registerData.version,
			this.plugins,
			true,
			remoteAddress,
			previousHostInfo?.publicAddress,
			previousHostInfo?.tokenValidAfter,
			0,
			false,
		);
		this._controller.hosts.set(this.info);

		this._checkPluginVersions();

		for (let event of ["connect", "drop", "resume", "close"] as const) {
			this.connector.on(event, () => this._controller.hooks.hostConnectionEvent.invoke(this, event));
		}

		this.connector.on("close", () => {
			// Update status to unknown for instances on this host.
			const instances: InstanceRecord[] = [];
			for (let instance of this._controller.instances.valuesMutable()) {
				if (instance.config.get("instance.assigned_host") !== this.id) {
					continue;
				}

				instances.push(instance);
				let prev = instance.status;
				instance.status = "unknown";
				this._controller.hooks.instanceStatusChanged.invoke(instance, prev);
			}
			this._controller.instances.records.setMany(instances);

			this.info.connected = false;
			this._controller.hosts.set(this.info);
		});

		this.handle(lib.GetFactorioCredentialsRequest, this.handleGetFactorioCredentialsRequest.bind(this));
		this.handle(lib.HostInfoUpdateEvent, this.handleHostInfoUpdateEvent.bind(this));
		this.handle(lib.InstanceStatusChangedEvent, this.handleInstanceStatusChangedEvent.bind(this));
		this.handle(lib.InstancesUpdateRequest, this.handleInstancesUpdateRequest.bind(this));
		this.handle(lib.InstanceSaveDetailsUpdatesEvent, this.handleInstanceSaveDetailsUpdatesEvent.bind(this));
		this.handle(lib.LogMessageEvent, this.handleLogMessageEvent.bind(this));
		this.handle(lib.InstancePlayerUpdateEvent, this.handleInstancePlayerUpdateEvent.bind(this));
		this.snoopEvent(lib.InstanceAdminlistUpdateEvent, this.snoopInstanceAdminlistUpdateEvent.bind(this));
		this.snoopEvent(lib.InstanceBanlistUpdateEvent, this.snoopInstanceBanlistUpdateEvent.bind(this));
		this.snoopEvent(lib.InstanceWhitelistUpdateEvent, this.snoopInstanceWhitelistUpdateEvent.bind(this));
	}

	validateIngress(message: lib.MessageRoutable) {
		let origin = this.connector.dst;
		switch (message.src.type) {
			case lib.Address.control:
			case lib.Address.controller:
				throw new lib.InvalidMessage(`Received message with invalid src ${message.src} from ${origin}`);

			case lib.Address.host:
				if (message.src.id !== origin.id) {
					throw new lib.InvalidMessage(
						`Received message with invalid src ${message.src} from ${origin}`
					);
				}
				break;

			case lib.Address.instance:
				let instance = this._controller.instances.get(message.src.id);
				if (!instance || instance.config.get("instance.assigned_host") !== this.id) {
					throw new lib.InvalidMessage(
						`Received message with invalid src ${message.src} from ${origin}`
					);
				}
				break;

			default:
				throw new Error("Should be unreachable");
		}
	}

	_checkPluginVersions() {
		let pluginInfos = new Map(this._controller.pluginInfos.map(i => [i.name, i]));
		for (let [name, version] of this.plugins) {
			let info = pluginInfos.get(name);
			if (!info) {
				logger.warn(
					`Host ${this.info.name} (${this.id}) has plugin ${name} ${version} which the ` +
					"controller does not have",
				);
				continue;
			}

			if (info.version !== version) {
				logger.warn(
					`Host ${this.info.name} (${this.id}) has plugin ${name} ${version} which does not match ` +
					`the version of this plugin on the controller (${info.version})`
				);
			}
		}

		for (let [name, info] of pluginInfos) {
			if (!this.plugins.has(name)) {
				logger.warn(`Host ${this.info.name} (${this.id}) is missing plugin ${name} ${info.version}`);
			}
		}
	}

	async prepareDisconnect() {
		await this._controller.hooks.prepareHostDisconnect.invoke(this);
		return await super.prepareDisconnect();
	}

	/**
	 * ID of the host this connection is connected to
	 *
	 * @returns {number} host ID.
	 */
	get id() {
		return this.info.id;
	}

	/**
	 * Remote address of the host this connection is connected to
	 *
	 * @returns {string} host remote address.
	 */
	get remoteAddress() {
		return this.info.remoteAddress;
	}

	async handleGetFactorioCredentialsRequest() {
		if (this._controller.config.get("controller.share_factorio_credentials_with_hosts")) {
			return {
				username: this._controller.config.get("controller.factorio_username") ?? undefined,
				token: this._controller.config.get("controller.factorio_token") ?? undefined,
			};
		}
		return {};
	}

	async handleHostInfoUpdateEvent(event: lib.HostInfoUpdateEvent) {
		this.info.name = event.update.name;
		this.info.publicAddress = event.update.publicAddress;
		this._controller.hosts.set(this.info);
	}

	async handleInstanceStatusChangedEvent(request: lib.InstanceStatusChangedEvent) {
		let instance = this._controller.instances.getMutable(request.instanceId);

		// It's possible to get updates from an instance that does not exist
		// or is not assigned to the host it originated from if it was
		// reassigned or deleted while the connection to the host it was
		// originally on was down at the time.
		if (!instance || instance.config.get("instance.assigned_host") !== this.id) {
			logger.warn(`Got bogus update for instance id ${request.instanceId}`);
			return;
		}

		// We may receive status changed where the status hasn't changed
		// from our perspective if the connection was down at the time it
		// changed.  Hosts also send status updates on assignInstance which
		// for hacky reason is also used to push config changes and
		// restablish status after a connection loss.
		if (
			instance.status === request.status
			&& instance.gamePort === request.gamePort
		) {
			return;
		}

		let prev = instance.status;
		instance.status = request.status;
		// If the instance has transitioned to running then updated the started at time
		if (request.status === "running" && prev === "starting") {
			instance.startedAtMs = Date.now();
		}

		instance.gamePort = request.gamePort;
		instance.factorioVersion = request.factorioVersion;
		this._controller.instances.records.set(instance);
		logger.verbose(`Instance ${instance.config.get("instance.name")} State: ${instance.status}`);
		await this._controller.hooks.instanceStatusChanged.invoke(instance, prev);
	}

	async handleInstancesUpdateRequest(request: lib.InstancesUpdateRequest) {
		// Push updated instance configs
		for (let instance of this._controller.instances.values()) {
			if (instance.config.get("instance.assigned_host") === this.id) {
				await this.send(
					new lib.InstanceAssignInternalRequest(instance.id, instance.config.toRemote("host"))
				);
			}
		}

		// Assign instances the host has but controller does not
		const instanceUpdates: InstanceRecord[] = [];
		for (let instanceData of request.instances) {
			const instanceConfig = new lib.InstanceConfig("controller");
			instanceConfig.update(instanceData.config, false, "host");

			let controllerInstance = this._controller.instances.getMutable(instanceConfig.get("instance.id"));
			if (controllerInstance) {
				// Check if this instance is assigned somewhere else.
				if (controllerInstance.config.get("instance.assigned_host") !== this.id) {
					await this.send(
						new lib.InstanceUnassignInternalRequest(instanceConfig.get("instance.id"))
					);
					continue;
				}

				// Already have this instance, update state instead
				if (controllerInstance.status !== instanceData.status) {
					let prev = controllerInstance.status;
					controllerInstance.status = instanceData.status;
					controllerInstance.gamePort = instanceData.gamePort;
					controllerInstance.factorioVersion = instanceData.factorioVersion;
					if (instanceData.status === "running" && prev === "starting") {
						controllerInstance.startedAtMs = Date.now();
					}
					instanceUpdates.push(controllerInstance);
					logger.verbose(`Instance ${instanceConfig.get("instance.name")} State: ${instanceData.status}`);
					await this._controller.hooks.instanceStatusChanged.invoke(controllerInstance, prev);
				}
				continue;
			}

			instanceConfig.set("instance.assigned_host", this.id);
			const newInstance = await this._controller.instances.createInstance(instanceConfig, true);
			newInstance.factorioVersion = instanceData.factorioVersion;
			newInstance.gamePort = instanceData.gamePort;
			newInstance.status = instanceData.status;
			instanceUpdates.push(newInstance);
		}

		// We suppress the changes made by createInstance so we must manually make them here
		this._controller.instances.records.setMany(instanceUpdates);
		await Promise.all(instanceUpdates.flatMap(newInstance => [
			this._controller.hooks.instanceStatusChanged.invoke(newInstance, undefined),
			this.send(
				new lib.InstanceAssignInternalRequest(
					newInstance.config.get("instance.id"), newInstance.config.toRemote("host")
				)
			),
		]));

		// Push lists to make sure they are in sync.
		let adminlist: Set<string> = new Set();
		let banlist: Map<string, string> = new Map();
		let whitelist: Set<string> = new Set();

		for (let user of this._controller.users.values()) {
			if (user.isAdmin) {
				adminlist.add(user.id);
			}
			if (user.isBanned) {
				banlist.set(user.id, user.banReason);
			}
			if (user.isWhitelisted) {
				whitelist.add(user.id);
			}
		}

		await this.send(new lib.SyncUserListsEvent(adminlist, banlist, whitelist));
	}

	async handleInstanceSaveDetailsUpdatesEvent(event: lib.InstanceSaveDetailsUpdatesEvent) {
		// A host may only report saves for instances assigned to it. Otherwise
		// any host could inject save records for, or wipe the save list of, an
		// instance on a different host. See handleInstanceStatusChangedEvent.
		const assignedHere = (instanceId: number) => {
			const instance = this._controller.instances.get(instanceId);
			return instance !== undefined && instance.config.get("instance.assigned_host") === this.id;
		};

		if (event.instanceId !== undefined && !assignedHere(event.instanceId)) {
			logger.warn(`Got bogus save updates for instance id ${event.instanceId}`);
			return;
		}
		if (event.updates.some(save => !assignedHere(save.instanceId))) {
			logger.warn(`Got bogus save updates from host id ${this.id}`);
			return;
		}

		const updates: lib.SaveDetails[] = [];
		const deletes: lib.SaveDetails[] = [];
		for (const save of event.updates) {
			const existingSave = this._controller.saves.get(save.id);
			if (existingSave && save.equals(existingSave)) {
				continue;
			}
			updates.push(save);
		}
		if (event.instanceId !== undefined) {
			const updatedSaves = new Set(event.updates.map(s => s.id));
			for (const save of this._controller.saves.values()) {
				if (save.instanceId === event.instanceId && !updatedSaves.has(save.id)) {
					deletes.push(save);
				}
			}
		}
		// Hosts eagerly send updates, which means we may get an update
		// where nothing actually changed.
		this._controller.saves.setMany(updates);
		this._controller.saves.deleteMany(deletes);
	}

	async handleLogMessageEvent(event: lib.LogMessageEvent) {
		this._controller.clusterLogger.log({
			...event.info,
			host_id: this.id,
			host_name: this.info.name,
		});
	}

	async handleInstancePlayerUpdateEvent(event: lib.InstancePlayerUpdateEvent, src: lib.Address) {
		let instanceId = src.id;
		const user = this._controller.users.getOrCreateUser(event.name);

		if (event.type === "join") {
			user.notifyJoin(instanceId);
		} else if (event.type === "leave") {
			user.notifyLeave(instanceId);
		}

		user.instanceStats.set(instanceId, event.stats);

		user.recalculatePlayerStats();
		user.saveRecord();

		let instance = this._controller.instances.get(instanceId)!;
		await this._controller.hooks.playerEvent.invoke(instance, {
			type: event.type,
			name: event.name,
			reason: event.reason,
			stats: event.stats,
		});
	}

	async snoopInstanceAdminlistUpdateEvent(event: lib.InstanceAdminlistUpdateEvent, src: lib.Address) {
		if (src.type !== lib.Address.instance) {
			return;
		}

		const user = this._controller.users.getOrCreateUser(event.name);
		user.set("isAdmin", event.admin);

		const instance = this._controller.instances.get(src.id)!;
		await this._controller.hooks.playerEvent.invoke(instance, {
			type: event.admin ? "promote" : "demote" as lib.PlayerEvent["type"],
			name: event.name,
			stats: user.instanceStats.get(src.id)!,
		});
	}

	async snoopInstanceBanlistUpdateEvent(event: lib.InstanceBanlistUpdateEvent, src: lib.Address) {
		if (src.type !== lib.Address.instance) {
			return;
		}

		const user = this._controller.users.getOrCreateUser(event.name);
		user.isBanned = event.banned;
		user.banReason = event.reason;
		user.saveRecord();

		const instance = this._controller.instances.get(src.id)!;
		await this._controller.hooks.playerEvent.invoke(instance, {
			type: event.banned ? "ban" : "unban" as lib.PlayerEvent["type"],
			name: event.name,
			stats: user.instanceStats.get(src.id)!,
			reason: event.reason,
		});
	}

	async snoopInstanceWhitelistUpdateEvent(event: lib.InstanceWhitelistUpdateEvent, src: lib.Address) {
		if (src.type !== lib.Address.instance) {
			return;
		}

		const user = this._controller.users.getOrCreateUser(event.name);
		user.set("isWhitelisted", event.whitelisted);

		const instance = this._controller.instances.get(src.id)!;
		await this._controller.hooks.playerEvent.invoke(instance, {
			type: event.whitelisted ? "whitelisted" : "unwhitelisted" as lib.PlayerEvent["type"],
			name: event.name,
			stats: user.instanceStats.get(src.id)!,
		});
	}
}