fix: throttle ip tracker

Signed-off-by: Innei <i@innei.in>

apps/core/src/common/filters/any-exception.filter.ts

@@ -91,7 +91,7 @@ export class AllExceptionsFilter implements ExceptionFilter {
         body: `IP: ${ip}`,
       })
 
-      return response.status(444).send({
+      return response.status(429).send({
         message: '请求过于频繁，请稍后再试',
       })
     }

apps/core/src/common/guards/throttler.guard.ts

@@ -1,18 +1,46 @@
 import type { ExecutionContext } from '@nestjs/common'
+import type { ThrottlerOptions } from '@nestjs/throttler'
+import type { FastifyBizRequest } from '~/transformers/get-req.transformer'
 
 import { Injectable } from '@nestjs/common'
-import { ThrottlerGuard } from '@nestjs/throttler'
+import { ThrottlerException, ThrottlerGuard } from '@nestjs/throttler'
 
 import { getNestExecutionContextRequest } from '~/transformers/get-req.transformer'
+import { getIp } from '~/utils'
 
 @Injectable()
 export class ExtendThrottlerGuard extends ThrottlerGuard {
-  async canActivate(context: ExecutionContext): Promise<boolean> {
+  protected async shouldSkip(context: ExecutionContext): Promise<boolean> {
     const req = getNestExecutionContextRequest(context)
 
     if (req.user) {
       return true
     }
-    return super.canActivate(context)
+    return super.shouldSkip(context)
+  }
+
+  protected async getTracker(req: FastifyBizRequest) {
+    return getIp(req.raw)
+  }
+}
+
+export class WsExtendThrottlerGuard extends ExtendThrottlerGuard {
+  async handleRequest(
+    context: ExecutionContext,
+    limit: number,
+    ttl: number,
+    throttler: ThrottlerOptions,
+  ): Promise<boolean> {
+    const client = context.switchToWs().getClient()
+    const ip = client._socket.remoteAddress
+    const key = this.generateKey(context, ip, throttler.name || 'ws-default')
+    const { totalHits } = await this.storageService.increment(key, ttl)
+
+    console.log('totalHits', totalHits)
+    if (totalHits > limit) {
+      throw new ThrottlerException()
+    }
+
+    return true
   }
 }

apps/core/src/modules/activity/activity.service.ts

@@ -38,7 +38,7 @@ import {
   parseRoomName,
 } from './activity.util'
 
-declare module '~/utils/socket.util' {
+declare module '~/types/socket-meta' {
   interface SocketMetadata {
     presence?: ActivityPresence
   }

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

@@ -7,9 +7,10 @@ import type {
 } from '@nestjs/websockets'
 import type SocketIO from 'socket.io'
 
-import { forwardRef, Inject } from '@nestjs/common'
+import { forwardRef, Inject, UseGuards } from '@nestjs/common'
 import { SubscribeMessage, WebSocketGateway } from '@nestjs/websockets'
 
+import { WsExtendThrottlerGuard } from '~/common/guards/throttler.guard'
 import { LOG_DIR } from '~/constants/path.constant'
 import { JWTService } from '~/processors/helper/helper.jwt.service'
 import { CacheService } from '~/processors/redis/cache.service'
@@ -22,6 +23,7 @@ import { createAuthGateway } from '../shared/auth.gateway'
 
 const AuthGateway = createAuthGateway({ namespace: 'admin', authway: 'jwt' })
 @WebSocketGateway<GatewayMetadata>({ namespace: 'admin' })
+@UseGuards(WsExtendThrottlerGuard)
 export class AdminEventsGateway
   extends AuthGateway
   implements OnGatewayConnection, OnGatewayDisconnect

apps/core/src/processors/gateway/gateway.service.ts

@@ -1,4 +1,4 @@
-import type { SocketMetadata } from '~/utils/socket.util'
+import type { SocketMetadata } from '~/types/socket-meta'
 import type { RemoteSocket, Socket } from 'socket.io'
 import type {
   DecorateAcknowledgementsWithMultipleResponses,

apps/core/src/processors/gateway/web/events.gateway.ts

@@ -14,6 +14,7 @@ import type {
 import type { SocketType } from '../gateway.service'
 import type { EventGatewayHooks } from './hook.interface'
 
+import { UseGuards } from '@nestjs/common'
 import {
   ConnectedSocket,
   MessageBody,
@@ -22,6 +23,7 @@ import {
   WebSocketServer,
 } from '@nestjs/websockets'
 
+import { WsExtendThrottlerGuard } from '~/common/guards/throttler.guard'
 import { BusinessEvents } from '~/constants/business-event.constant'
 import { RedisKeys } from '~/constants/cache.constant'
 import { CacheService } from '~/processors/redis/cache.service'
@@ -33,7 +35,7 @@ import { BroadcastBaseGateway } from '../base.gateway'
 import { GatewayService } from '../gateway.service'
 import { MessageEventDto, SupportedMessageEvent } from './dtos/message'
 
-declare module '~/utils/socket.util' {
+declare module '~/types/socket-meta' {
   interface SocketMetadata {
     sessionId: string
 
@@ -42,6 +44,8 @@ declare module '~/utils/socket.util' {
 }
 
 const namespace = 'web'
+
+@UseGuards(WsExtendThrottlerGuard)
 @WebSocketGateway<GatewayMetadata>({
   namespace,
 })

apps/core/src/utils/ip.util.ts

@@ -23,6 +23,7 @@ export const getIp = (request: FastifyRequest | IncomingMessage) => {
     headers['X-Real-IP'] ||
     headers['x-real-ip'] ||
     req?.ip ||
+    req?.ips?.[0] ||
     req?.raw?.connection?.remoteAddress ||
     req?.raw?.socket?.remoteAddress ||
     undefined