BC-9047 Update ws connection behavior

src/infra/authorization/authorization.service.spec.ts

@@ -1,6 +1,7 @@
 import { createMock, DeepMocked } from '@golevelup/ts-jest';
 import { Test, TestingModule } from '@nestjs/testing';
 import { HttpRequest } from 'uWebSockets.js';
+import { WebSocketCloseCode } from '../../shared/type/websocket-close-code.js';
 import { Logger } from '../logger/index.js';
 import { AuthorizationApi, AuthorizedReponse } from './authorization-api-client/index.js';
 import { AuthorizationService } from './authorization.service.js';
@@ -82,7 +83,7 @@ describe(AuthorizationService.name, () => {
 
 				const expectedResult = {
 					error: {
-						code: 4401,
+						code: WebSocketCloseCode.Unauthorized,
 						reason: 'Unauthorized',
 					},
 					hasWriteAccess: false,
@@ -108,7 +109,7 @@ describe(AuthorizationService.name, () => {
 
 				const expectedResult = {
 					error: {
-						code: 4500,
+						code: WebSocketCloseCode.InternalError,
 						reason: 'RoomId not found',
 					},
 					hasWriteAccess: false,
@@ -133,8 +134,8 @@ describe(AuthorizationService.name, () => {
 				const { req } = setupRequest('roomId', 'other=ABC');
 				const expectedResult = {
 					error: {
-						code: 4500,
-						reason: 'JWT token not found',
+						code: WebSocketCloseCode.Unauthorized,
+						reason: 'JWT not found',
 					},
 					hasWriteAccess: false,
 					room: null,
@@ -162,7 +163,7 @@ describe(AuthorizationService.name, () => {
 
 				const expectedResult = {
 					error: {
-						code: 4500,
+						code: WebSocketCloseCode.InternalError,
 						reason: 'testError',
 					},
 					hasWriteAccess: false,

src/infra/authorization/authorization.service.ts

@@ -1,5 +1,6 @@
 import { Injectable } from '@nestjs/common';
 import { HttpRequest } from 'uWebSockets.js';
+import { WebSocketCloseCode } from '../../shared/type/websocket-close-code.js';
 import { Logger } from '../logger/index.js';
 import {
 	AuthorizationApi,
@@ -28,7 +29,11 @@ export class AuthorizationService {
 
 			response = await this.fetchAuthorization(room, token);
 		} catch (error) {
-			response = this.createErrorResponsePayload(4500, error.message);
+			if (error.message === 'JWT not found') {
+				response = this.createErrorResponsePayload(WebSocketCloseCode.Unauthorized, 'JWT not found');
+			} else {
+				response = this.createErrorResponsePayload(WebSocketCloseCode.InternalError, error.message);
+			}
 		}
 
 		return response;
@@ -45,13 +50,13 @@ export class AuthorizationService {
 	}
 
 	private getToken(req: HttpRequest): string {
-		const jwtToken = this.getCookie(req, 'jwt');
+		const token = this.getCookie(req, 'jwt');
 
-		if (!jwtToken) {
-			throw new Error('JWT token not found');
+		if (!token) {
+			throw new Error('JWT not found');
 		}
 
-		return jwtToken;
+		return token;
 	}
 
 	private async fetchAuthorization(room: string, token: string): Promise<ResponsePayload> {
@@ -80,7 +85,7 @@ export class AuthorizationService {
 
 		const { isAuthorized, userId } = response;
 		if (!isAuthorized) {
-			return this.createErrorResponsePayload(4401, 'Unauthorized');
+			return this.createErrorResponsePayload(WebSocketCloseCode.Unauthorized, 'Unauthorized');
 		}
 
 		return this.createResponsePayload(room, userId);

src/modules/server/api/test/websocket.api.spec.ts

@@ -8,6 +8,7 @@ import { WebsocketProvider } from 'y-websocket';
 import { Doc, encodeStateAsUpdateV2 } from 'yjs';
 import { ResponsePayloadBuilder } from '../../../../infra//authorization/response.builder.js';
 import { AuthorizationService } from '../../../../infra/authorization/index.js';
+import { WebSocketCloseCode } from '../../../../shared/type/websocket-close-code.js';
 import { ServerModule } from '../../server.module.js';
 import { TldrawServerConfig } from '../../tldraw-server.config.js';
 
@@ -184,7 +185,7 @@ describe('Websocket Api Test', () => {
 				const randomString = Math.random().toString(36).substring(7);
 				const room = randomString;
 
-				const errorResponse = ResponsePayloadBuilder.buildWithError(4401, 'Unauthorized');
+				const errorResponse = ResponsePayloadBuilder.buildWithError(WebSocketCloseCode.Unauthorized, 'Unauthorized');
 				authorizationService.hasPermission.mockResolvedValueOnce(errorResponse);
 
 				const { ydoc: client1Doc, provider } = createWsClient(room);
@@ -209,7 +210,7 @@ describe('Websocket Api Test', () => {
 				// @ts-ignore
 				expect(error.reason).toBe('Unauthorized');
 				// @ts-ignore
-				expect(error.code).toBe(4401);
+				expect(error.code).toBe(WebSocketCloseCode.Unauthorized);
 
 				provider.awareness.destroy();
 				provider.destroy();
@@ -246,7 +247,7 @@ describe('Websocket Api Test', () => {
 				// @ts-ignore
 				expect(error.reason).toBe('Missing room or userid');
 				// @ts-ignore
-				expect(error.code).toBe(1008);
+				expect(error.code).toBe(WebSocketCloseCode.InternalError);
 
 				provider.awareness.destroy();
 				provider.destroy();
@@ -281,7 +282,7 @@ describe('Websocket Api Test', () => {
 			// @ts-ignore
 			expect(error.reason).toBe('Internal Server Error');
 			// @ts-ignore
-			expect(error.code).toBe(1011);
+			expect(error.code).toBe(WebSocketCloseCode.InternalError);
 		});
 	});
 });

src/modules/server/api/websocket.gateway.ts

@@ -12,6 +12,7 @@ import { Logger } from '../../../infra/logger/index.js';
 import { MetricsService } from '../../../infra/metrics/index.js';
 import { RedisAdapter } from '../../../infra/redis/index.js';
 import { YRedisClient, YRedisDoc, YRedisService, YRedisUser, YRedisUserFactory } from '../../../infra/y-redis/index.js';
+import { WebSocketCloseCode } from '../../../shared/type/websocket-close-code.js';
 import { REDIS_FOR_SUBSCRIBE_OF_DELETION, UWS } from '../server.const.js';
 import { TldrawServerConfig } from '../tldraw-server.config.js';
 
@@ -21,14 +22,6 @@ interface RequestHeaderInfos {
 	headerWsProtocol: string;
 }
 
-// https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
-enum WebSocketErrorCodes {
-	InternalError = 1011,
-	PolicyViolation = 1008,
-	TldrawPolicyViolation = 4401,
-	TldrawInternalError = 4500,
-}
-
 @Injectable()
 export class WebsocketGateway implements OnModuleInit, OnModuleDestroy {
 	public constructor(
@@ -120,7 +113,7 @@ export class WebsocketGateway implements OnModuleInit, OnModuleDestroy {
 			}
 
 			if (user.room === null || user.userid === null) {
-				ws.end(WebSocketErrorCodes.PolicyViolation, 'Missing room or userid');
+				ws.end(WebSocketCloseCode.InternalError, 'Missing room or userid');
 
 				return;
 			}
@@ -149,7 +142,7 @@ export class WebsocketGateway implements OnModuleInit, OnModuleDestroy {
 			this.yRedisService.ensureLatestContentSubscription(yRedisDoc, user);
 		} catch (error) {
 			this.logger.warning(error);
-			ws.end(WebSocketErrorCodes.InternalError, 'Internal Server Error');
+			ws.end(WebSocketCloseCode.InternalError, 'Internal Server Error');
 		}
 	}
 
@@ -176,7 +169,7 @@ export class WebsocketGateway implements OnModuleInit, OnModuleDestroy {
 			const user = ws.getUserData();
 
 			if (!user.hasWriteAccess || !user.room) {
-				ws.end(WebSocketErrorCodes.TldrawPolicyViolation, 'User has no write access or room is missing');
+				ws.end(WebSocketCloseCode.Unauthorized, 'User has no write access or room is missing');
 
 				return;
 			}
@@ -188,7 +181,7 @@ export class WebsocketGateway implements OnModuleInit, OnModuleDestroy {
 			}
 		} catch (error) {
 			this.logger.warning(error);
-			ws.end(WebSocketErrorCodes.InternalError);
+			ws.end(WebSocketCloseCode.InternalError);
 		}
 	}
 
@@ -207,7 +200,7 @@ export class WebsocketGateway implements OnModuleInit, OnModuleDestroy {
 			MetricsService.openConnectionsGauge.dec();
 		} catch (error) {
 			this.logger.warning(error);
-			ws.end(WebSocketErrorCodes.InternalError);
+			ws.end(WebSocketCloseCode.InternalError);
 		}
 	}
 

src/shared/type/websocket-close-code.ts

@@ -0,0 +1,5 @@
+// https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
+export enum WebSocketCloseCode {
+	InternalError = 1011,
+	Unauthorized = 4401,
+}