Make JWT key rotation cluster-safe with lock-gated cleanup and active-key convergence

2026-07-15 17:16:51 +01:00
parent 6f6243688d
commit bc9b4c2e6f
5 changed files with 229 additions and 3 deletions
@@ -2,6 +2,7 @@ package stirling.software.proprietary.security.repository;

 import java.time.LocalDateTime;
 import java.util.List;
+import java.util.Optional;

 import org.springframework.data.jpa.repository.JpaRepository;
 import org.springframework.stereotype.Repository;
@@ -15,6 +16,11 @@ public interface JwtSigningKeyRepository extends JpaRepository<JwtSigningKeyEnti
    /** Newest first, so the most recently created key is the active signing key. */
    List<JwtSigningKeyEntity> findAllByOrderByCreatedAtDesc();

+    /**
+     * The current active signing key: the single newest row. Used for cheap cluster convergence.
+     */
+    Optional<JwtSigningKeyEntity> findFirstByOrderByCreatedAtDesc();
+
    /** Keys created before the cutoff, eligible for rotation cleanup. */
    List<JwtSigningKeyEntity> findByCreatedAtBefore(LocalDateTime cutoff);
 }
@@ -3,8 +3,10 @@ package stirling.software.proprietary.security.service;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.time.LocalDateTime;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;

 import org.springframework.beans.factory.annotation.Autowired;
@@ -17,6 +19,8 @@ import jakarta.annotation.PostConstruct;

 import lombok.extern.slf4j.Slf4j;

+import stirling.software.common.cluster.DistributedLock;
+import stirling.software.common.cluster.DistributedLock.LockHandle;
 import stirling.software.common.configuration.InstallationPathConfig;
 import stirling.software.common.model.ApplicationProperties;
 import stirling.software.proprietary.security.model.JwtVerificationKey;
@@ -26,15 +30,23 @@ import stirling.software.proprietary.security.model.JwtVerificationKey;
@ConditionalOnBooleanProperty("v2")
 public class KeyPairCleanupService {

+    // Cluster-wide single-writer: keys live in the shared DB, so only one node may prune + rotate
+    // per cycle. Otherwise every node runs this and they race to delete each other's keys.
+    private static final String CLEANUP_LOCK = "jwt-key-cleanup";
+    private static final Duration LOCK_LEASE = Duration.ofMinutes(5);
+
    private final KeyPersistenceService keyPersistenceService;
    private final ApplicationProperties.Security.Jwt jwtProperties;
+    private final DistributedLock distributedLock;

    @Autowired
    public KeyPairCleanupService(
            KeyPersistenceService keyPersistenceService,
-            ApplicationProperties applicationProperties) {
+            ApplicationProperties applicationProperties,
+            DistributedLock distributedLock) {
        this.keyPersistenceService = keyPersistenceService;
        this.jwtProperties = applicationProperties.getSecurity().getJwt();
+        this.distributedLock = distributedLock;
    }

    @Transactional
@@ -44,7 +56,28 @@ public class KeyPairCleanupService {
        if (!jwtProperties.isEnableKeyCleanup() || !keyPersistenceService.isKeystoreEnabled()) {
            return;
        }
+        // A lock-backend error must never fail this @PostConstruct/scheduled run: degrade to
+        // "skip this cycle" so a transient Valkey blip can't stop a node from booting.
+        Optional<LockHandle> lock;
+        try {
+            lock = distributedLock.tryAcquire(CLEANUP_LOCK, LOCK_LEASE);
+        } catch (RuntimeException e) {
+            log.warn(
+                    "Could not acquire the JWT key-cleanup lock ({}); skipping this cycle",
+                    e.getMessage());
+            return;
+        }
+        // No lock means another node is already pruning; skip until the next tick.
+        if (lock.isEmpty()) {
+            log.debug("Another node holds the JWT key-cleanup lock; skipping this cycle");
+            return;
+        }
+        try (LockHandle held = lock.get()) {
+            runCleanup();
+        }
+    }

+    private void runCleanup() {
        LocalDateTime cutoffDate =
                LocalDateTime.now().minusDays(jwtProperties.getKeyRetentionDays());

@@ -23,9 +23,11 @@ import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;

+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.cache.Cache;
 import org.springframework.cache.CacheManager;
 import org.springframework.context.annotation.DependsOn;
+import org.springframework.scheduling.annotation.Scheduled;
 import org.springframework.stereotype.Service;

 import jakarta.annotation.PostConstruct;
@@ -51,6 +53,7 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
    private final ApplicationProperties.Security.Jwt jwtProperties;
    private final Cache verifyingKeyCache;
    private final JwtSigningKeyRepository keyRepository;
+    private final boolean clusterEnabled;

    // kid -> KeyPair; safe to cache since key material is immutable.
    private final Map<String, KeyPair> keyPairCache = new ConcurrentHashMap<>();
@@ -60,10 +63,12 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
    public KeyPersistenceService(
            ApplicationProperties applicationProperties,
            CacheManager cacheManager,
-            JwtSigningKeyRepository keyRepository) {
+            JwtSigningKeyRepository keyRepository,
+            @Value("${cluster.enabled:false}") boolean clusterEnabled) {
        this.jwtProperties = applicationProperties.getSecurity().getJwt();
        this.verifyingKeyCache = cacheManager.getCache("verifyingKeys");
        this.keyRepository = keyRepository;
+        this.clusterEnabled = clusterEnabled;
    }

    @PostConstruct
@@ -81,6 +86,40 @@ public class KeyPersistenceService implements KeyPersistenceServiceInterface {
        }
    }

+    /**
+     * Cluster convergence: adopt the newest signing key in the shared DB as this node's active key.
+     * Runs on every node so a key a peer just minted becomes the shared active signer within one
+     * interval, keeping cluster rotation equivalent to single-node. Cluster-only: a single node
+     * always holds its own newest key, so this is skipped entirely off-cluster.
+     */
+    @Scheduled(fixedDelayString = "${stirling.security.jwt.activeKeyReloadMs:300000}")
+    public void reloadActiveKeyFromDb() {
+        if (!clusterEnabled || !isKeystoreEnabled()) {
+            return;
+        }
+        try {
+            Optional<JwtSigningKeyEntity> newestOpt =
+                    keyRepository.findFirstByOrderByCreatedAtDesc();
+            if (newestOpt.isEmpty()) {
+                return;
+            }
+            JwtSigningKeyEntity newest = newestOpt.get();
+            JwtVerificationKey current = activeKey;
+            if (current != null && newest.getKeyId().equals(current.getKeyId())) {
+                return;
+            }
+            JwtVerificationKey adopted =
+                    new JwtVerificationKey(newest.getKeyId(), newest.getVerifyingKey());
+            verifyingKeyCache.put(newest.getKeyId(), adopted);
+            activeKey = adopted;
+            log.info(
+                    "Adopted newest JWT signing key {} from the shared DB as active",
+                    newest.getKeyId());
+        } catch (Exception e) {
+            log.warn("Could not reload active JWT key from the shared DB: {}", e.getMessage());
+        }
+    }
+
    /** Load every signing key from the shared DB into the caches; most recent becomes active. */
    private void loadKeysFromDb() {
        List<JwtSigningKeyEntity> keys = keyRepository.findAllByOrderByCreatedAtDesc();
@@ -0,0 +1,102 @@
+package stirling.software.proprietary.security.service;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.lenient;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.List;
+import java.util.Optional;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import stirling.software.common.cluster.DistributedLock;
+import stirling.software.common.cluster.DistributedLock.LockHandle;
+import stirling.software.common.model.ApplicationProperties;
+import stirling.software.proprietary.security.model.JwtVerificationKey;
+
+/** Cluster safety: pruning JWT keys is single-writer, gated on the shared cleanup lock. */
+@ExtendWith(MockitoExtension.class)
+class KeyPairCleanupServiceTest {
+
+    @Mock private KeyPersistenceService keyPersistenceService;
+    @Mock private ApplicationProperties applicationProperties;
+    @Mock private ApplicationProperties.Security security;
+    @Mock private ApplicationProperties.Security.Jwt jwtProperties;
+    @Mock private DistributedLock distributedLock;
+    @Mock private LockHandle lockHandle;
+
+    private KeyPairCleanupService cleanupService;
+
+    @BeforeEach
+    void setUp() {
+        lenient().when(applicationProperties.getSecurity()).thenReturn(security);
+        lenient().when(security.getJwt()).thenReturn(jwtProperties);
+        lenient().when(jwtProperties.isEnableKeyCleanup()).thenReturn(true);
+        lenient().when(keyPersistenceService.isKeystoreEnabled()).thenReturn(true);
+        cleanupService =
+                new KeyPairCleanupService(
+                        keyPersistenceService, applicationProperties, distributedLock);
+    }
+
+    @Test
+    void skipsPruningWhenAnotherNodeHoldsTheLock() {
+        when(distributedLock.tryAcquire(any(), any())).thenReturn(Optional.empty());
+
+        cleanupService.cleanup();
+
+        // No node-local pruning happened; the lock holder owns this cycle.
+        verify(keyPersistenceService, never()).getKeysEligibleForCleanup(any());
+        verify(keyPersistenceService, never()).refreshActiveKeyPair();
+    }
+
+    @Test
+    void prunesAndRotatesWhenLockAcquiredThenReleasesIt() {
+        when(distributedLock.tryAcquire(any(), any())).thenReturn(Optional.of(lockHandle));
+        when(keyPersistenceService.getKeysEligibleForCleanup(any()))
+                .thenReturn(List.of(new JwtVerificationKey("old-key", "cHVi")));
+
+        cleanupService.cleanup();
+
+        verify(keyPersistenceService).removeKey("old-key");
+        verify(keyPersistenceService).refreshActiveKeyPair();
+        verify(lockHandle).close();
+    }
+
+    @Test
+    void releasesTheLockEvenWhenNoKeysAreEligible() {
+        when(distributedLock.tryAcquire(any(), any())).thenReturn(Optional.of(lockHandle));
+        when(keyPersistenceService.getKeysEligibleForCleanup(any())).thenReturn(List.of());
+
+        cleanupService.cleanup();
+
+        verify(keyPersistenceService, never()).refreshActiveKeyPair();
+        verify(lockHandle).close();
+    }
+
+    @Test
+    void skipsPruningWhenTheLockBackendErrors() {
+        // A Valkey blip at boot must not fail startup: tryAcquire throwing degrades to skip.
+        when(distributedLock.tryAcquire(any(), any()))
+                .thenThrow(new RuntimeException("valkey unreachable"));
+
+        cleanupService.cleanup();
+
+        verify(keyPersistenceService, never()).getKeysEligibleForCleanup(any());
+        verify(keyPersistenceService, never()).refreshActiveKeyPair();
+    }
+
+    @Test
+    void doesNothingWhenCleanupDisabled() {
+        when(jwtProperties.isEnableKeyCleanup()).thenReturn(false);
+
+        cleanupService.cleanup();
+
+        verify(distributedLock, never()).tryAcquire(any(), any());
+    }
+}
@@ -56,8 +56,9 @@ class KeyPersistenceServiceInterfaceTest {
        lenient().when(security.getJwt()).thenReturn(jwtConfig);
        lenient().when(jwtConfig.isEnableKeystore()).thenReturn(true);
        lenient().when(keyRepository.save(any())).thenAnswer(inv -> inv.getArgument(0));
+        // clusterEnabled=true so the convergence-reload path is exercised.
        keyPersistenceService =
-                new KeyPersistenceService(applicationProperties, cacheManager, keyRepository);
+                new KeyPersistenceService(applicationProperties, cacheManager, keyRepository, true);
    }

    private JwtSigningKeyEntity entityFrom(String keyId) {
@@ -132,4 +133,49 @@ class KeyPersistenceServiceInterfaceTest {
        assertEquals(1, stale.size());
        assertEquals("old-key", stale.get(0).getKeyId());
    }
+
+    @Test
+    void reloadAdoptsTheNewestKeyAPeerMinted() {
+        // Boot with our own key active, then a peer mints a newer one in the shared DB.
+        when(keyRepository.count()).thenReturn(1L);
+        when(keyRepository.findAllByOrderByCreatedAtDesc())
+                .thenReturn(List.of(entityFrom("jwt-key-local-old")));
+        keyPersistenceService.initializeKeystore();
+        assertEquals("jwt-key-local-old", keyPersistenceService.getActiveKey().getKeyId());
+
+        when(keyRepository.findFirstByOrderByCreatedAtDesc())
+                .thenReturn(Optional.of(entityFrom("jwt-key-peer-new")));
+
+        keyPersistenceService.reloadActiveKeyFromDb();
+
+        // Converged: this node now signs with the peer's newer key.
+        assertEquals("jwt-key-peer-new", keyPersistenceService.getActiveKey().getKeyId());
+    }
+
+    @Test
+    void reloadDoesNothingOffCluster() {
+        KeyPersistenceService singleNode =
+                new KeyPersistenceService(
+                        applicationProperties, cacheManager, keyRepository, false);
+
+        singleNode.reloadActiveKeyFromDb();
+
+        // Off-cluster the DB is never consulted for convergence.
+        verify(keyRepository, org.mockito.Mockito.never()).findFirstByOrderByCreatedAtDesc();
+    }
+
+    @Test
+    void reloadIsANoOpWhenAlreadyHoldingTheNewestKey() {
+        when(keyRepository.count()).thenReturn(1L);
+        when(keyRepository.findAllByOrderByCreatedAtDesc())
+                .thenReturn(List.of(entityFrom("jwt-key-current")));
+        keyPersistenceService.initializeKeystore();
+
+        when(keyRepository.findFirstByOrderByCreatedAtDesc())
+                .thenReturn(Optional.of(entityFrom("jwt-key-current")));
+
+        keyPersistenceService.reloadActiveKeyFromDb();
+
+        assertEquals("jwt-key-current", keyPersistenceService.getActiveKey().getKeyId());
+    }
 }