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