add lock to read by Hash

src/main/java/com/ddf/vodsystem/repositories/RefreshTokenRepository.java

@@ -1,15 +1,17 @@
 package com.ddf.vodsystem.repositories;
 
 import com.ddf.vodsystem.entities.RefreshToken;
+import jakarta.persistence.LockModeType;
 import org.springframework.data.jpa.repository.JpaRepository;
+import org.springframework.data.jpa.repository.Lock;
 import org.springframework.data.jpa.repository.Query;
 import org.springframework.stereotype.Repository;
 
-import java.sql.Ref;
 import java.util.Optional;
 
 @Repository
 public interface RefreshTokenRepository extends JpaRepository<RefreshToken, Long> {
+    @Lock(LockModeType.PESSIMISTIC_WRITE)
     @Query("SELECT r FROM RefreshToken r WHERE r.tokenHash = ?1")
     Optional<RefreshToken> findByHash(byte[] tokenHash);
 }

src/main/java/com/ddf/vodsystem/security/TokenService.java

@@ -16,7 +16,6 @@ import java.security.SecureRandom;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.HexFormat;
-import java.util.Optional;
 
 @Service
 public class TokenService {
@@ -83,21 +82,17 @@ public class TokenService {
                 .orElseThrow(() -> new NotAuthenticated("No such refresh token"));
 
         // Find Token Family
-        Optional<TokenFamily> tokenFamily = tokenFamilyRepository.findById(refreshToken.getTokenFamily().getId());
+        TokenFamily tokenFamily = refreshToken.getTokenFamily();
 
-        if (tokenFamily.isEmpty()) {
+        if (tokenFamily.getRevokedAt() != null) {
-            throw new IllegalStateException("Token must be tied to a family");
-        }
-
-        if (tokenFamily.get().getRevokedAt() != null) {
             throw new NotAuthenticated("Token family has been revoked");
         }
 
         // Reuse Check
         if (refreshToken.getUsedAt() != null) {
             if (now.isAfter(refreshToken.getUsedAt().plus(refreshGracePeriod))) {
-                tokenFamily.get().setRevokedAt(Instant.now());
+                tokenFamily.setRevokedAt(Instant.now());
-                tokenFamilyRepository.save(tokenFamily.get());
+                tokenFamilyRepository.save(tokenFamily);
             }
 
             throw new NotAuthenticated("Refresh token already being used");
@@ -106,7 +101,7 @@ public class TokenService {
         // Expiry Check
         if (
                 now.isAfter(refreshToken.getExpiresAt()) ||
-                now.isAfter(tokenFamily.get().getExpiresAt())
+                now.isAfter(tokenFamily.getExpiresAt())
         ) {
             throw new NotAuthenticated("Refresh token or family expired");
         }
@@ -118,10 +113,10 @@ public class TokenService {
         byte[] newRawToken = generateRandomBytes(32);
         createRefreshToken(
                 newRawToken,
-                tokenFamily.get()
+                tokenFamily
         );
 
-        return new Rotation(tokenFamily.get().getUser().getId(), HEX.formatHex(newRawToken));
+        return new Rotation(tokenFamily.getUser().getId(), HEX.formatHex(newRawToken));
     }
 
     public record Rotation(Long userId, String newRefreshToken) {}
@@ -152,7 +147,7 @@ public class TokenService {
         refreshToken.setTokenHash(sha256(token));
         refreshToken.setCreatedAt(now);
         refreshToken.setExpiresAt(now.plusMillis(refreshTokenExpirationMs.toMillis()));
-        return refreshTokenRepository.saveAndFlush(refreshToken);
+        return refreshTokenRepository.save(refreshToken);
     }
 
 }