Add 404 for drawing authorization proccess

apps/server/src/modules/tldraw/controller/api-test/tldraw.ws.api.spec.ts

@@ -2,7 +2,7 @@ import { WsAdapter } from '@nestjs/platform-ws';
 import { Test } from '@nestjs/testing';
 import WebSocket from 'ws';
 import { TextEncoder } from 'util';
-import { INestApplication } from '@nestjs/common';
+import { INestApplication, NotFoundException } from '@nestjs/common';
 import { throwError } from 'rxjs';
 import { createMock, DeepMocked } from '@golevelup/ts-jest';
 import { HttpService } from '@nestjs/axios';
@@ -151,7 +151,7 @@ describe('WebSocketController (WsAdapter)', () => {
 			const error = new Error('unknown error');
 
 			httpGetCallSpy.mockReturnValueOnce(throwError(() => error));
-			ws = await TestConnection.setupWs(wsUrl, 'TEST', { jwt: 'mock-wrong-jwt' });
+			ws = await TestConnection.setupWs(wsUrl, 'TEST', { headers: { cookie: { jwt: 'jwt-mocked' } } });
 
 			expect(wsCloseSpy).toHaveBeenCalledWith(
 				WsCloseCodeEnum.WS_CLIENT_UNAUTHORISED_CONNECTION_CODE,
@@ -179,7 +179,7 @@ describe('WebSocketController (WsAdapter)', () => {
 			const { setupConnectionSpy, wsCloseSpy } = setup();
 			const { buffer } = getMessage();
 
-			ws = await TestConnection.setupWs(wsUrl, '', { jwt: 'jwt-mocked' });
+			ws = await TestConnection.setupWs(wsUrl, '', { headers: { cookie: { jwt: 'jwt-mocked' } } });
 			ws.send(buffer);
 
 			expect(wsCloseSpy).toHaveBeenCalledWith(
@@ -192,6 +192,25 @@ describe('WebSocketController (WsAdapter)', () => {
 			ws.close();
 		});
 
+		it(`should close for not existing docName resource`, async () => {
+			const { setupConnectionSpy, wsCloseSpy } = setup();
+			const authorizeConnectionSpy = jest.spyOn(wsService, 'authorizeConnection');
+			authorizeConnectionSpy.mockImplementationOnce(() => {
+				throw new NotFoundException('Resource not found');
+			});
+			ws = await TestConnection.setupWs(wsUrl, 'GLOBAL', { headers: { cookie: { jwt: 'jwt-mocked' } } });
+
+			expect(wsCloseSpy).toHaveBeenCalledWith(
+				WsCloseCodeEnum.WS_CLIENT_NOT_FOUND_CODE,
+				WsCloseMessageEnum.WS_CLIENT_NOT_FOUND_MESSAGE
+			);
+
+			authorizeConnectionSpy.mockRestore();
+			wsCloseSpy.mockRestore();
+			setupConnectionSpy.mockRestore();
+			ws.close();
+		});
+
 		it(`should close for not authorizing connection`, async () => {
 			const { setupConnectionSpy, wsCloseSpy } = setup();
 			const { buffer } = getMessage();
@@ -200,7 +219,7 @@ describe('WebSocketController (WsAdapter)', () => {
 			const error = new Error('unknown error');
 			httpGetCallSpy.mockReturnValueOnce(throwError(() => error));
 
-			ws = await TestConnection.setupWs(wsUrl, 'TEST', { jwt: 'jwt-mocked' });
+			ws = await TestConnection.setupWs(wsUrl, 'TEST', { headers: { cookie: { jwt: 'jwt-mocked' } } });
 			ws.send(buffer);
 
 			expect(wsCloseSpy).toHaveBeenCalledWith(
@@ -222,7 +241,7 @@ describe('WebSocketController (WsAdapter)', () => {
 				.spyOn(wsService, 'authorizeConnection')
 				.mockImplementationOnce(() => Promise.resolve());
 
-			ws = await TestConnection.setupWs(wsUrl, 'TEST', { jwt: 'jwt-mocked' });
+			ws = await TestConnection.setupWs(wsUrl, 'TEST', { headers: { cookie: { jwt: 'jwt-mocked' } } });
 			ws.send(buffer);
 
 			expect(setupConnectionSpy).toHaveBeenCalledWith(expect.anything(), 'TEST');
@@ -244,7 +263,7 @@ describe('WebSocketController (WsAdapter)', () => {
 				throw new Error('unknown error');
 			});
 
-			ws = await TestConnection.setupWs(wsUrl, 'TEST', { jwt: 'jwt-mocked' });
+			ws = await TestConnection.setupWs(wsUrl, 'TEST', { headers: { cookie: { jwt: 'jwt-mocked' } } });
 			ws.send(buffer);
 
 			expect(setupConnectionSpy).toHaveBeenCalledWith(expect.anything(), 'TEST');

apps/server/src/modules/tldraw/controller/tldraw.ws.ts

@@ -3,7 +3,7 @@ import { Server, WebSocket } from 'ws';
 import { Request } from 'express';
 import { ConfigService } from '@nestjs/config';
 import cookie from 'cookie';
-import { BadRequestException } from '@nestjs/common';
+import { BadRequestException, NotFoundException } from '@nestjs/common';
 import { Logger } from '@src/core/logger';
 import { WebsocketCloseErrorLoggable } from '../loggable/websocket-close-error.loggable';
 import { TldrawConfig, SOCKET_PORT } from '../config';
@@ -28,12 +28,21 @@ export class TldrawWs implements OnGatewayInit, OnGatewayConnection {
 			try {
 				await this.tldrawWsService.authorizeConnection(docName, cookies?.jwt);
 			} catch (err) {
-				this.closeClientAndLogError(
-					client,
-					WsCloseCodeEnum.WS_CLIENT_UNAUTHORISED_CONNECTION_CODE,
-					WsCloseMessageEnum.WS_CLIENT_UNAUTHORISED_CONNECTION_MESSAGE,
-					err as Error
-				);
+				if (err instanceof NotFoundException) {
+					this.closeClientAndLogError(
+						client,
+						WsCloseCodeEnum.WS_CLIENT_NOT_FOUND_CODE,
+						WsCloseMessageEnum.WS_CLIENT_NOT_FOUND_MESSAGE,
+						err as Error
+					);
+				} else {
+					this.closeClientAndLogError(
+						client,
+						WsCloseCodeEnum.WS_CLIENT_UNAUTHORISED_CONNECTION_CODE,
+						WsCloseMessageEnum.WS_CLIENT_UNAUTHORISED_CONNECTION_MESSAGE,
+						err as Error
+					);
+				}
 				return;
 			}
 			try {

apps/server/src/modules/tldraw/types/ws-close-enum.ts

@@ -1,10 +1,12 @@
 export enum WsCloseCodeEnum {
 	WS_CLIENT_BAD_REQUEST_CODE = 4400,
 	WS_CLIENT_UNAUTHORISED_CONNECTION_CODE = 4401,
+	WS_CLIENT_NOT_FOUND_CODE = 4404,
 	WS_CLIENT_ESTABLISHING_CONNECTION_CODE = 4500,
 }
 export enum WsCloseMessageEnum {
 	WS_CLIENT_BAD_REQUEST_MESSAGE = 'Document name is mandatory in url or Tldraw Tool is turned off.',
 	WS_CLIENT_UNAUTHORISED_CONNECTION_MESSAGE = "Unauthorised connection - you don't have permission to this drawing.",
+	WS_CLIENT_NOT_FOUND_MESSAGE = 'Drawing not found.',
 	WS_CLIENT_ESTABLISHING_CONNECTION_MESSAGE = 'Unable to establish websocket connection. Try again later.',
 }