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