dmenu-api/src/modules/auth/services/tokens.service.ts

import { createHash, randomBytes } from 'node:crypto';
import { EntityManager } from '@mikro-orm/postgresql';
import { LockMode } from '@mikro-orm/core';
import { Injectable, Logger, UnauthorizedException } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { JwtService } from '@nestjs/jwt';
import dayjs from 'dayjs';

import { AuthMessage } from '../../../common/enums/message.enum';
import { RefreshToken, RefreshTokenType } from '../entities/refresh-token.entity';
import { IAdminTokenPayload, ITokenPayload } from '../interfaces/IToken-payload';
import { Restaurant } from 'src/modules/restaurants/entities/restaurant.entity';

@Injectable()
export class TokensService {
  private readonly logger = new Logger(TokensService.name);
  constructor(
    private readonly configService: ConfigService,
    private readonly jwtService: JwtService,
    private readonly em: EntityManager,
  ) {}

  async generateAccessAndRefreshToken(
    ownerId: string,
    restaurantId: string,
    isAdmin: boolean,
    slug: string,
    em?: EntityManager,
  ) {
    const entityManager = em || this.em;
    const refreshExpire = this.configService.getOrThrow<number>('REFRESH_TOKEN_EXPIRE');
    const accessExpire = this.configService.getOrThrow<number>('JWT_EXPIRATION_TIME');

    const payload: ITokenPayload | IAdminTokenPayload = isAdmin
      ? { adminId: ownerId, restId: restaurantId }
      : { userId: ownerId, restId: restaurantId, slug };

    const accessToken = await this.generateAccessToken(payload, accessExpire);
    const refreshToken = this.generateRefreshToken();
    const type = isAdmin ? RefreshTokenType.ADMIN : RefreshTokenType.USER;

    await this.storeRefreshToken(ownerId, restaurantId, refreshToken, type, entityManager);

    return {
      accessToken: { token: accessToken, expire: dayjs().add(accessExpire, 'seconds').valueOf() },
      refreshToken: { token: refreshToken, expire: dayjs().add(refreshExpire, 'day').valueOf() },
    };
  }

  private generateAccessToken(payload: ITokenPayload | IAdminTokenPayload, expiresIn: number) {
    // Ensure expiresIn is passed as a string with time unit for reliability
    // JWT library accepts: number (seconds) or string with unit (e.g., "3600s", "1h")
    // Using string format is more explicit and prevents unit confusion
    return this.jwtService.signAsync(payload, { expiresIn: `${expiresIn}s` });
  }

  async storeRefreshToken(
    ownerId: string,
    restId: string,
    refreshToken: string,
    type: RefreshTokenType,
    em?: EntityManager,
  ) {
    const entityManager = em || this.em;
    const refreshExpire = this.configService.getOrThrow<number>('REFRESH_TOKEN_EXPIRE');
    const expiresAt = dayjs().add(refreshExpire, 'day').toDate();

    const hashedToken = this.hashToken(refreshToken);

    const token = entityManager.create(RefreshToken, {
      hashedToken,
      ownerId,
      restId,
      type,
      expiresAt,
    });

    if (em) {
      // Within transaction, just persist (flush will be called by transaction)
      entityManager.persist(token);
    } else {
      // Outside transaction, persist and flush
      return entityManager.persistAndFlush(token);
    }
  }

  async refreshToken(oldRefreshToken: string) {
    const hashedToken = this.hashToken(oldRefreshToken);

    // Use transaction to ensure atomicity and prevent race conditions
    return this.em.transactional(async em => {
      // Lock the token row to prevent concurrent refresh attempts
      // Using pessimistic write lock (FOR UPDATE) to prevent race conditions
      const token = await em.findOne(RefreshToken, { hashedToken }, { lockMode: LockMode.PESSIMISTIC_WRITE });

      if (!token) {
        throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN);
      }

      // Check expiration
      if (dayjs(token.expiresAt).isBefore(dayjs())) {
        em.remove(token);
        await em.flush();
        throw new UnauthorizedException(AuthMessage.REFRESH_TOKEN_EXPIRED);
      }

      // Store token data before removal
      const ownerId = token.ownerId;
      const restId = token.restId;
      const isAdmin = token.type === RefreshTokenType.ADMIN;

      // Verify restaurant still exists
      const restaurant = await em.findOne(Restaurant, { id: restId });
      if (!restaurant) {
        throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN);
      }

      try {
        // Generate new tokens first (before removing old token)
        // Pass the transaction EntityManager to ensure operations are within the transaction
        const tokens = await this.generateAccessAndRefreshToken(ownerId, restaurant.id, isAdmin, restaurant.slug, em);

        // Remove old token only after new token is created
        em.remove(token);

        // Persist changes atomically
        await em.flush();

        return tokens;
      } catch (error) {
        this.logger.error('Failed to generate new tokens after refresh', error);
        // Transaction will rollback automatically, preserving the old token
        throw new UnauthorizedException('Failed to refresh token');
      }
    });
  }

  private generateRefreshToken() {
    return randomBytes(32).toString('hex');
  }

  private hashToken(token: string): string {
    return createHash('sha256').update(token).digest('hex');
  }
}