refactor: use existing auth message type for token sync

Author: 0xb4lamx

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

@@ -10,7 +10,6 @@ import { MessageReceiver } from "./MessageReceiver.ts";
 import { MessageSender } from "./MessageSender.ts";
 import { AuthenticationMessage } from "./OutgoingMessages/AuthenticationMessage.ts";
 import { AwarenessMessage } from "./OutgoingMessages/AwarenessMessage.ts";
-import { TokenSyncMessage } from "./OutgoingMessages/TokenSyncMessage.ts";
 import { StatelessMessage } from "./OutgoingMessages/StatelessMessage.ts";
 import { SyncStepOneMessage } from "./OutgoingMessages/SyncStepOneMessage.ts";
 import { UpdateMessage } from "./OutgoingMessages/UpdateMessage.ts";
@@ -306,16 +305,18 @@ export class HocuspocusProvider extends EventEmitter {
 	}
 
 	async sendToken() {
+		let token: string | null;
 		try {
-			const token = await this.getToken();
-
-			this.send(TokenSyncMessage, {
-				token: token ?? "",
-				documentName: this.configuration.name,
-			});
+			token = await this.getToken();
 		} catch (error) {
-			console.error("Failed to getToken() during sendToken():", 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) {

packages/provider/src/MessageReceiver.ts

@@ -38,10 +38,6 @@ export class MessageReceiver {
 				this.applyQueryAwarenessMessage(provider);
 				break;
 
-			case MessageType.TokenSync:
-				this.applyTokenSyncMessage(provider);
-				break;
-
 			case MessageType.Stateless:
 				provider.receiveStateless(readVarString(message.decoder));
 				break;
@@ -121,6 +117,7 @@ export class MessageReceiver {
 
 		readAuthMessage(
 			message.decoder,
+			provider.sendToken.bind(provider),
 			provider.permissionDeniedHandler.bind(provider),
 			provider.authenticatedHandler.bind(provider),
 		);
@@ -139,9 +136,4 @@ export class MessageReceiver {
 			),
 		);
 	}
-
-	private applyTokenSyncMessage(provider: HocuspocusProvider) {
-		// Server is requesting the current token, send it back
-		provider.sendToken();
-	}
 }

packages/provider/src/OutgoingMessages/TokenSyncMessage.ts

@@ -8,7 +8,6 @@ import type { OutgoingMessage } from "./OutgoingMessage.ts";
 import type { AuthenticationMessage } from "./OutgoingMessages/AuthenticationMessage.ts";
 import type { AwarenessMessage } from "./OutgoingMessages/AwarenessMessage.ts";
 import type { QueryAwarenessMessage } from "./OutgoingMessages/QueryAwarenessMessage.ts";
-import type { TokenSyncMessage } from "./OutgoingMessages/TokenSyncMessage.ts";
 import type { SyncStepOneMessage } from "./OutgoingMessages/SyncStepOneMessage.ts";
 import type { SyncStepTwoMessage } from "./OutgoingMessages/SyncStepTwoMessage.ts";
 import type { UpdateMessage } from "./OutgoingMessages/UpdateMessage.ts";
@@ -21,7 +20,6 @@ export enum MessageType {
 	Stateless = 5,
 	CLOSE = 7,
 	SyncStatus = 8,
-	TokenSync = 9,
 }
 
 export enum WebSocketStatus {
@@ -55,7 +53,6 @@ export type ConstructableOutgoingMessage =
 	| Constructable<AuthenticationMessage>
 	| Constructable<AwarenessMessage>
 	| Constructable<QueryAwarenessMessage>
-	| Constructable<TokenSyncMessage>
 	| Constructable<SyncStepOneMessage>
 	| Constructable<SyncStepTwoMessage>
 	| Constructable<UpdateMessage>;

packages/provider/src/types.ts

@@ -252,16 +252,16 @@ export class ClientConnection {
 		connection.onTokenSyncCallback(async (payload) => {
 			try {
 				return await this.hooks("onTokenSync", {
-          ...hookPayload,
-          ...payload,
-          connection,
-          documentName,
-        }, (contextAdditions: any) => {
-          hookPayload.context = {
-            ...hookPayload.context,
-            ...contextAdditions,
-          };
-        });
+					...hookPayload,
+					...payload,
+					connection,
+					documentName,
+				}, (contextAdditions: any) => {
+					hookPayload.context = {
+						...hookPayload.context,
+						...contextAdditions,
+					};
+				});
 			} catch (err: any) {
 				const error = err || Forbidden;
 				const message = new OutgoingMessage(documentName).writePermissionDenied(

packages/server/src/ClientConnection.ts

@@ -154,7 +154,7 @@ export class Connection {
 	 * Request current token from the client
 	 */
 	public requestToken(): void {
-		const message = new OutgoingMessage(this.document.name).writeTokenSync();
+		const message = new OutgoingMessage(this.document.name).writeTokenSyncRequest();
 
 		this.send(message.toUint8Array());
 	}

packages/server/src/Connection.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;
@@ -77,12 +78,6 @@ export class MessageReceiver {
 
 				break;
 			}
-			case MessageType.TokenSync: {
-        connection?.callbacks.onTokenSyncCallback({
-					token: message.readVarString(),
-				});
-				break;
-			}
 			case MessageType.Stateless: {
 				connection?.callbacks.statelessCallback({
 					connection,
@@ -109,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/MessageReceiver.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,11 +70,12 @@ export class OutgoingMessage {
 		return this;
 	}
 
-	writeTokenSync(): OutgoingMessage {
-		this.type = MessageType.TokenSync;
+	writeTokenSyncRequest(): OutgoingMessage {
+		this.type = MessageType.Auth;
 		this.category = "TokenSync";
 
-		writeVarUint(this.encoder, MessageType.TokenSync);
+		writeVarUint(this.encoder, MessageType.Auth);
+		writeTokenSyncRequest(this.encoder);
 
 		return this;
 	}

packages/server/src/OutgoingMessage.ts

@@ -20,7 +20,6 @@ export enum MessageType {
 	BroadcastStateless = 6,
 	CLOSE = 7,
 	SyncStatus = 8,
-	TokenSync = 9,
 }
 
 export interface AwarenessUpdate {
@@ -180,15 +179,15 @@ export interface onAuthenticatePayload {
 
 export interface onTokenSyncPayload {
 	context: any;
-  document: Document;
+	document: Document;
 	documentName: string;
 	instance: Hocuspocus;
 	requestHeaders: IncomingHttpHeaders;
 	requestParameters: URLSearchParams;
 	socketId: string;
 	token: string;
 	connectionConfig: ConnectionConfiguration;
-  connection: Connection;
+	connection: Connection;
 }
 
 export interface onCreateDocumentPayload {