diff --git a/src/models/server.ts b/src/models/server.ts index 8639281..65f14ea 100644 --- a/src/models/server.ts +++ b/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 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 { + 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({ 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({ 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 { + 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('Server', ServerSchema); diff --git a/src/provisioning.ts b/src/provisioning.ts index 43a4fc7..1bfe597 100644 --- a/src/provisioning.ts +++ b/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 { $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) { diff --git a/src/services/nasc/routes/ac.ts b/src/services/nasc/routes/ac.ts index 1d16a65..1c5fa11 100644 --- a/src/services/nasc/routes/ac.ts +++ b/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, diff --git a/src/services/nnas/routes/provider.ts b/src/services/nnas/routes/provider.ts index 133acbf..f59b1df 100644 --- a/src/services/nnas/routes/provider.ts +++ b/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()); diff --git a/src/types/mongoose/server.ts b/src/types/mongoose/server.ts index 7984918..9f21429 100644 --- a/src/types/mongoose/server.ts +++ b/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; +} interface IServerQueryHelpers {}