refactor: Prepare websocket session and ready data in api module

backend/src/JobContext.ts

@@ -9,6 +9,7 @@ import { Jobs, Modules } from "./types/Modules";
 import { StationType } from "./schemas/station";
 import { UserRole, UserSchema } from "./schemas/user";
 import permissions from "./permissions";
+import { Models } from "./types/Models";
 
 export default class JobContext {
 	public readonly job: Job;
@@ -73,12 +74,16 @@ export default class JobContext {
 		}).execute();
 	}
 
+	public async getModel(model: keyof Models) {
+		return this.executeJob("data", "getModel", model);
+	}
+
 	public async getUser(refresh = false) {
 		if (!this.session?.userId) throw new Error("No user found for session");
 
 		if (this.user && !refresh) return this.user;
 
-		const User = await this.executeJob("data", "getModel", "user");
+		const User = await this.getModel("user");
 
 		this.user = await User.findById(this.session.userId);
 
@@ -102,11 +107,7 @@ export default class JobContext {
 		const roles: (UserRole | "owner" | "dj")[] = [user.role];
 
 		if (scope?.stationId) {
-			const Station = await this.executeJob(
-				"data",
-				"getModel",
-				"station"
-			);
+			const Station = await this.getModel("station");
 
 			const station = await Station.findById(scope.stationId);
 
@@ -115,7 +116,7 @@ export default class JobContext {
 				station.owner === this.session.userId
 			)
 				roles.push("owner");
-			if (station.djs.find(dj => dj === this.session?.userId))
+			else if (station.djs.find(dj => dj === this.session?.userId))
 				roles.push("dj");
 		}
 

backend/src/modules/APIModule.ts

@@ -1,6 +1,10 @@
+import config from "config";
+import { isObjectIdOrHexString } from "mongoose";
+import { IncomingMessage } from "node:http";
 import JobContext from "../JobContext";
 import BaseModule from "../BaseModule";
 import { Jobs, Modules, UniqueMethods } from "../types/Modules";
+import WebSocket from "../WebSocket";
 
 export default class APIModule extends BaseModule {
 	/**
@@ -49,31 +53,116 @@ export default class APIModule extends BaseModule {
 			moduleName,
 			jobName,
 			payload,
-			socketId,
 			sessionId
 		}: {
 			moduleName: ModuleNameType;
 			jobName: JobNameType;
 			payload: PayloadType;
-			socketId?: string;
 			sessionId?: string;
 		}
 	): Promise<ReturnType> {
 		let session;
 		if (sessionId) {
-			const Session = await context.executeJob(
-				"data",
-				"getModel",
-				"session"
-			);
-			session = await Session.findOneAndUpdate(
-				{ _id: sessionId },
-				{ $addToSet: { socketIds: socketId } }
-			);
+			const Session = await context.getModel("session");
+
+			session = await Session.findByIdAndUpdate(sessionId, {
+				updatedAt: Date.now()
+			});
 		}
 
 		return context.executeJob(moduleName, jobName, payload, { session });
 	}
+
+	/**
+	 * getCookieValueFromHeader - Get value of a cookie from cookie header string
+	 */
+	private getCookieValueFromHeader(cookieName: string, header: string) {
+		const cookie = header
+			.split("; ")
+			.find(
+				cookie =>
+					cookie.substring(0, cookie.indexOf("=")) === cookieName
+			);
+
+		return cookie?.substring(cookie.indexOf("=") + 1, cookie.length);
+	}
+
+	/**
+	 * prepareWebsocket - Prepare websocket connection
+	 */
+	public async prepareWebsocket(
+		context: JobContext,
+		{ socket, request }: { socket: WebSocket; request: IncomingMessage }
+	) {
+		const socketId = request.headers["sec-websocket-key"];
+		socket.setSocketId(socketId);
+
+		let sessionId = request.headers.cookie
+			? this.getCookieValueFromHeader(
+					config.get<string>("cookie"),
+					request.headers.cookie
+			  )
+			: undefined;
+
+		if (sessionId && isObjectIdOrHexString(sessionId))
+			socket.setSessionId(sessionId);
+		else sessionId = undefined;
+
+		let user;
+		if (sessionId) {
+			const Session = await context.getModel("session");
+
+			const session = await Session.findByIdAndUpdate(sessionId, {
+				updatedAt: Date.now()
+			});
+
+			if (session) {
+				context.setSession(session);
+
+				user = await context.getUser().catch(() => undefined);
+			}
+		}
+
+		return {
+			config: {
+				cookie: config.get("cookie"),
+				sitename: config.get("sitename"),
+				recaptcha: {
+					enabled: config.get("apis.recaptcha.enabled"),
+					key: config.get("apis.recaptcha.key")
+				},
+				githubAuthentication: config.get("apis.github.enabled"),
+				messages: config.get("messages"),
+				christmas: config.get("christmas"),
+				footerLinks: config.get("footerLinks"),
+				shortcutOverrides: config.get("shortcutOverrides"),
+				registrationDisabled: config.get("registrationDisabled"),
+				mailEnabled: config.get("mail.enabled"),
+				discogsEnabled: config.get("apis.discogs.enabled"),
+				experimental: {
+					changable_listen_mode: config.get(
+						"experimental.changable_listen_mode"
+					),
+					media_session: config.get("experimental.media_session"),
+					disable_youtube_search: config.get(
+						"experimental.disable_youtube_search"
+					),
+					station_history: config.get("experimental.station_history"),
+					soundcloud: config.get("experimental.soundcloud"),
+					spotify: config.get("experimental.spotify")
+				}
+			},
+			user: user
+				? {
+						loggedIn: true,
+						role: user.role,
+						username: user.username,
+						email: user.email.address,
+						userId: user._id
+				  }
+				: { loggedIn: false }
+		};
+	}
 }
 
 export type APIModuleJobs = {

backend/src/modules/WebSocketModule.ts

@@ -2,9 +2,12 @@ import config from "config";
 import express from "express";
 import http, { Server, IncomingMessage } from "node:http";
 import { RawData, WebSocketServer } from "ws";
+import { Types } from "mongoose";
 import BaseModule from "../BaseModule";
 import { UniqueMethods } from "../types/Modules";
 import WebSocket from "../WebSocket";
+import JobContext from "../JobContext";
+import Job from "../Job";
 
 export default class WebSocketModule extends BaseModule {
 	private httpServer?: Server;
@@ -87,22 +90,15 @@ export default class WebSocketModule extends BaseModule {
 			return;
 		}
 
-		socket.log({ type: "debug", message: "WebSocket #ID connected" });
+		const readyData = await new Job("prepareWebsocket", "api", {
+			socket,
+			request
+		}).execute();
 
-		socket.setSocketId(request.headers["sec-websocket-key"]);
-
-		const sessionCookie = request.headers.cookie
-			?.split("; ")
-			.find(
-				cookie =>
-					cookie.substring(0, cookie.indexOf("=")) ===
-					config.get("cookie")
-			);
-		const sessionId = sessionCookie?.substring(
-			sessionCookie.indexOf("=") + 1,
-			sessionCookie.length
-		);
-		socket.setSessionId(sessionId);
+		socket.log({
+			type: "debug",
+			message: `WebSocket opened #${socket.getSocketId()}`
+		});
 
 		socket.on("error", error =>
 			socket.log({
@@ -112,41 +108,12 @@ export default class WebSocketModule extends BaseModule {
 			})
 		);
 
-		socket.on("close", () =>
-			socket.log({ type: "debug", message: "WebSocket #ID closed" })
-		);
-
-		const readyData = {
-			config: {
-				cookie: config.get("cookie"),
-				sitename: config.get("sitename"),
-				recaptcha: {
-					enabled: config.get("apis.recaptcha.enabled"),
-					key: config.get("apis.recaptcha.key")
-				},
-				githubAuthentication: config.get("apis.github.enabled"),
-				messages: config.get("messages"),
-				christmas: config.get("christmas"),
-				footerLinks: config.get("footerLinks"),
-				shortcutOverrides: config.get("shortcutOverrides"),
-				registrationDisabled: config.get("registrationDisabled"),
-				mailEnabled: config.get("mail.enabled"),
-				discogsEnabled: config.get("apis.discogs.enabled"),
-				experimental: {
-					changable_listen_mode: config.get(
-						"experimental.changable_listen_mode"
-					),
-					media_session: config.get("experimental.media_session"),
-					disable_youtube_search: config.get(
-						"experimental.disable_youtube_search"
-					),
-					station_history: config.get("experimental.station_history"),
-					soundcloud: config.get("experimental.soundcloud"),
-					spotify: config.get("experimental.spotify")
-				}
-			},
-			user: { loggedIn: false }
-		};
+		socket.on("close", async () => {
+			socket.log({
+				type: "debug",
+				message: `WebSocket closed #${socket.getSocketId()}`
+			});
+		});
 
 		socket.dispatch("ready", readyData);
 
@@ -176,7 +143,6 @@ export default class WebSocketModule extends BaseModule {
 				moduleName,
 				jobName,
 				payload,
-				socketId: socket.getSocketId(),
 				sessionId: socket.getSessionId()
 			});
 
@@ -190,6 +156,40 @@ export default class WebSocketModule extends BaseModule {
 		}
 	}
 
+	/**
+	 * getSockets - Get websocket clients
+	 */
+	public async getSockets(context: JobContext) {
+		return this.wsServer?.clients;
+	}
+
+	/**
+	 * getSocket - Get websocket client
+	 */
+	public async getSocket(
+		context: JobContext,
+		{
+			socketId,
+			sessionId
+		}: { socketId?: string; sessionId?: Types.ObjectId }
+	) {
+		if (!this.wsServer) return null;
+
+		for (const clients of this.wsServer.clients.entries() as IterableIterator<
+			[WebSocket, WebSocket]
+		>) {
+			const socket = clients.find(socket => {
+				if (socket.getSocketId() === socketId) return true;
+				if (socket.getSessionId() === sessionId) return true;
+				return false;
+			});
+
+			if (socket) return socket;
+		}
+
+		return null;
+	}
+
 	/**
 	 * shutdown - Shutdown websocket module
 	 */

backend/src/schemas/session.ts

@@ -3,7 +3,6 @@ import { BaseSchema } from "../types/Schemas";
 
 export interface SessionSchema extends BaseSchema {
 	userId: Types.ObjectId;
-	socketIds: string[];
 }
 
 export type SessionModel = Model<SessionSchema>;
@@ -13,8 +12,7 @@ export const schema = new Schema<SessionSchema, SessionModel>({
 		type: SchemaTypes.ObjectId,
 		ref: "user",
 		required: true
-	},
-	socketIds: [SchemaTypes.String]
+	}
 });
 
 export type SessionSchemaType = typeof schema;

backend/src/schemas/user.ts

@@ -95,7 +95,8 @@ export const schema = new Schema<UserSchema, UserModel>(
 			},
 			verificationToken: {
 				type: SchemaTypes.String,
-				required: false
+				required: false,
+				select: false
 			}
 		},
 		avatar: {
@@ -115,29 +116,62 @@ export const schema = new Schema<UserSchema, UserModel>(
 			}
 		},
 		services: {
-			password: {
-				password: SchemaTypes.String,
-				reset: {
-					code: {
-						type: SchemaTypes.String,
-						minLength: 8,
-						maxLength: 8
+			type: {
+				password: {
+					type: {
+						password: {
+							type: SchemaTypes.String,
+							required: true,
+							select: false
+						},
+						reset: {
+							code: {
+								type: SchemaTypes.String,
+								minLength: 8,
+								maxLength: 8,
+								required: false,
+								select: false
+							},
+							expires: {
+								type: SchemaTypes.Date,
+								required: false,
+								select: false
+							}
+						},
+						set: {
+							code: {
+								type: SchemaTypes.String,
+								minLength: 8,
+								maxLength: 8,
+								required: false,
+								select: false
+							},
+							expires: {
+								type: SchemaTypes.Date,
+								required: false,
+								select: false
+							}
+						}
 					},
-					expires: { type: SchemaTypes.Date }
+					required: false
 				},
-				set: {
-					code: {
-						type: SchemaTypes.String,
-						minLength: 8,
-						maxLength: 8
+				github: {
+					type: {
+						id: {
+							type: SchemaTypes.Number,
+							required: true,
+							select: false
+						},
+						access_token: {
+							type: SchemaTypes.String,
+							required: true,
+							select: false
+						}
 					},
-					expires: { type: SchemaTypes.Date }
+					required: false
 				}
 			},
-			github: {
-				id: SchemaTypes.Number,
-				access_token: SchemaTypes.String
-			}
+			required: true
 		},
 		statistics: {
 			songsRequested: {