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('ACCESS_TOKEN_EXPIRE'); const refreshExpire = this.configService.getOrThrow('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('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'); } }