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

import { randomBytes } from 'node:crypto';

import { EntityManager } from '@mikro-orm/postgresql';
import { BadRequestException, Injectable, Logger, UnauthorizedException } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { JwtService } from '@nestjs/jwt';
import dayjs from 'dayjs';

import { AuthMessage, UserMessage } from '../../../common/enums/message.enum';
import { RefreshToken } from '../../users/entities/refresh-token.entity';
import { User } from '../../users/entities/user.entity';
import { IAdminTokenPayload, ITokenPayload } from '../interfaces/IToken-payload';
import { Restaurant } from 'src/modules/restaurants/entities/restaurant.entity';
import { RestRepository } from 'src/modules/restaurants/repositories/rest.repository';
import { Admin } from 'src/modules/admin/entities/admin.entity';
import { CacheService } from 'src/modules/utils/cache.service';

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

  async generateTokens(user: User, rest: Restaurant, em?: EntityManager) {
    return this.generateAccessAndRefreshToken(
      {
        userId: user.id,
        restId: rest.id,
      },
      em,
    );
  }

  async generateTokensAdmin(admin: Admin, rest: Restaurant, em?: EntityManager) {
    return this.generateAccessAndRefreshTokenAdmin(
      {
        adminId: admin.id,
        restId: rest.id,
      },
      em,
    );
  }

  private async generateAccessAndRefreshToken(payload: ITokenPayload, em?: EntityManager) {
    const accessExpire = this.configService.getOrThrow<number>('ACCESS_TOKEN_EXPIRE');
    const refreshExpire = this.configService.getOrThrow<number>('REFRESH_TOKEN_EXPIRE');

    const tokenPayload: ITokenPayload = {
      userId: payload.userId,
      restId: payload.restId,
    };
    const accessToken = await this.jwtService.signAsync(tokenPayload, { expiresIn: `${accessExpire}m` });
    const refreshToken = this.generateRefreshToken();

    await this.storeRefreshToken(payload.userId, payload.restId, refreshToken, em);

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

  private async generateAccessAndRefreshTokenAdmin(payload: IAdminTokenPayload, em?: EntityManager) {
    const accessExpire = this.configService.getOrThrow<number>('ACCESS_TOKEN_EXPIRE');
    const refreshExpire = this.configService.getOrThrow<number>('REFRESH_TOKEN_EXPIRE');

    const tokenPayload: IAdminTokenPayload = {
      adminId: payload.adminId,
      restId: payload.restId,
    };
    const accessToken = await this.jwtService.signAsync(tokenPayload, { expiresIn: `${accessExpire}m` });
    const refreshToken = this.generateRefreshToken();

    await this.storeRefreshToken(payload.adminId, payload.restId, refreshToken, em);
    return {
      accessToken: { token: accessToken, expire: dayjs().add(accessExpire, 'minute').valueOf() },
      refreshToken: { token: refreshToken, expire: dayjs().add(refreshExpire, 'day').valueOf() },
    };
  }

  async storeRefreshToken(userId: string, restId: string, refreshToken: string, em?: EntityManager) {
    const entityManager = em || this.em.fork();
    const isExternalTransaction = !!em && em.isInTransaction();

    try {
      // Only start transaction if no external EntityManager with transaction is provided
      if (!isExternalTransaction && !entityManager.isInTransaction()) {
        await entityManager.begin();
      }

      const refreshExpire = this.configService.getOrThrow<number>('REFRESH_TOKEN_EXPIRE');
      const expiresAt = dayjs().add(refreshExpire, 'day').toDate();

      const user = await entityManager.findOne(User, { id: userId });
      if (!user) throw new BadRequestException(UserMessage.USER_NOT_FOUND);

      const restaurant = await entityManager.findOne(Restaurant, { id: restId });
      if (!restaurant) throw new BadRequestException(UserMessage.Rest_NOT_FOUND);

      const token = entityManager.create(RefreshToken, {
        token: refreshToken,
        restaurant,
        user,
        expiresAt,
      });

      // Use persist instead of persistAndFlush when in transaction
      if (isExternalTransaction) {
        entityManager.persist(token);
      } else {
        await entityManager.persistAndFlush(token);
        if (entityManager.isInTransaction()) {
          await entityManager.commit();
        }
      }
    } catch (error) {
      this.logger.error(error);
      // Only rollback if we started the transaction
      if (!isExternalTransaction && entityManager.isInTransaction()) {
        await entityManager.rollback();
      }
      throw error;
    }
  }

  async refreshToken(oldRefreshToken: string) {
    const entityManager = this.em.fork();

    try {
      await entityManager.begin();
      console.log('oldRefreshToken', oldRefreshToken);
      const token = await entityManager.findOne(
        RefreshToken,
        { token: oldRefreshToken },
        {
          populate: ['user', 'restaurant'],
        },
      );

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

      // const rest = await this.restRepository.findOne({ id: token.user.id });
      // if (!rest) throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN);

      if (dayjs(token.expiresAt).isBefore(dayjs())) {
        entityManager.remove(token);
        await entityManager.flush();
        await entityManager.commit();
        throw new UnauthorizedException(AuthMessage.REFRESH_TOKEN_EXPIRED);
      }

      // Remove the old token
      entityManager.remove(token);

      // Generate new tokens within the same transaction
      const tokens = await this.generateTokens(token.user, token.restaurant, entityManager);

      // Commit the transaction
      await entityManager.flush();
      await entityManager.commit();

      return tokens;
    } catch (error) {
      this.logger.error(error);
      if (entityManager.isInTransaction()) {
        await entityManager.rollback();
      }
      throw error;
    }
  }

  async invalidateRefreshToken(userId: string) {
    const entityManager = this.em.fork();

    try {
      await entityManager.begin();

      await entityManager.nativeDelete(RefreshToken, { user: { id: userId } });
      this.logger.log(`Successfully deleted all refresh tokens for user ${userId}`);

      await entityManager.commit();
    } catch (error) {
      this.logger.error(error);
      if (entityManager.isInTransaction()) {
        await entityManager.rollback();
      }
      throw error;
    }
  }

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