diff --git a/src/modules/auth/services/tokens.service.ts b/src/modules/auth/services/tokens.service.ts index a482ecf..eff89e0 100755 --- a/src/modules/auth/services/tokens.service.ts +++ b/src/modules/auth/services/tokens.service.ts @@ -1,5 +1,6 @@ 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'; @@ -19,7 +20,14 @@ export class TokensService { private readonly em: EntityManager, ) {} - async generateAccessAndRefreshToken(ownerId: string, restaurantId: string, isAdmin: boolean, slug: string) { + 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'); @@ -31,7 +39,7 @@ export class TokensService { const refreshToken = this.generateRefreshToken(); const type = isAdmin ? RefreshTokenType.ADMIN : RefreshTokenType.USER; - await this.storeRefreshToken(ownerId, restaurantId, refreshToken, type); + await this.storeRefreshToken(ownerId, restaurantId, refreshToken, type, entityManager); return { accessToken: { token: accessToken, expire: dayjs().add(accessExpire, 'seconds').valueOf() }, @@ -46,13 +54,20 @@ export class TokensService { return this.jwtService.signAsync(payload, { expiresIn: `${expiresIn}s` }); } - async storeRefreshToken(ownerId: string, restId: string, refreshToken: string, type: RefreshTokenType) { + 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 = this.em.create(RefreshToken, { + const token = entityManager.create(RefreshToken, { hashedToken, ownerId, restId, @@ -60,46 +75,64 @@ export class TokensService { expiresAt, }); - return this.em.persistAndFlush(token); + 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); - const token = await this.em.findOne(RefreshToken, { hashedToken }); - if (!token) throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN); + // 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 (dayjs(token.expiresAt).isBefore(dayjs())) { - this.em.remove(token); - await this.em.flush(); - throw new UnauthorizedException(AuthMessage.REFRESH_TOKEN_EXPIRED); - } + if (!token) { + throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN); + } - // Store token data before removal - const ownerId = token.ownerId; - const restId = token.restId; + // Check expiration + if (dayjs(token.expiresAt).isBefore(dayjs())) { + em.remove(token); + await em.flush(); + throw new UnauthorizedException(AuthMessage.REFRESH_TOKEN_EXPIRED); + } - const restaurant = await this.em.findOne(Restaurant, { id: restId }); - if (!restaurant) { - throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN); - } - // Remove the old token - this.em.remove(token); - const isAdmin = token.type === RefreshTokenType.ADMIN; - try { - // Generate new tokens - const tokens = await this.generateAccessAndRefreshToken(ownerId, restId, isAdmin, restaurant.slug); + // Store token data before removal + const ownerId = token.ownerId; + const restId = token.restId; + const isAdmin = token.type === RefreshTokenType.ADMIN; - // Commit the transaction - await this.em.flush(); + // Verify restaurant still exists + const restaurant = await em.findOne(Restaurant, { id: restId }); + if (!restaurant) { + throw new UnauthorizedException(AuthMessage.INVALID_REFRESH_TOKEN); + } - return tokens; - } catch (error) { - // If token generation fails, the old token is already removed - // This is acceptable as the token was already used - this.logger.error('Failed to generate new tokens after refresh', error); - throw new UnauthorizedException('Failed to 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() {