From 988170dcf6df315f56b627d992e61005618ba739 Mon Sep 17 00:00:00 2001 From: nav Date: Wed, 8 Apr 2026 23:24:33 +0000 Subject: [PATCH] YNX-1003: Implement single-identifier single-active-connection policy - Refactor transport to track temp connections separately from authenticated - Add assignIdentifierToTemp() for hello phase (pre-auth) - Add promoteToAuthenticated() that closes old connection only after new one auths - Add removeTempConnection() for cleanup on auth failure - Update runtime to use new API: assignIdentifierToTemp() on hello, promoteToAuthenticated() on auth_success This prevents an attacker from kicking an authenticated connection with just a hello message. --- plugin/core/runtime.ts | 6 +-- plugin/core/transport.ts | 104 ++++++++++++++++++++++++++++----------- 2 files changed, 79 insertions(+), 31 deletions(-) diff --git a/plugin/core/runtime.ts b/plugin/core/runtime.ts index e72612d..8c8dceb 100644 --- a/plugin/core/runtime.ts +++ b/plugin/core/runtime.ts @@ -254,7 +254,7 @@ export class YonexusServerRuntime { const record = this.ensureClientRecord(helloIdentifier); record.updatedAt = this.now(); - this.options.transport.registerConnection(helloIdentifier, connection.ws); + this.options.transport.assignIdentifierToTemp(connection.ws, helloIdentifier); this.registry.sessions.set(helloIdentifier, { identifier: helloIdentifier, socket: connection.ws, @@ -361,7 +361,7 @@ export class YonexusServerRuntime { } if (connection.identifier !== identifier) { - this.options.transport.registerConnection(identifier, connection.ws); + this.options.transport.assignIdentifierToTemp(connection.ws, identifier); } const session = this.registry.sessions.get(identifier); @@ -540,7 +540,7 @@ export class YonexusServerRuntime { session.lastActivityAt = now; session.publicKey = publicKey; } - this.options.transport.markAuthenticated(identifier); + this.options.transport.promoteToAuthenticated(identifier, connection.ws); this.options.transport.sendToConnection( { ...connection, identifier }, encodeBuiltin( diff --git a/plugin/core/transport.ts b/plugin/core/transport.ts index a2834eb..e5409b3 100644 --- a/plugin/core/transport.ts +++ b/plugin/core/transport.ts @@ -9,6 +9,17 @@ export interface ClientConnection { isAuthenticated: boolean; } +/** + * Temporary connection tracking before authentication. + * Connections remain in this state until successfully authenticated. + */ +interface TempConnection { + readonly ws: WebSocket; + readonly connectedAt: number; + /** The identifier claimed during hello, if any */ + assignedIdentifier: string | null; +} + export interface ServerTransport { readonly isRunning: boolean; readonly connections: ReadonlyMap; @@ -18,8 +29,23 @@ export interface ServerTransport { sendToConnection(connection: ClientConnection, message: string): boolean; broadcast(message: string): void; closeConnection(identifier: string, code?: number, reason?: string): boolean; - registerConnection(identifier: string, ws: WebSocket): boolean; - markAuthenticated(identifier: string): boolean; + /** + * Promote a temp connection to authenticated status. + * This implements the single-identifier-single-active-connection policy: + * - If another authenticated connection exists for this identifier, it is closed + * - The connection is moved from temp to authenticated registry + */ + promoteToAuthenticated(identifier: string, ws: WebSocket): boolean; + /** + * Remove a temp connection without promoting it. + * Called when authentication fails or connection closes before auth. + */ + removeTempConnection(ws: WebSocket): void; + /** + * Assign an identifier to a temp connection during hello processing. + * This does NOT register the connection as authenticated yet. + */ + assignIdentifierToTemp(ws: WebSocket, identifier: string): void; } export type MessageHandler = (connection: ClientConnection, message: string) => void; @@ -36,7 +62,7 @@ export interface ServerTransportOptions { export class YonexusServerTransport implements ServerTransport { private wss: WebSocketServer | null = null; private _connections = new Map(); - private tempConnections = new Set(); + private tempConnections = new Map(); private options: ServerTransportOptions; private _isRunning = false; @@ -87,14 +113,15 @@ export class YonexusServerTransport implements ServerTransport { return; } - // Close all connections + // Close all authenticated connections for (const conn of this._connections.values()) { conn.ws.close(1000, "Server shutting down"); } this._connections.clear(); - for (const ws of this.tempConnections) { - ws.close(1000, "Server shutting down"); + // Close all temp connections + for (const temp of this.tempConnections.values()) { + temp.ws.close(1000, "Server shutting down"); } this.tempConnections.clear(); @@ -144,39 +171,54 @@ export class YonexusServerTransport implements ServerTransport { return true; } - registerConnection(identifier: string, ws: WebSocket): boolean { - // Check if already connected - if (this._connections.has(identifier)) { - // Replace old connection - const oldConn = this._connections.get(identifier)!; - oldConn.ws.close(1008, "New connection established"); + promoteToAuthenticated(identifier: string, ws: WebSocket): boolean { + // Verify the connection exists in temp connections + const tempConn = this.tempConnections.get(ws); + if (!tempConn) { + return false; } - // Remove from temp connections if present + // Check if already have an authenticated connection for this identifier + // If so, close it (single-identifier-single-active-connection policy) + const existingConn = this._connections.get(identifier); + if (existingConn) { + existingConn.ws.close(1008, "Connection replaced by new authenticated session"); + this._connections.delete(identifier); + } + + // Remove from temp connections this.tempConnections.delete(ws); + // Register the new authenticated connection const conn: ClientConnection = { identifier, ws, - connectedAt: Math.floor(Date.now() / 1000), - isAuthenticated: false + connectedAt: tempConn.connectedAt, + isAuthenticated: true }; this._connections.set(identifier, conn); return true; } - markAuthenticated(identifier: string): boolean { - const conn = this._connections.get(identifier); - if (!conn) { - return false; + removeTempConnection(ws: WebSocket): void { + this.tempConnections.delete(ws); + } + + assignIdentifierToTemp(ws: WebSocket, identifier: string): void { + const tempConn = this.tempConnections.get(ws); + if (tempConn) { + tempConn.assignedIdentifier = identifier; } - conn.isAuthenticated = true; - return true; } private handleConnection(ws: WebSocket, _req: import("http").IncomingMessage): void { - this.tempConnections.add(ws); + // Store as temp connection until authenticated + this.tempConnections.set(ws, { + ws, + connectedAt: Math.floor(Date.now() / 1000), + assignedIdentifier: null + }); const tempConn: ClientConnection = { identifier: null, @@ -187,12 +229,18 @@ export class YonexusServerTransport implements ServerTransport { ws.on("message", (data: RawData) => { const message = data.toString("utf8"); - // Try to get identifier from registered connections + // Try to get identifier from temp connections first, then authenticated connections let identifier: string | null = null; - for (const [id, conn] of this._connections) { - if (conn.ws === ws) { - identifier = id; - break; + const tempData = this.tempConnections.get(ws); + if (tempData) { + identifier = tempData.assignedIdentifier; + } + if (!identifier) { + for (const [id, conn] of this._connections) { + if (conn.ws === ws) { + identifier = id; + break; + } } } @@ -203,7 +251,7 @@ export class YonexusServerTransport implements ServerTransport { ws.on("close", (code: number, reason: Buffer) => { this.tempConnections.delete(ws); - // Find and remove from registered connections + // Find and remove from authenticated connections for (const [id, conn] of this._connections) { if (conn.ws === ws) { this._connections.delete(id);