whale-town-end/src/business/location_broadcast/websocket_auth.guard.ts

/**
 * WebSocket认证守卫
 * 
 * 功能描述：
 * - 验证WebSocket连接中的JWT令牌
 * - 提取用户信息并添加到WebSocket客户端上下文
 * - 保护需要认证的WebSocket事件处理器
 * - 处理WebSocket特有的认证流程
 * 
 * 职责分离：
 * - 专注于WebSocket环境下的JWT令牌验证
 * - 提供统一的WebSocket认证守卫机制
 * - 处理WebSocket认证失败的异常情况
 * - 支持实时通信的安全认证
 * 
 * 技术实现：
 * - 从WebSocket消息中提取JWT令牌
 * - 使用现有的LoginCore服务进行令牌验证
 * - 将用户信息附加到WebSocket客户端对象
 * - 提供错误处理和日志记录
 * 
 * 最近修改：
 * - 2026-01-09: 重构为原生WebSocket - 适配原生WebSocket接口 (修改者: moyin)
 * 
 * @author moyin
 * @version 2.0.0
 * @since 2026-01-08
 * @lastModified 2026-01-09
 */

import { Injectable, CanActivate, ExecutionContext, Logger } from '@nestjs/common';
import { WsException } from '@nestjs/websockets';
import { LoginCoreService, JwtPayload } from '../../core/login_core/login_core.service';

/**
 * 扩展的WebSocket客户端接口，包含用户信息
 * 
 * 职责：
 * - 扩展原生WebSocket接口
 * - 添加用户认证信息到客户端对象
 * - 提供类型安全的用户数据访问
 */
export interface AuthenticatedSocket extends WebSocket {
  /** 客户端ID */
  id: string;
  /** 认证用户信息 */
  user?: JwtPayload;
  /** 用户ID（便于快速访问） */
  userId?: string;
  /** 认证时间戳 */
  authenticatedAt?: number;
  /** 会话ID集合 */
  sessionIds?: Set<string>;
  /** 连接超时 */
  connectionTimeout?: NodeJS.Timeout;
  /** 心跳状态 */
  isAlive?: boolean;
}

@Injectable()
export class WebSocketAuthGuard implements CanActivate {
  private readonly logger = new Logger(WebSocketAuthGuard.name);

  constructor(private readonly loginCoreService: LoginCoreService) {}

  /**
   * WebSocket JWT令牌验证和用户认证
   * 
   * 技术实现：
   * 1. 从WebSocket客户端获取认证信息
   * 2. 提取JWT令牌（支持多种提取方式）
   * 3. 验证令牌的有效性和签名
   * 4. 解码令牌获取用户信息
   * 5. 将用户信息添加到Socket客户端对象
   * 6. 记录认证成功或失败的日志
   * 7. 返回认证结果或抛出WebSocket异常
   * 
   * @param context 执行上下文，包含WebSocket客户端信息
   * @returns Promise<boolean> 认证是否成功
   * @throws WsException 当令牌缺失或无效时
   */
  async canActivate(context: ExecutionContext): Promise<boolean> {
    const client = context.switchToWs().getClient<AuthenticatedSocket>();
    const data = context.switchToWs().getData();

    this.logAuthStart(client, context);

    try {
      const token = this.extractToken(client, data);
      
      if (!token) {
        this.handleMissingToken(client);
      }

      // 如果是缓存的认证信息，直接返回成功
      if (token === 'cached' && client.user && client.userId) {
        this.logger.debug('使用缓存的认证信息', {
          socketId: client.id,
          userId: client.userId,
        });
        return true;
      }

      const payload = await this.loginCoreService.verifyToken(token, 'access');
      this.attachUserToClient(client, payload);
      this.logAuthSuccess(client, payload);

      return true;

    } catch (error) {
      this.handleAuthError(client, error);
    }
  }

  /**
   * 记录认证开始日志
   * 
   * @param client WebSocket客户端
   * @param context 执行上下文
   * @private
   */
  private logAuthStart(client: AuthenticatedSocket, context: ExecutionContext): void {
    this.logger.log('开始WebSocket认证验证', {
      operation: 'websocket_auth',
      socketId: client.id,
      eventName: context.getHandler().name,
      timestamp: new Date().toISOString()
    });
  }

  /**
   * 处理缺少令牌的情况
   * 
   * @param client WebSocket客户端
   * @throws WsException
   * @private
   */
  private handleMissingToken(client: AuthenticatedSocket): never {
    this.logger.warn('WebSocket认证失败：缺少认证令牌', {
      operation: 'websocket_auth',
      socketId: client.id,
      reason: 'missing_token'
    });
    
    throw new WsException({
      type: 'error',
      code: 'INVALID_TOKEN',
      message: '缺少认证令牌',
      timestamp: Date.now()
    });
  }

  /**
   * 将用户信息附加到客户端
   * 
   * @param client WebSocket客户端
   * @param payload JWT载荷
   * @private
   */
  private attachUserToClient(client: AuthenticatedSocket, payload: JwtPayload): void {
    client.user = payload;
    client.userId = payload.sub;
    client.authenticatedAt = Date.now();
  }

  /**
   * 记录认证成功日志
   * 
   * @param client WebSocket客户端
   * @param payload JWT载荷
   * @private
   */
  private logAuthSuccess(client: AuthenticatedSocket, payload: JwtPayload): void {
    this.logger.log('WebSocket认证成功', {
      operation: 'websocket_auth',
      socketId: client.id,
      userId: payload.sub,
      username: payload.username,
      role: payload.role,
      timestamp: new Date().toISOString()
    });
  }

  /**
   * 处理认证错误
   * 
   * @param client WebSocket客户端
   * @param error 错误对象
   * @throws WsException
   * @private
   */
  private handleAuthError(client: AuthenticatedSocket, error: any): never {
    this.logger.error('WebSocket认证失败', {
      operation: 'websocket_auth',
      socketId: client.id,
      error: error instanceof Error ? error.message : String(error),
      timestamp: new Date().toISOString()
    }, error instanceof Error ? error.stack : undefined);

    // 如果已经是WsException，直接抛出
    if (error instanceof WsException) {
      throw error;
    }

    // 转换为WebSocket异常
    throw new WsException({
      type: 'error',
      code: 'INVALID_TOKEN',
      message: '无效的认证令牌',
      details: {
        reason: error instanceof Error ? error.message : String(error)
      },
      timestamp: Date.now()
    });
  }

  /**
   * 从WebSocket连接中提取JWT令牌
   * 
   * 技术实现：
   * 1. 优先从消息数据中提取token字段
   * 2. 检查是否已经认证过（用于后续消息）
   * 3. 从URL查询参数中提取token（如果可用）
   * 
   * 支持的令牌传递方式：
   * - 消息数据: { token: "jwt_token" }
   * - 缓存认证: 使用已验证的用户信息
   * 
   * @param client WebSocket客户端对象
   * @param data 消息数据
   * @returns JWT令牌字符串或undefined
   */
  private extractToken(client: AuthenticatedSocket, data: any): string | undefined {
    // 1. 优先从消息数据中提取token
    if (data && typeof data === 'object' && data.token) {
      this.logger.debug('从消息数据中提取到token', {
        socketId: client.id,
        source: 'message_data'
      });
      return data.token;
    }

    // 2. 检查是否已经认证过（用于后续消息）
    if (client.user && client.userId) {
      this.logger.debug('使用已认证的用户信息', {
        socketId: client.id,
        userId: client.userId,
        source: 'cached_auth'
      });
      return 'cached'; // 返回特殊标识，表示使用缓存的认证信息
    }

    this.logger.warn('未找到有效的认证令牌', {
      socketId: client.id,
      availableSources: {
        messageData: !!data?.token,
        cachedAuth: !!(client.user && client.userId)
      }
    });

    return undefined;
  }

  /**
   * 清理客户端的认证信息
   * 
   * @param client WebSocket客户端
   */
  static clearAuthentication(client: AuthenticatedSocket): void {
    delete client.user;
    delete client.userId;
    delete client.authenticatedAt;
  }
}