Merge pull request #318 from PretendoNetwork/dev

Release

src/models/server.ts

@@ -1,10 +1,61 @@
+import dgram from 'node:dgram';
+import crypto from 'node:crypto';
 import { Schema, model } from 'mongoose';
 import uniqueValidator from 'mongoose-unique-validator';
-import type { IServer, IServerMethods, ServerModel } from '@/types/mongoose/server';
+import { LOG_WARN } from '@/logger';
+import type { IServer, IServerConnectInfo, IServerMethods, ServerModel } from '@/types/mongoose/server';
+
+// * Kinda ugly to slap this in with the Mongoose stuff but it's fine for now
+// TODO - Maybe move this one day?
+const socket = dgram.createSocket('udp4');
+const pendingHealthCheckRequests = new Map<string, () => void>();
+
+socket.on('message', (msg: Buffer, _rinfo: dgram.RemoteInfo) => {
+	const uuid = msg.toString();
+	const resolve = pendingHealthCheckRequests.get(uuid);
+
+	if (resolve) {
+		resolve();
+	}
+});
+
+socket.bind();
+
+function healthCheck(target: { host: string; port: number }): Promise<string> {
+	return new Promise((resolve, reject) => {
+		const uuid = crypto.randomUUID();
+
+		const timeout = setTimeout(() => {
+			pendingHealthCheckRequests.delete(uuid);
+			reject(new Error('No valid response received'));
+		}, 2 * 1000); // TODO - Make this configurable? 2 seconds seems fine for now
+
+		pendingHealthCheckRequests.set(uuid, () => {
+			clearTimeout(timeout);
+			pendingHealthCheckRequests.delete(uuid);
+			resolve(target.host);
+		});
+
+		socket.send(Buffer.from(uuid), target.port, target.host, (error) => {
+			if (error) {
+				clearTimeout(timeout);
+				pendingHealthCheckRequests.delete(uuid);
+				reject(error);
+			}
+		});
+	});
+}
 
 const ServerSchema = new Schema<IServer, ServerModel, IServerMethods>({
 	client_id: String,
-	ip: String,
+	ip: {
+		type: String,
+		required: false
+	},
+	ip_list: {
+		type: [String],
+		required: false
+	},
 	port: Number,
 	service_name: String,
 	service_type: String,
@@ -13,9 +64,60 @@ const ServerSchema = new Schema<IServer, ServerModel, IServerMethods>({
 	access_mode: String,
 	maintenance_mode: Boolean,
 	device: Number,
-	aes_key: String
+	aes_key: String,
+	health_check_port: {
+		type: Number,
+		required: false
+	}
 });
 
 ServerSchema.plugin(uniqueValidator, { message: '{PATH} already in use.' });
 
+ServerSchema.method('getServerConnectInfo', async function (): Promise<IServerConnectInfo> {
+	const ipList = [this.ip_list, this.ip].flat().filter((v): v is string => !!v);
+	if (ipList.length === 0) {
+		throw new Error(`No IP configured for server ${this._id}`);
+	}
+
+	const randomIP = ipList[Math.floor(Math.random() * ipList.length)];
+
+	if (!this.health_check_port) {
+		return {
+			ip: randomIP,
+			port: this.port
+		};
+	}
+
+	// * Remove the random IP from the race pool to remove the duplicate health check
+	const healthCheckTargets = ipList.filter(ip => ip !== randomIP).map(ip => ({
+		host: ip,
+		port: this.health_check_port!
+	}));
+
+	// * Default to the random IP in case nothing responded in time
+	// * and just Hope For The Best:tm:
+	let target = randomIP;
+
+	// * Check the random IP and start the race at the same time, preferring
+	// * the result of the random IP should it succeed. Worst case scenario
+	// * this takes 2 seconds to complete
+	const [randomResult, raceResult] = await Promise.allSettled([
+		healthCheck({ host: randomIP, port: this.health_check_port! }),
+		Promise.race(healthCheckTargets.map(target => healthCheck(target)))
+	]);
+
+	if (randomResult.status === 'rejected') {
+		if (raceResult.status === 'fulfilled') {
+			target = raceResult.value;
+		} else {
+			LOG_WARN(`Server ${this.service_name} failed to find healthy NEX server. Using the randomly selected IP ${target}`);
+		}
+	}
+
+	return {
+		ip: target,
+		port: this.port
+	};
+});
+
 export const Server = model<IServer, ServerModel>('Server', ServerSchema);

src/provisioning.ts

@@ -13,8 +13,10 @@ const serverProvisioningSchema = z.object({
 	servers: z.array(z.object({
 		id: z.string(),
 		name: z.string(),
-		ip: z.string(),
-		port: z.coerce.number()
+		ip: z.string().optional(),
+		ipList: z.array(z.string()).optional(),
+		port: z.coerce.number(),
+		health_check_port: z.coerce.number().optional()
 	}))
 });
 
@@ -40,8 +42,10 @@ export async function handleServerProvisioning(): Promise<void> {
 			$set: {
 				_id: id,
 				service_name: server.name,
+				ipList: server.ipList,
 				ip: server.ip,
-				port: server.port
+				port: server.port,
+				health_check_port: server.health_check_port
 			}
 		});
 		if (!result) {

src/services/nasc/routes/ac.ts

@@ -51,7 +51,8 @@ router.post('/', async (request: express.Request, response: express.Response): P
 		return;
 	}
 
-	if (action === 'LOGIN' && server.port <= 0 && server.ip !== '0.0.0.0') {
+	const connectInfo = await server.getServerConnectInfo();
+	if (action === 'LOGIN' && connectInfo.port <= 0 && connectInfo.ip !== '0.0.0.0') {
 		// * Addresses of 0.0.0.0:0 are allowed
 		// * They are expected for titles with no NEX server
 		response.status(200).send(nascError('110').toString());
@@ -85,8 +86,9 @@ async function processLoginRequest(server: HydratedServerDocument, pid: number,
 	const nexTokenBuffer = await generateToken(server.aes_key, tokenOptions);
 	const nexToken = nintendoBase64Encode(nexTokenBuffer || '');
 
+	const connectInfo = await server.getServerConnectInfo();
 	return new URLSearchParams({
-		locator: nintendoBase64Encode(`${server.ip}:${server.port}`),
+		locator: nintendoBase64Encode(`${connectInfo.ip}:${connectInfo.port}`),
 		retry: nintendoBase64Encode('0'),
 		returncd: nintendoBase64Encode('001'),
 		token: nexToken,

src/services/nnas/routes/provider.ts

@@ -229,12 +229,13 @@ router.get('/nex_token/@me', async (request: express.Request, response: express.
 		nexToken = Buffer.from(nexToken || '', 'base64').toString('hex');
 	}
 
+	const connectInfo = await server.getServerConnectInfo();
 	response.send(xmlbuilder.create({
 		nex_token: {
-			host: server.ip,
+			host: connectInfo.ip,
 			nex_password: nexAccount.password,
 			pid: nexAccount.pid,
-			port: server.port,
+			port: connectInfo.port,
 			token: nexToken
 		}
 	}).end());

src/types/mongoose/server.ts

@@ -2,7 +2,8 @@ import type { Model, HydratedDocument } from 'mongoose';
 
 export interface IServer {
 	client_id: string;
-	ip: string;
+	ip?: string;
+	ip_list?: string[];
 	port: number;
 	service_name: string;
 	service_type: string;
@@ -12,9 +13,17 @@ export interface IServer {
 	maintenance_mode: boolean;
 	device: number;
 	aes_key: string;
+	health_check_port?: number;
 }
 
-export interface IServerMethods {}
+export interface IServerConnectInfo {
+	ip: string;
+	port: number;
+}
+
+export interface IServerMethods {
+	getServerConnectInfo(): Promise<IServerConnectInfo>;
+}
 
 interface IServerQueryHelpers {}