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 { ITokenPayload } from '../interfaces/IToken-payload';

@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 generateTokens(user: User, em?: EntityManager) {
    return this.generateAccessAndRefreshToken(
      {
        id: user.id,
        role: user.role.name,
      },
      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 accessToken = await this.jwtService.signAsync(payload, { expiresIn: `${accessExpire}m` });
    const refreshToken = await this.generateRefreshToken();

    await this.storeRefreshToken(payload.id, 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, 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 token = entityManager.create(RefreshToken, {
        token: refreshToken,
        user,
        expiresAt,
      });

      // Use persist instead of persistAndFlush when in transaction
      if (isExternalTransaction) {
        await 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();

      const token = await entityManager.findOne(
        RefreshToken,
        { token: oldRefreshToken },
        {
          populate: ['user', 'user.role'],
        },
      );

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

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

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

      // Generate new tokens within the same transaction
      const tokens = await this.generateTokens(token.user, 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 async generateRefreshToken() {
    return randomBytes(32).toString('hex');
  }
}