refresh token bug

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