refactor: websocket repository (#23228)

server/src/repositories/event.repository.ts

@@ -1,27 +1,15 @@
 import { Injectable } from '@nestjs/common';
 import { ModuleRef, Reflector } from '@nestjs/core';
-import {
-  OnGatewayConnection,
-  OnGatewayDisconnect,
-  OnGatewayInit,
-  WebSocketGateway,
-  WebSocketServer,
-} from '@nestjs/websockets';
 import { ClassConstructor } from 'class-transformer';
 import _ from 'lodash';
-import { Server, Socket } from 'socket.io';
+import { Socket } from 'socket.io';
 import { SystemConfig } from 'src/config';
 import { EventConfig } from 'src/decorators';
-import { AssetResponseDto } from 'src/dtos/asset-response.dto';
 import { AuthDto } from 'src/dtos/auth.dto';
-import { NotificationDto } from 'src/dtos/notification.dto';
-import { ReleaseNotification, ServerVersionResponseDto } from 'src/dtos/server.dto';
-import { SyncAssetExifV1, SyncAssetV1 } from 'src/dtos/sync.dto';
 import { ImmichWorker, JobStatus, MetadataKey, QueueName, UserAvatarColor, UserStatus } from 'src/enum';
 import { ConfigRepository } from 'src/repositories/config.repository';
 import { LoggingRepository } from 'src/repositories/logging.repository';
 import { JobItem, JobSource } from 'src/types';
-import { handlePromiseError } from 'src/utils/misc';
 
 type EmitHandlers = Partial<{ [T in EmitEvent]: Array<EventItem<T>> }>;
 
@@ -130,33 +118,11 @@ type UserEvent = {
   profileChangedAt: Date;
 };
 
-export const serverEvents = ['ConfigUpdate'] as const;
-export type ServerEvents = (typeof serverEvents)[number];
-
 export type EmitEvent = keyof EventMap;
 export type EmitHandler<T extends EmitEvent> = (...args: ArgsOf<T>) => Promise<void> | void;
 export type ArgOf<T extends EmitEvent> = EventMap[T][0];
 export type ArgsOf<T extends EmitEvent> = EventMap[T];
 
-export interface ClientEventMap {
-  on_upload_success: [AssetResponseDto];
-  on_user_delete: [string];
-  on_asset_delete: [string];
-  on_asset_trash: [string[]];
-  on_asset_update: [AssetResponseDto];
-  on_asset_hidden: [string];
-  on_asset_restore: [string[]];
-  on_asset_stack_update: string[];
-  on_person_thumbnail: [string];
-  on_server_version: [ServerVersionResponseDto];
-  on_config_update: [];
-  on_new_release: [ReleaseNotification];
-  on_notification: [NotificationDto];
-  on_session_delete: [string];
-
-  AssetUploadReadyV1: [{ asset: SyncAssetV1; exif: SyncAssetExifV1 }];
-}
-
 export type EventItem<T extends EmitEvent> = {
   event: T;
   handler: EmitHandler<T>;
@@ -165,18 +131,9 @@ export type EventItem<T extends EmitEvent> = {
 
 export type AuthFn = (client: Socket) => Promise<AuthDto>;
 
-@WebSocketGateway({
-  cors: true,
-  path: '/api/socket.io',
-  transports: ['websocket'],
-})
 @Injectable()
-export class EventRepository implements OnGatewayConnection, OnGatewayDisconnect, OnGatewayInit {
+export class EventRepository {
   private emitHandlers: EmitHandlers = {};
-  private authFn?: AuthFn;
-
-  @WebSocketServer()
-  private server?: Server;
 
   constructor(
     private moduleRef: ModuleRef,
@@ -237,38 +194,6 @@ export class EventRepository implements OnGatewayConnection, OnGatewayDisconnect
     }
   }
 
-  afterInit(server: Server) {
-    this.logger.log('Initialized websocket server');
-
-    for (const event of serverEvents) {
-      server.on(event, (...args: ArgsOf<any>) => {
-        this.logger.debug(`Server event: ${event} (receive)`);
-        handlePromiseError(this.onEvent({ name: event, args, server: true }), this.logger);
-      });
-    }
-  }
-
-  async handleConnection(client: Socket) {
-    try {
-      this.logger.log(`Websocket Connect:    ${client.id}`);
-      const auth = await this.authenticate(client);
-      await client.join(auth.user.id);
-      if (auth.session) {
-        await client.join(auth.session.id);
-      }
-      await this.onEvent({ name: 'WebsocketConnect', args: [{ userId: auth.user.id }], server: false });
-    } catch (error: Error | any) {
-      this.logger.error(`Websocket connection error: ${error}`, error?.stack);
-      client.emit('error', 'unauthorized');
-      client.disconnect();
-    }
-  }
-
-  async handleDisconnect(client: Socket) {
-    this.logger.log(`Websocket Disconnect: ${client.id}`);
-    await client.leave(client.nsp.name);
-  }
-
   private addHandler<T extends EmitEvent>(item: Item<T>): void {
     const event = item.event;
 
@@ -283,7 +208,7 @@ export class EventRepository implements OnGatewayConnection, OnGatewayDisconnect
     return this.onEvent({ name: event, args, server: false });
   }
 
-  private async onEvent<T extends EmitEvent>(event: { name: T; args: ArgsOf<T>; server: boolean }): Promise<void> {
+  async onEvent<T extends EmitEvent>(event: { name: T; args: ArgsOf<T>; server: boolean }): Promise<void> {
     const handlers = this.emitHandlers[event.name] || [];
     for (const { handler, server } of handlers) {
       // exclude handlers that ignore server events
@@ -294,29 +219,4 @@ export class EventRepository implements OnGatewayConnection, OnGatewayDisconnect
       await handler(...event.args);
     }
   }
-
-  clientSend<T extends keyof ClientEventMap>(event: T, room: string, ...data: ClientEventMap[T]) {
-    this.server?.to(room).emit(event, ...data);
-  }
-
-  clientBroadcast<T extends keyof ClientEventMap>(event: T, ...data: ClientEventMap[T]) {
-    this.server?.emit(event, ...data);
-  }
-
-  serverSend<T extends ServerEvents>(event: T, ...args: ArgsOf<T>): void {
-    this.logger.debug(`Server event: ${event} (send)`);
-    this.server?.serverSideEmit(event, ...args);
-  }
-
-  setAuthFn(fn: (client: Socket) => Promise<AuthDto>) {
-    this.authFn = fn;
-  }
-
-  private async authenticate(client: Socket) {
-    if (!this.authFn) {
-      throw new Error('Auth function not set');
-    }
-
-    return this.authFn(client);
-  }
 }

server/src/repositories/index.ts

@@ -44,6 +44,7 @@ import { TrashRepository } from 'src/repositories/trash.repository';
 import { UserRepository } from 'src/repositories/user.repository';
 import { VersionHistoryRepository } from 'src/repositories/version-history.repository';
 import { ViewRepository } from 'src/repositories/view-repository';
+import { WebsocketRepository } from 'src/repositories/websocket.repository';
 
 export const repositories = [
   AccessRepository,
@@ -92,4 +93,5 @@ export const repositories = [
   UserRepository,
   ViewRepository,
   VersionHistoryRepository,
+  WebsocketRepository,
 ];

server/src/repositories/websocket.repository.ts

@@ -0,0 +1,118 @@
+import { Injectable } from '@nestjs/common';
+import {
+  OnGatewayConnection,
+  OnGatewayDisconnect,
+  OnGatewayInit,
+  WebSocketGateway,
+  WebSocketServer,
+} from '@nestjs/websockets';
+import { Server, Socket } from 'socket.io';
+import { AssetResponseDto } from 'src/dtos/asset-response.dto';
+import { AuthDto } from 'src/dtos/auth.dto';
+import { NotificationDto } from 'src/dtos/notification.dto';
+import { ReleaseNotification, ServerVersionResponseDto } from 'src/dtos/server.dto';
+import { SyncAssetExifV1, SyncAssetV1 } from 'src/dtos/sync.dto';
+import { ArgsOf, EventRepository } from 'src/repositories/event.repository';
+import { LoggingRepository } from 'src/repositories/logging.repository';
+import { handlePromiseError } from 'src/utils/misc';
+
+export const serverEvents = ['ConfigUpdate'] as const;
+export type ServerEvents = (typeof serverEvents)[number];
+
+export interface ClientEventMap {
+  on_upload_success: [AssetResponseDto];
+  on_user_delete: [string];
+  on_asset_delete: [string];
+  on_asset_trash: [string[]];
+  on_asset_update: [AssetResponseDto];
+  on_asset_hidden: [string];
+  on_asset_restore: [string[]];
+  on_asset_stack_update: string[];
+  on_person_thumbnail: [string];
+  on_server_version: [ServerVersionResponseDto];
+  on_config_update: [];
+  on_new_release: [ReleaseNotification];
+  on_notification: [NotificationDto];
+  on_session_delete: [string];
+
+  AssetUploadReadyV1: [{ asset: SyncAssetV1; exif: SyncAssetExifV1 }];
+}
+
+export type AuthFn = (client: Socket) => Promise<AuthDto>;
+
+@WebSocketGateway({
+  cors: true,
+  path: '/api/socket.io',
+  transports: ['websocket'],
+})
+@Injectable()
+export class WebsocketRepository implements OnGatewayConnection, OnGatewayDisconnect, OnGatewayInit {
+  private authFn?: AuthFn;
+
+  @WebSocketServer()
+  private server?: Server;
+
+  constructor(
+    private eventRepository: EventRepository,
+    private logger: LoggingRepository,
+  ) {
+    this.logger.setContext(WebsocketRepository.name);
+  }
+
+  afterInit(server: Server) {
+    this.logger.log('Initialized websocket server');
+
+    for (const event of serverEvents) {
+      server.on(event, (...args: ArgsOf<any>) => {
+        this.logger.debug(`Server event: ${event} (receive)`);
+        handlePromiseError(this.eventRepository.onEvent({ name: event, args, server: true }), this.logger);
+      });
+    }
+  }
+
+  async handleConnection(client: Socket) {
+    try {
+      this.logger.log(`Websocket Connect:    ${client.id}`);
+      const auth = await this.authenticate(client);
+      await client.join(auth.user.id);
+      if (auth.session) {
+        await client.join(auth.session.id);
+      }
+      await this.eventRepository.emit('WebsocketConnect', { userId: auth.user.id });
+    } catch (error: Error | any) {
+      this.logger.error(`Websocket connection error: ${error}`, error?.stack);
+      client.emit('error', 'unauthorized');
+      client.disconnect();
+    }
+  }
+
+  async handleDisconnect(client: Socket) {
+    this.logger.log(`Websocket Disconnect: ${client.id}`);
+    await client.leave(client.nsp.name);
+  }
+
+  clientSend<T extends keyof ClientEventMap>(event: T, room: string, ...data: ClientEventMap[T]) {
+    this.server?.to(room).emit(event, ...data);
+  }
+
+  clientBroadcast<T extends keyof ClientEventMap>(event: T, ...data: ClientEventMap[T]) {
+    this.server?.emit(event, ...data);
+  }
+
+  serverSend<T extends ServerEvents>(event: T, ...args: ArgsOf<T>): void {
+    this.logger.debug(`Server event: ${event} (send)`);
+    this.server?.serverSideEmit(event, ...args);
+  }
+
+  setAuthFn(fn: (client: Socket) => Promise<AuthDto>) {
+    this.authFn = fn;
+  }
+
+  private async authenticate(client: Socket) {
+    if (!this.authFn) {
+      throw new Error('Auth function not set');
+    }
+
+    return this.authFn(client);
+  }
+}