dmail-api/src/modules/email-gateway/services/websocket-auth.service.ts

import { Injectable, Logger } from "@nestjs/common";
import { JsonWebTokenError, JwtService, NotBeforeError, TokenExpiredError } from "@nestjs/jwt";
import { Socket } from "socket.io";

import { WebSocketMessage } from "../../../common/enums/message.enum";
import { ITokenPayload } from "../../auth/interfaces/IToken-payload";
import { WEBSOCKET_EVENTS, WebSocketEvent } from "../../email/constants/email-events.constant";
import { AuthenticationResult, WebSocketErrorContext } from "../../email/interfaces/websocket.interface";
import { extractTokenFromClient } from "../../utils/services/extract-token.utils";

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

  constructor(private readonly jwtService: JwtService) {}

  /**
   * Authenticates a WebSocket client and returns the result
   *
   * @param client - Socket.IO client to authenticate
   * @returns Promise with authentication result
   */
  async authenticateClient(client: Socket): Promise<AuthenticationResult> {
    try {
      const token = extractTokenFromClient(client);

      if (!token) {
        this.logAuthenticationAttempt(client, false, "No token provided");
        return {
          success: false,
          error: WebSocketMessage.AUTHENTICATION_REQUIRED,
        };
      }

      const payload = await this.jwtService.verifyAsync<ITokenPayload>(token);

      if (!payload?.id) {
        this.logAuthenticationAttempt(client, false, "Invalid token payload");
        return {
          success: false,
          error: WebSocketMessage.INVALID_TOKEN,
        };
      }

      // Store user data in client
      client.data.user = payload;

      this.logAuthenticationAttempt(client, true, undefined, payload.id);

      return {
        success: true,
        user: payload,
      };
    } catch (error) {
      const errorMessage = this.getAuthErrorMessage(error);
      this.logAuthenticationAttempt(client, false, errorMessage);

      return {
        success: false,
        error: errorMessage,
      };
    }
  }

  /**
   * Emits authentication failure event to client and disconnects
   *
   * @param client - Socket.IO client
   * @param error - Error message to send
   */
  handleAuthenticationFailure(client: Socket, error: string): void {
    client.emit(WEBSOCKET_EVENTS.CONNECTION_ERROR, {
      message: error,
      timestamp: new Date().toISOString(),
    });

    // Graceful disconnect with delay to ensure message is received
    setTimeout(() => {
      client.disconnect(true);
    }, 100);
  }

  /**
   * Emits successful authentication event to client
   *
   * @param client - Socket.IO client
   * @param user - Authenticated user payload
   */
  emitAuthenticationSuccess(client: Socket, user: ITokenPayload): void {
    client.emit(WEBSOCKET_EVENTS.CONNECTION_ESTABLISHED, {
      message: WebSocketMessage.AUTHENTICATED,
      userId: user.id,
      timestamp: new Date().toISOString(),
    });
  }

  /**
   * Creates error context object for logging and monitoring
   *
   * @param client - Socket.IO client
   * @param event - Event name (optional)
   * @param data - Event data (optional)
   * @returns WebSocket error context
   */
  createErrorContext(client: Socket, event?: WebSocketEvent, data?: unknown): WebSocketErrorContext {
    return {
      clientId: client.id,
      userId: client.data?.user?.id,
      sessionId: client.data?.sessionId,
      event,
      data,
      timestamp: new Date(),
    };
  }

  /**
   * Validates if client is authenticated
   *
   * @param client - Socket.IO client to validate
   * @returns True if client has valid user data
   */
  isClientAuthenticated(client: Socket): boolean {
    return !!client.data?.user?.id;
  }

  /**
   * Gets user ID from authenticated client
   *
   * @param client - Socket.IO client
   * @returns User ID or undefined if not authenticated
   */
  getUserId(client: Socket): string | undefined {
    return client.data?.user?.id;
  }

  /**
   * Logs authentication attempts with structured data
   */
  private logAuthenticationAttempt(client: Socket, success: boolean, error?: string, userId?: string): void {
    const logData = {
      clientId: client.id,
      clientIP: client.handshake.address,
      success,
      userId,
      error,
      timestamp: new Date().toISOString(),
    };

    if (success) {
      this.logger.log(`WebSocket authentication successful`, logData);
    } else {
      this.logger.warn(`WebSocket authentication failed`, logData);
    }
  }

  /**
   * Maps JWT errors to user-friendly messages
   */
  private getAuthErrorMessage(error: unknown): string {
    if (error instanceof TokenExpiredError) {
      return WebSocketMessage.TOKEN_EXPIRED;
    }

    if (error instanceof JsonWebTokenError) {
      return WebSocketMessage.INVALID_TOKEN;
    }

    if (error instanceof NotBeforeError) {
      return WebSocketMessage.INVALID_TOKEN;
    }

    return WebSocketMessage.INVALID_TOKEN;
  }
}