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 { AuthMessage } from '../../../common/enums/message.enum'; import { RefreshToken } from '../entities/refresh-token.entity'; import { IAdminTokenPayload } from '../interfaces/IToken-payload'; import { JwtService } from '@nestjs/jwt'; import dayjs from 'dayjs'; @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( adminId: string, businessId: string, em?: EntityManager, ) { const refreshExpire = this.configService.getOrThrow('REFRESH_TOKEN_EXPIRE'); const accessExpire = this.configService.getOrThrow('JWT_EXPIRATION_TIME'); const payload: IAdminTokenPayload = { adminId, businessId } const accessToken = await this.generateAccessToken(payload, accessExpire); const refreshToken = this.generateRefreshToken(); // Only pass em if it's a transaction manager (not this.em), otherwise pass undefined to trigger flush await this.storeRefreshToken(adminId, businessId, refreshToken, em); return { accessToken: { token: accessToken, expire: dayjs().add(accessExpire, 'seconds').valueOf() }, refreshToken: { token: refreshToken, expire: dayjs().add(refreshExpire, 'day').valueOf() }, }; } private generateAccessToken(payload: 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, em?: EntityManager, ) { const entityManager = em || this.em; const refreshExpire = this.configService.getOrThrow('REFRESH_TOKEN_EXPIRE'); const expiresAt = dayjs().add(refreshExpire, 'day').toDate(); const hashedToken = this.hashToken(refreshToken); const token = entityManager.create(RefreshToken, { hashedToken, adminId: ownerId, businessId: restId, expiresAt, }); if (em) { // Within transaction, just persist (flush will be called by transaction) entityManager.persist(token); } else { // Outside transaction, persist and flush immediately await 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 adminId = token.adminId; const businessId = token.businessId; 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(adminId, businessId, 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'); } }