From 40fc80537bcd5b0c501af12b50c0620f23c2d7d9 Mon Sep 17 00:00:00 2001
From: Dylan De Faoite <dylanseandefaoite@gmail.com>
Date: Wed, 7 Oct 2026 18:31:47 +0100
Subject: [PATCH] add lock to read by Hash

---
 .../repositories/RefreshTokenRepository.java  |  4 +++-
 .../ddf/vodsystem/security/TokenService.java  | 21 +++++++------------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/main/java/com/ddf/vodsystem/repositories/RefreshTokenRepository.java b/src/main/java/com/ddf/vodsystem/repositories/RefreshTokenRepository.java
index 4b925cd..c8542d6 100644
--- a/src/main/java/com/ddf/vodsystem/repositories/RefreshTokenRepository.java
+++ b/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);
 }
diff --git a/src/main/java/com/ddf/vodsystem/security/TokenService.java b/src/main/java/com/ddf/vodsystem/security/TokenService.java
index 099d0d8..96721e3 100644
--- a/src/main/java/com/ddf/vodsystem/security/TokenService.java
+++ b/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()) {
-            throw new IllegalStateException("Token must be tied to a family");
-        }
-
-        if (tokenFamily.get().getRevokedAt() != null) {
+        if (tokenFamily.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());
-                tokenFamilyRepository.save(tokenFamily.get());
+                tokenFamily.setRevokedAt(Instant.now());
+                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);
     }
 
 }