From 10dd3c2489c5759d479fc9b963bb9d0e285e1fab Mon Sep 17 00:00:00 2001
From: Elijah R <elijah@elijahr.dev>
Date: Fri, 21 Mar 2025 20:29:16 -0400
Subject: [PATCH] cvmts: Move initial protocol selection to transport layer

---
 cvmts/src/User.ts            |  4 ++--
 cvmts/src/net/ws/WSServer.ts | 14 ++++++++++----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/cvmts/src/User.ts b/cvmts/src/User.ts
index 45111a6..e3d2054 100644
--- a/cvmts/src/User.ts
+++ b/cvmts/src/User.ts
@@ -37,7 +37,7 @@ export class User {
 
 	private logger = pino({ name: 'CVMTS.User' });
 
-	constructor(socket: NetworkClient, ip: IPData, config: IConfig, username?: string, node?: string) {
+	constructor(socket: NetworkClient, protocol: string, ip: IPData, config: IConfig, username?: string, node?: string) {
 		this.IP = ip;
 		this.connectedToNode = false;
 		this.viewMode = -1;
@@ -47,7 +47,7 @@ export class User {
 		this.Capabilities = new CollabVMCapabilities();
 
 		// All clients default to the Guacamole protocol.
-		this.protocol = TheProtocolManager.createProtocol('guacamole', this);
+		this.protocol = TheProtocolManager.createProtocol(protocol, this);
 
 		this.socket.on('disconnect', () => {
 			// Unref the ip data for this connection
diff --git a/cvmts/src/net/ws/WSServer.ts b/cvmts/src/net/ws/WSServer.ts
index 8436e63..3fd7b67 100644
--- a/cvmts/src/net/ws/WSServer.ts
+++ b/cvmts/src/net/ws/WSServer.ts
@@ -11,6 +11,10 @@ import { User } from '../../User.js';
 import pino from 'pino';
 import { BanManager } from '../../BanManager.js';
 
+const kAllowedProtocols = [
+	"guacamole" // Regular ol' collabvm1 protocol
+]
+
 export default class WSServer extends EventEmitter implements NetworkServer {
 	private httpServer: http.Server;
 	private wsServer: WebSocketServer;
@@ -50,7 +54,9 @@ export default class WSServer extends EventEmitter implements NetworkServer {
 			socket.destroy();
 		};
 
-		if (req.headers['sec-websocket-protocol'] !== 'guacamole') {
+		let protocol = req.headers['sec-websocket-protocol'];
+
+		if (!protocol || kAllowedProtocols.indexOf(protocol) === -1) {
 			killConnection();
 			return;
 		}
@@ -131,14 +137,14 @@ export default class WSServer extends EventEmitter implements NetworkServer {
 
 		this.wsServer.handleUpgrade(req, socket, head, (ws: WebSocket) => {
 			this.wsServer.emit('connection', ws, req);
-			this.onConnection(ws, req, ip);
+			this.onConnection(ws, req, ip, protocol);
 		});
 	}
 
-	private onConnection(ws: WebSocket, req: http.IncomingMessage, ip: string) {
+	private onConnection(ws: WebSocket, req: http.IncomingMessage, ip: string, protocol: string) {
 		let client = new WSClient(ws, ip);
 		this.clients.push(client);
-		let user = new User(client, IPDataManager.GetIPData(ip), this.Config);
+		let user = new User(client, protocol, IPDataManager.GetIPData(ip), this.Config);
 
 		this.emit('connect', user);