refactor: auth gateway

src/modules/auth/auth.controller.ts

@@ -12,6 +12,7 @@ import {
   Controller,
   Delete,
   Get,
+  NotFoundException,
   Post,
   Query,
   Scope,
@@ -88,9 +89,23 @@ export class AuthController {
   @Auth()
   async deleteToken(@Query() query: MongoIdDto) {
     const { id } = query
+    const token = await this.authService
+      .getAllAccessToken()
+      .then((models) =>
+        models.find((model) => {
+          return (model as any).id === id
+        }),
+      )
+      .then((model) => {
+        return model?.token
+      })
+
+    if (!token) {
+      throw new NotFoundException(`token ${id} is not found`)
+    }
     await this.authService.deleteToken(id)
 
-    this.eventEmitter.emit(EventBusEvents.TokenExpired, id)
+    this.eventEmitter.emit(EventBusEvents.TokenExpired, token)
     return 'OK'
   }
 }

src/modules/auth/auth.service.ts

@@ -59,7 +59,7 @@ export class AuthService {
     return tokens.map((token) => ({
       // @ts-ignore
       id: token._id,
-      ...omit(token, ['_id', '__v', 'token']),
+      ...omit(token, ['_id', '__v']),
     })) as any as TokenModel[]
   }
 

src/modules/pty/pty.gateway.ts

@@ -18,9 +18,10 @@ import { DATA_DIR } from '~/constants/path.constant'
 import { AuthService } from '~/modules/auth/auth.service'
 import { ConfigsService } from '~/modules/configs/configs.service'
 import { CacheService } from '~/processors/cache/cache.service'
-import { AuthGateway } from '~/processors/gateway/admin/auth.gateway'
+import { createAuthGateway } from '~/processors/gateway/shared/auth.gateway'
 import { getIp, getRedisKey } from '~/utils'
 
+const AuthGateway = createAuthGateway({ namespace: 'pty', authway: 'jwt' })
 @WebSocketGateway<GatewayMetadata>({ namespace: 'pty' })
 export class PTYGateway
   extends AuthGateway

src/processors/gateway/admin/auth.gateway.ts

@@ -1,97 +0,0 @@
-import { Namespace, Socket } from 'socket.io'
-
-import { OnEvent } from '@nestjs/event-emitter'
-import { JwtService } from '@nestjs/jwt'
-import {
-  OnGatewayConnection,
-  OnGatewayDisconnect,
-  WebSocketServer,
-} from '@nestjs/websockets'
-
-import { EventBusEvents } from '~/constants/event-bus.constant'
-import { AuthService } from '~/modules/auth/auth.service'
-
-import { BusinessEvents } from '../../../constants/business-event.constant'
-import { BoardcastBaseGateway } from '../base.gateway'
-
-export abstract class AuthGateway
-  extends BoardcastBaseGateway
-  implements OnGatewayConnection, OnGatewayDisconnect
-{
-  constructor(
-    protected readonly jwtService: JwtService,
-    protected readonly authService: AuthService,
-  ) {
-    super()
-  }
-
-  @WebSocketServer()
-  protected namespace: Namespace
-
-  async authFailed(client: Socket) {
-    client.send(
-      this.gatewayMessageFormat(BusinessEvents.AUTH_FAILED, '认证失败'),
-    )
-    client.disconnect()
-  }
-
-  async authToken(token: string): Promise<boolean> {
-    if (typeof token !== 'string') {
-      return false
-    }
-    // first check this token is custom token in user
-    const verifyCustomToken = await this.authService.verifyCustomToken(token)
-    if (verifyCustomToken) {
-      return true
-    } else {
-      // if not, then verify jwt token
-      try {
-        const payload = this.jwtService.verify(token)
-        const user = await this.authService.verifyPayload(payload)
-        if (!user) {
-          return false
-        }
-      } catch {
-        return false
-      }
-      // is not crash, is verify
-      return true
-    }
-  }
-  async handleConnection(client: Socket) {
-    const token =
-      client.handshake.query.token || client.handshake.headers['authorization']
-    if (!token) {
-      return this.authFailed(client)
-    }
-    if (!(await this.authToken(token as string))) {
-      return this.authFailed(client)
-    }
-
-    super.handleConnect(client)
-
-    const sid = client.id
-    this.tokenSocketIdMap.set(token.toString(), sid)
-  }
-
-  handleDisconnect(client: Socket) {
-    super.handleDisconnect(client)
-  }
-  tokenSocketIdMap = new Map<string, string>()
-
-  @OnEvent(EventBusEvents.TokenExpired)
-  handleTokenExpired(token: string) {
-    const server = this.namespace.server
-    const sid = this.tokenSocketIdMap.get(token)
-    if (!sid) {
-      return false
-    }
-    const socket = server.of('/admin').sockets.get(sid)
-    if (socket) {
-      socket.disconnect()
-      super.handleDisconnect(socket)
-      return true
-    }
-    return false
-  }
-}

src/processors/gateway/admin/events.gateway.ts

@@ -17,8 +17,9 @@ import { CacheService } from '~/processors/cache/cache.service'
 
 import { BusinessEvents } from '../../../constants/business-event.constant'
 import { AuthService } from '../../../modules/auth/auth.service'
-import { AuthGateway } from './auth.gateway'
+import { createAuthGateway } from '../shared/auth.gateway'
 
+const AuthGateway = createAuthGateway({ namespace: 'admin', authway: 'jwt' })
 @WebSocketGateway<GatewayMetadata>({ namespace: 'admin' })
 export class AdminEventsGateway
   extends AuthGateway

src/processors/gateway/gateway.module.ts

@@ -11,12 +11,23 @@ import { Global, Module } from '@nestjs/common'
 import { AuthModule } from '../../modules/auth/auth.module'
 import { AdminEventsGateway } from './admin/events.gateway'
 import { SharedGateway } from './shared/events.gateway'
+import { SystemEventsGateway } from './system/events.gateway'
 import { WebEventsGateway } from './web/events.gateway'
 
 @Global()
 @Module({
   imports: [AuthModule],
-  providers: [AdminEventsGateway, WebEventsGateway, SharedGateway],
-  exports: [AdminEventsGateway, WebEventsGateway, SharedGateway],
+  providers: [
+    AdminEventsGateway,
+    WebEventsGateway,
+    SharedGateway,
+    SystemEventsGateway,
+  ],
+  exports: [
+    AdminEventsGateway,
+    WebEventsGateway,
+    SharedGateway,
+    SystemEventsGateway,
+  ],
 })
 export class GatewayModule {}

src/processors/gateway/shared/auth.gateway.ts

@@ -0,0 +1,134 @@
+import { Namespace, Socket } from 'socket.io'
+
+import { OnEvent } from '@nestjs/event-emitter'
+import { JwtService } from '@nestjs/jwt'
+import {
+  OnGatewayConnection,
+  OnGatewayDisconnect,
+  WebSocketServer,
+} from '@nestjs/websockets'
+
+import { EventBusEvents } from '~/constants/event-bus.constant'
+import { AuthService } from '~/modules/auth/auth.service'
+
+import { BusinessEvents } from '../../../constants/business-event.constant'
+import { BoardcastBaseGateway } from '../base.gateway'
+
+export type AuthGatewayOptions = {
+  namespace: string
+  authway?: 'jwt' | 'custom-token' | 'all'
+}
+
+// @ts-ignore
+export interface IAuthGateway
+  extends OnGatewayConnection,
+    OnGatewayDisconnect,
+    BoardcastBaseGateway {}
+
+export const createAuthGateway = (
+  options: AuthGatewayOptions,
+): new (...args: any[]) => IAuthGateway => {
+  const { namespace, authway = 'all' } = options
+  class AuthGateway extends BoardcastBaseGateway implements IAuthGateway {
+    constructor(
+      protected readonly jwtService: JwtService,
+      protected readonly authService: AuthService,
+    ) {
+      super()
+    }
+
+    @WebSocketServer()
+    protected namespace: Namespace
+
+    async authFailed(client: Socket) {
+      client.send(
+        this.gatewayMessageFormat(BusinessEvents.AUTH_FAILED, '认证失败'),
+      )
+      client.disconnect()
+    }
+
+    async authToken(token: string): Promise<boolean> {
+      if (typeof token !== 'string') {
+        return false
+      }
+      const validCustomToken = async () => {
+        const verifyCustomToken = await this.authService.verifyCustomToken(
+          token,
+        )
+        if (verifyCustomToken) {
+          return true
+        }
+        return false
+      }
+
+      const validJwt = async () => {
+        try {
+          const payload = this.jwtService.verify(token)
+          const user = await this.authService.verifyPayload(payload)
+          if (!user) {
+            return false
+          }
+        } catch {
+          return false
+        }
+        // is not crash, is verify
+        return true
+      }
+
+      switch (authway) {
+        case 'custom-token': {
+          return await validCustomToken()
+        }
+        case 'jwt': {
+          return await validJwt()
+        }
+        case 'all': {
+          const validCustomTokenResult = await validCustomToken()
+          return validCustomTokenResult || (await validJwt())
+        }
+      }
+    }
+
+    async handleConnection(client: Socket) {
+      const token =
+        client.handshake.query.token ||
+        client.handshake.headers['authorization']
+      if (!token) {
+        return this.authFailed(client)
+      }
+      if (!(await this.authToken(token as string))) {
+        return this.authFailed(client)
+      }
+
+      super.handleConnect(client)
+
+      const sid = client.id
+      this.tokenSocketIdMap.set(token.toString(), sid)
+    }
+
+    handleDisconnect(client: Socket) {
+      super.handleDisconnect(client)
+    }
+    tokenSocketIdMap = new Map<string, string>()
+
+    @OnEvent(EventBusEvents.TokenExpired)
+    handleTokenExpired(token: string) {
+      // consola.debug(`token expired: ${token}`)
+
+      const server = this.namespace.server
+      const sid = this.tokenSocketIdMap.get(token)
+      if (!sid) {
+        return false
+      }
+      const socket = server.of(`/${namespace}`).sockets.get(sid)
+      if (socket) {
+        socket.disconnect()
+        super.handleDisconnect(socket)
+        return true
+      }
+      return false
+    }
+  }
+
+  return AuthGateway
+}

src/processors/gateway/system/events.gateway.ts

@@ -0,0 +1,88 @@
+import { resolve } from 'path'
+import SocketIO, { Socket } from 'socket.io'
+
+import { JwtService } from '@nestjs/jwt'
+import {
+  GatewayMetadata,
+  OnGatewayConnection,
+  OnGatewayDisconnect,
+  SubscribeMessage,
+  WebSocketGateway,
+} from '@nestjs/websockets'
+import { Emitter } from '@socket.io/redis-emitter'
+
+import { LOG_DIR } from '~/constants/path.constant'
+import { getTodayLogFilePath } from '~/global/consola.global'
+import { CacheService } from '~/processors/cache/cache.service'
+
+import { BusinessEvents } from '../../../constants/business-event.constant'
+import { AuthService } from '../../../modules/auth/auth.service'
+import { createAuthGateway } from '../shared/auth.gateway'
+
+const AuthGateway = createAuthGateway({
+  namespace: 'admin',
+  authway: 'custom-token',
+})
+
+@WebSocketGateway<GatewayMetadata>({ namespace: 'system' })
+export class SystemEventsGateway
+  extends AuthGateway
+  implements OnGatewayConnection, OnGatewayDisconnect
+{
+  constructor(
+    protected readonly jwtService: JwtService,
+    protected readonly authService: AuthService,
+    private readonly cacheService: CacheService,
+  ) {
+    super(jwtService, authService)
+  }
+
+  subscribeSocketToHandlerMap = new WeakMap<Socket, Function>()
+
+  @SubscribeMessage('log')
+  async subscribeStdOut(client: Socket, data?: { prevLog?: boolean }) {
+    const { prevLog = true } = data || {}
+    if (this.subscribeSocketToHandlerMap.has(client)) {
+      return
+    }
+
+    const handler = (data) => {
+      client.send(this.gatewayMessageFormat(BusinessEvents.STDOUT, data))
+    }
+
+    this.subscribeSocketToHandlerMap.set(client, handler)
+    if (prevLog) {
+      const stream = fs
+        .createReadStream(resolve(LOG_DIR, getTodayLogFilePath()), {
+          encoding: 'utf-8',
+          highWaterMark: 32 * 1024,
+        })
+        .on('data', handler)
+        .on('end', () => {
+          this.cacheService.subscribe('log', handler)
+          stream.close()
+        })
+    } else {
+      this.cacheService.subscribe('log', handler)
+    }
+  }
+
+  @SubscribeMessage('unlog')
+  unsubscribeStdOut(client: Socket) {
+    const cb = this.subscribeSocketToHandlerMap.get(client)
+    if (cb) {
+      this.cacheService.unsubscribe('log', cb as any)
+    }
+    this.subscribeSocketToHandlerMap.delete(client)
+  }
+
+  handleDisconnect(client: SocketIO.Socket) {
+    super.handleDisconnect(client)
+    this.unsubscribeStdOut(client)
+  }
+
+  override broadcast(event: BusinessEvents, data: any) {
+    const client = new Emitter(this.cacheService.getClient())
+    client.of('/admin').emit('message', this.gatewayMessageFormat(event, data))
+  }
+}

src/processors/helper/helper.event.service.ts

@@ -8,6 +8,7 @@ import { EventBusEvents } from '~/constants/event-bus.constant'
 
 import { AdminEventsGateway } from '../gateway/admin/events.gateway'
 import { BoardcastBaseGateway } from '../gateway/base.gateway'
+import { SystemEventsGateway } from '../gateway/system/events.gateway'
 import { WebEventsGateway } from '../gateway/web/events.gateway'
 
 export type EventManagerOptions = {
@@ -30,6 +31,7 @@ export class EventManagerService {
     private readonly webGateway: WebEventsGateway,
 
     private readonly adminGateway: AdminEventsGateway,
+    private readonly systemGateway: SystemEventsGateway,
 
     private readonly emitter2: EventEmitter2,
   ) {
@@ -42,15 +44,33 @@ export class EventManagerService {
 
   private mapScopeToInstance: Record<
     EventScope,
-    (WebEventsGateway | AdminEventsGateway | EventEmitter2)[]
+    (
+      | WebEventsGateway
+      | AdminEventsGateway
+      | EventEmitter2
+      | SystemEventsGateway
+    )[]
   > = {
-    [EventScope.ALL]: [this.webGateway, this.adminGateway, this.emitter2],
+    [EventScope.ALL]: [
+      this.webGateway,
+      this.adminGateway,
+      this.emitter2,
+      this.systemGateway,
+    ],
     [EventScope.TO_VISITOR]: [this.webGateway],
     [EventScope.TO_ADMIN]: [this.adminGateway],
-    [EventScope.TO_SYSTEM]: [this.emitter2],
+    [EventScope.TO_SYSTEM]: [this.emitter2, this.systemGateway],
     [EventScope.TO_VISITOR_ADMIN]: [this.webGateway, this.adminGateway],
-    [EventScope.TO_SYSTEM_VISITOR]: [this.emitter2, this.webGateway],
-    [EventScope.TO_SYSTEM_ADMIN]: [this.emitter2, this.adminGateway],
+    [EventScope.TO_SYSTEM_VISITOR]: [
+      this.emitter2,
+      this.webGateway,
+      this.systemGateway,
+    ],
+    [EventScope.TO_SYSTEM_ADMIN]: [
+      this.emitter2,
+      this.adminGateway,
+      this.systemGateway,
+    ],
   }
 
   #key = 'event-manager'