feat: add onTokenSync hook for auth token re-sync (#1001)

* init token sync

* refactor: use existing auth message type for token sync

* refactor(provider): sentToken onOpen and improve testing

* docs(onTokenSync): init onTokenSync in docs

* refactor: close WSconn if onTokenSync throws an err

* refactor: close connection not the ws

Author: 0xb4lamx

docs/server/hooks.md

@@ -34,6 +34,7 @@ By way of illustration, if a user isn’t allowed to connect: Just throw an erro
 | `onConnect`                | When a connection is established          | [Read more](/server/hooks#on-connect)                 |
 | `connected`                | After a connection has been establied     | [Read more](/server/hooks#connected)                  |
 | `onAuthenticate`           | When authentication is required           | [Read more](/server/hooks#on-authenticate)            |
+| `onTokenSync`              | When token synchronization occurs         | [Read more](/server/hooks#on-token-sync)              |
 | `onAwarenessUpdate`        | When awareness changed                    | [Read more](/server/hooks#on-awareness-update)        |
 | `onLoadDocument`           | During the creation of a new document     | [Read more](/server/hooks#on-load-document)           |
 | `afterLoadDocument`        | After a document is created               | [Read more](/server/hooks#after-load-document)        |
@@ -213,6 +214,73 @@ const server = new Server({
 server.listen();
 ```
 
+### onTokenSync
+
+The `onTokenSync` hook is called when the server receives a token synchronization request from a connected provider. This enables the server to validate user tokens during active sessions without requiring a full reconnection.
+
+**Hook payload**
+
+The `data` passed to the `onTokenSync` hook has the following attributes:
+
+```js
+const data = {
+  context: any,
+  document: Doc,
+  documentName: string,
+  instance: Hocuspocus,
+  requestHeaders: IncomingHttpHeaders,
+  requestParameters: URLSearchParams,
+  socketId: string,
+  token: string,
+  connectionConfig: {
+    readOnly: boolean,
+  },
+  connection: Connection,
+};
+```
+
+**Example**
+
+```js
+import { Server } from "@hocuspocus/server";
+
+const server = new Server({
+  async onTokenSync({ token, context, connection }) {
+    // Validate the current token
+    const isValid = await validateToken(token);
+    
+    if (!isValid) {
+      throw new Error("Token has expired or is invalid");
+    }
+    
+    // Update permissions if changed
+    const permissions = await getUserPermissions(context.userId);
+    if (permissions.readOnly !== connection.readOnly) {
+      connection.readOnly = permissions.readOnly;
+    }
+    
+    return { lastTokenSync: new Date() };
+  },
+});
+
+server.listen();
+```
+
+**Provider Usage**
+
+```js
+// Provider sends token to server
+provider.sendToken();
+```
+
+**Server Usage**
+
+```js
+// Server requests token from provider
+const connection = document.connections.values().next().value?.connection;
+connection.requestToken();
+```
+
 ### onAwarenessUpdate
 
 The `onAwarenessUpdate` hooks are called when awareness changed ([Provider Awareness API](/provider/events)).

packages/common/src/auth.ts

@@ -1,7 +1,7 @@
 import * as encoding from "lib0/encoding";
 import * as decoding from "lib0/decoding";
 
-enum AuthMessageType {
+export enum AuthMessageType {
 	Token = 0,
 	PermissionDenied = 1,
 	Authenticated = 2,
@@ -31,12 +31,23 @@ export const writeAuthenticated = (
 	encoding.writeVarString(encoder, scope);
 };
 
+export const writeTokenSyncRequest = (
+	encoder: encoding.Encoder,
+) => {
+	encoding.writeVarUint(encoder, AuthMessageType.Token);
+};
+
 export const readAuthMessage = (
 	decoder: decoding.Decoder,
+	sendToken: () => void,
 	permissionDeniedHandler: (reason: string) => void,
 	authenticatedHandler: (scope: string) => void,
 ) => {
 	switch (decoding.readVarUint(decoder)) {
+		case AuthMessageType.Token: {
+			sendToken();
+			break;
+		}
 		case AuthMessageType.PermissionDenied: {
 			permissionDeniedHandler(decoding.readVarString(decoder));
 			break;

packages/provider/src/HocuspocusProvider.ts

@@ -304,6 +304,21 @@ export class HocuspocusProvider extends EventEmitter {
 		});
 	}
 
+	async sendToken() {
+		let token: string | null;
+		try {
+			token = await this.getToken();
+		} catch (error) {
+			this.permissionDeniedHandler(`Failed to get token during sendToken(): ${error}`);
+			return;
+		}
+
+		this.send(AuthenticationMessage, {
+			token: token ?? "",
+			documentName: this.configuration.name,
+		});
+	}
+
 	documentUpdateHandler(update: Uint8Array, origin: any) {
 		if (origin === this) {
 			return;
@@ -374,20 +389,7 @@ export class HocuspocusProvider extends EventEmitter {
 		this.isAuthenticated = false;
 
 		this.emit("open", { event });
-
-		let token: string | null;
-		try {
-			token = await this.getToken();
-		} catch (error) {
-			this.permissionDeniedHandler(`Failed to get token: ${error}`);
-			return;
-		}
-
-		this.send(AuthenticationMessage, {
-			token: token ?? "",
-			documentName: this.configuration.name,
-		});
-
+		await this.sendToken();
 		this.startSync();
 	}
 

packages/provider/src/MessageReceiver.ts

@@ -117,6 +117,7 @@ export class MessageReceiver {
 
 		readAuthMessage(
 			message.decoder,
+			provider.sendToken.bind(provider),
 			provider.permissionDeniedHandler.bind(provider),
 			provider.authenticatedHandler.bind(provider),
 		);

packages/server/src/ClientConnection.ts

@@ -249,6 +249,26 @@ export class ClientConnection {
 			this.documentConnectionsEstablished.delete(documentName);
 		});
 
+		connection.onTokenSyncCallback(async (payload) => {
+			try {
+				return await this.hooks("onTokenSync", {
+					...hookPayload,
+					...payload,
+					connection,
+					documentName,
+				}, (contextAdditions: any) => {
+					hookPayload.context = {
+						...hookPayload.context,
+						...contextAdditions,
+					};
+				});
+			} catch (err: any) {
+				console.error(err);
+				const error = { ...Unauthorized, ...err };
+				connection.close({ code: error.code, reason: error.reason });
+			}
+		});
+
 		this.documentConnections[documentName] = connection;
 
 		// If the WebSocket has already disconnected (wow, that was fast) – then

packages/server/src/Connection.ts

@@ -9,7 +9,7 @@ import type Document from "./Document.ts";
 import { IncomingMessage } from "./IncomingMessage.ts";
 import { MessageReceiver } from "./MessageReceiver.ts";
 import { OutgoingMessage } from "./OutgoingMessage.ts";
-import type { beforeSyncPayload, onStatelessPayload } from "./types.ts";
+import type { beforeSyncPayload, onTokenSyncPayload, onStatelessPayload } from "./types.ts";
 
 export class Connection {
 	webSocket: WebSocket;
@@ -29,6 +29,7 @@ export class Connection {
 			payload: Pick<beforeSyncPayload, "type" | "payload">,
 		) => Promise.resolve(),
 		statelessCallback: (payload: onStatelessPayload) => Promise.resolve(),
+		onTokenSyncCallback: (payload: Partial<onTokenSyncPayload>) => Promise.resolve(),
 	};
 
 	socketId: string;
@@ -106,6 +107,17 @@ export class Connection {
 		return this;
 	}
 
+  /**
+	 * Set a callback that will be triggered when on token sync message is received
+	 */
+	onTokenSyncCallback(
+		callback: (payload: onTokenSyncPayload) => Promise<void>,
+	): Connection {
+		this.callbacks.onTokenSyncCallback = callback;
+
+		return this;
+	}
+
 	/**
 	 * Send the given message
 	 */
@@ -138,6 +150,15 @@ export class Connection {
 		this.send(message.toUint8Array());
 	}
 
+	/**
+	 * Request current token from the client
+	 */
+	public requestToken(): void {
+		const message = new OutgoingMessage(this.document.name).writeTokenSyncRequest();
+
+		this.send(message.toUint8Array());
+	}
+
 	/**
 	 * Graceful wrapper around the WebSocket close method.
 	 */

packages/server/src/Hocuspocus.ts

@@ -106,6 +106,7 @@ export class Hocuspocus {
 			onConnect: this.configuration.onConnect,
 			connected: this.configuration.connected,
 			onAuthenticate: this.configuration.onAuthenticate,
+			onTokenSync: this.configuration.onTokenSync,
 			onLoadDocument: this.configuration.onLoadDocument,
 			afterLoadDocument: this.configuration.afterLoadDocument,
 			beforeHandleMessage: this.configuration.beforeHandleMessage,

packages/server/src/MessageReceiver.ts

@@ -15,6 +15,7 @@ import type Document from "./Document.ts";
 import type { IncomingMessage } from "./IncomingMessage.ts";
 import { OutgoingMessage } from "./OutgoingMessage.ts";
 import { MessageType } from "./types.ts";
+import { AuthMessageType } from "@hocuspocus/common";
 
 export class MessageReceiver {
 	message: IncomingMessage;
@@ -103,11 +104,19 @@ export class MessageReceiver {
 				break;
 			}
 
-			case MessageType.Auth:
+			case MessageType.Auth: {
+				const authType = message.readVarUint();
+				if (authType === AuthMessageType.Token) {
+					connection?.callbacks.onTokenSyncCallback({
+						token: message.readVarString(),
+					});
+					break;
+  				}
 				console.error(
 					"Received an authentication message on a connection that is already fully authenticated. Probably your provider has been destroyed + recreated really fast.",
 				);
 				break;
+			}
 
 			default:
 				console.error(

packages/server/src/OutgoingMessage.ts

@@ -10,7 +10,7 @@ import type { Awareness } from "y-protocols/awareness";
 import { encodeAwarenessUpdate } from "y-protocols/awareness";
 import { writeSyncStep1, writeUpdate } from "y-protocols/sync";
 
-import { writeAuthenticated, writePermissionDenied } from "@hocuspocus/common";
+import { writeAuthenticated, writePermissionDenied, writeTokenSyncRequest } from "@hocuspocus/common";
 import type Document from "./Document.ts";
 import { MessageType } from "./types.ts";
 
@@ -70,6 +70,16 @@ export class OutgoingMessage {
 		return this;
 	}
 
+	writeTokenSyncRequest(): OutgoingMessage {
+		this.type = MessageType.Auth;
+		this.category = "TokenSync";
+
+		writeVarUint(this.encoder, MessageType.Auth);
+		writeTokenSyncRequest(this.encoder);
+
+		return this;
+	}
+
 	writeAuthenticated(readonly: boolean): OutgoingMessage {
 		this.type = MessageType.Auth;
 		this.category = "Authenticated";

packages/server/src/types.ts

@@ -42,6 +42,7 @@ export interface Extension {
 	onConnect?(data: onConnectPayload): Promise<any>;
 	connected?(data: connectedPayload): Promise<any>;
 	onAuthenticate?(data: onAuthenticatePayload): Promise<any>;
+	onTokenSync?(data: onTokenSyncPayload): Promise<any>;
 	onCreateDocument?(data: onCreateDocumentPayload): Promise<any>;
 	onLoadDocument?(data: onLoadDocumentPayload): Promise<any>;
 	afterLoadDocument?(data: afterLoadDocumentPayload): Promise<any>;
@@ -69,6 +70,7 @@ export type HookName =
 	| "onConnect"
 	| "connected"
 	| "onAuthenticate"
+	| "onTokenSync"
 	| "onCreateDocument"
 	| "onLoadDocument"
 	| "afterLoadDocument"
@@ -93,6 +95,7 @@ export type HookPayloadByName = {
 	onConnect: onConnectPayload;
 	connected: connectedPayload;
 	onAuthenticate: onAuthenticatePayload;
+	onTokenSync: onTokenSyncPayload;
 	onCreateDocument: onCreateDocumentPayload;
 	onLoadDocument: onLoadDocumentPayload;
 	afterLoadDocument: afterLoadDocumentPayload;
@@ -174,6 +177,19 @@ export interface onAuthenticatePayload {
 	connectionConfig: ConnectionConfiguration;
 }
 
+export interface onTokenSyncPayload {
+	context: any;
+	document: Document;
+	documentName: string;
+	instance: Hocuspocus;
+	requestHeaders: IncomingHttpHeaders;
+	requestParameters: URLSearchParams;
+	socketId: string;
+	token: string;
+	connectionConfig: ConnectionConfiguration;
+	connection: Connection;
+}
+
 export interface onCreateDocumentPayload {
 	context: any;
 	documentName: string;