fix(authz): keep a rotated refresh token valid for a short grace window CU-17tkuw62gnb

openframe-authorization-service-core/src/main/java/com/openframe/authz/service/auth/MongoAuthorizationService.java

@@ -1,9 +1,12 @@
 package com.openframe.authz.service.auth;
 
 import com.openframe.data.document.oauth.MongoOAuth2Authorization;
+import com.openframe.data.document.oauth.RetiredRefreshToken;
 import com.openframe.data.repository.oauth.MongoOAuth2AuthorizationRepository;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.security.oauth2.core.OAuth2RefreshToken;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
 import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
 import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationCode;
@@ -12,6 +15,9 @@
 import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
 import org.springframework.stereotype.Service;
 
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.stream.Stream;
@@ -27,14 +33,49 @@
 
     private static final OAuth2TokenType AUTH_CODE = new OAuth2TokenType("code");
 
+    @Value("${openframe.authz.refresh-token-grace-seconds:60}")
+    private long refreshTokenGraceSeconds;
+
     @Override
     public void save(OAuth2Authorization authorization) {
         log.debug("Saving authorization: {}", authorization.getId());
 
         MongoOAuth2Authorization entity = MongoAuthorizationMapper.toEntity(authorization);
+        entity.setRetiredRefreshTokens(retiredTokensAfterSaving(entity));
         repository.save(entity);
     }
 
+    /**
+     * Rotation replaces the refresh token, and the client only learns the new one from a response it
+     * may never receive: a cancelled request or a dropped connection, while the server still finishes.
+     * Keeping the replaced token for a short window lets that client's next refresh recover the session
+     * instead of ending it. A list rather than one slot, so a retry that rotates again doesn't kill the
+     * token a concurrent request has just received.
+     */
+    private List<RetiredRefreshToken> retiredTokensAfterSaving(MongoOAuth2Authorization entity) {
+        MongoOAuth2Authorization stored = repository.findById(entity.getId()).orElse(null);
+        if (stored == null) {
+            return null;
+        }
+        Instant now = Instant.now();
+        List<RetiredRefreshToken> retired = new ArrayList<>();
+        if (stored.getRetiredRefreshTokens() != null) {
+            stored.getRetiredRefreshTokens().stream()
+                    .filter(token -> isWithinGrace(token, now))
+                    .filter(token -> !token.value().equals(entity.getRefreshTokenValue()))
+                    .forEach(retired::add);
+        }
+        String replaced = stored.getRefreshTokenValue();
+        if (replaced != null && !replaced.equals(entity.getRefreshTokenValue())) {
+            retired.add(new RetiredRefreshToken(replaced, now));
+        }
+        return retired.isEmpty() ? null : retired;
+    }
+
+    private boolean isWithinGrace(RetiredRefreshToken token, Instant now) {
+        return token.retiredAt().plusSeconds(refreshTokenGraceSeconds).isAfter(now);
+    }
+
     @Override
     public void remove(OAuth2Authorization authorization) {
         repository.deleteById(authorization.getId());
@@ -97,10 +138,16 @@
                     .or(() -> repository.findByRefreshTokenValue(token))
                     .or(() -> repository.findByAuthorizationCodeValue(token))
                     .or(() -> repository.findByState(token));
+            if (found.isEmpty()) {
+                return findByRetiredRefreshToken(token);
+            }
         } else if (OAuth2TokenType.ACCESS_TOKEN.equals(tokenType)) {
             found = repository.findByAccessTokenValue(token);
         } else if (OAuth2TokenType.REFRESH_TOKEN.equals(tokenType)) {
             found = repository.findByRefreshTokenValue(token);
+            if (found.isEmpty()) {
+                return findByRetiredRefreshToken(token);
+            }
         } else if (AUTH_CODE.equals(tokenType)) {
             found = repository.findByAuthorizationCodeValue(token);
         } else {
@@ -110,6 +157,28 @@
         return found.map(entity -> MongoAuthorizationMapper.toDomain(entity, registeredClientRepository))
                 .orElse(null);
     }
+
+    /**
+     * The presented value goes onto the returned authorization because Spring's revocation looks the
+     * token up by value on it; the metadata stays the current token's, so a revoked session stays revoked.
+     */
+    private OAuth2Authorization findByRetiredRefreshToken(String token) {
+        Instant now = Instant.now();
+        return repository.findByRetiredRefreshTokensValue(token)
+                .filter(entity -> entity.getRetiredRefreshTokens().stream()
+                        .anyMatch(retired -> retired.value().equals(token) && isWithinGrace(retired, now)))
+                .map(entity -> {
+                    log.info("Accepted a rotated refresh token within the grace window: authorization {}", entity.getId());
+                    OAuth2Authorization authorization = MongoAuthorizationMapper.toDomain(entity, registeredClientRepository);
+                    OAuth2Authorization.Token<OAuth2RefreshToken> current = authorization.getRefreshToken();
+                    OAuth2RefreshToken presented = new OAuth2RefreshToken(
+                            token, current.getToken().getIssuedAt(), current.getToken().getExpiresAt());
+                    return OAuth2Authorization.from(authorization)
+                            .token(presented, metadata -> metadata.putAll(current.getMetadata()))
+                            .build();
+                })
+                .orElse(null);
+    }
 }
 
 

openframe-authorization-service-core/src/test/java/com/openframe/authz/service/auth/MongoAuthorizationServiceRefreshGraceTest.java

@@ -0,0 +1,231 @@
+package com.openframe.authz.service.auth;
+
+import com.openframe.data.document.oauth.MongoOAuth2Authorization;
+import com.openframe.data.document.oauth.RetiredRefreshToken;
+import com.openframe.data.repository.oauth.MongoOAuth2AuthorizationRepository;
+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 org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
+import org.springframework.security.oauth2.core.AuthorizationGrantType;
+import org.springframework.security.oauth2.core.ClientAuthenticationMethod;
+import org.springframework.security.oauth2.core.OAuth2RefreshToken;
+import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
+import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
+import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
+import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+@MockitoSettings(strictness = Strictness.LENIENT)
+class MongoAuthorizationServiceRefreshGraceTest {
+
+    private static final String ID = "authorization-1";
+    private static final String PRINCIPAL = "user@example.com";
+    private static final long GRACE_SECONDS = 60;
+
+    @Mock
+    private MongoOAuth2AuthorizationRepository repository;
+    @Mock
+    private RegisteredClientRepository clients;
+
+    private final Map<String, MongoOAuth2Authorization> store = new HashMap<>();
+    private MongoAuthorizationService service;
+
+    @BeforeEach
+    void setUp() {
+        when(clients.findById(anyString())).thenReturn(RegisteredClient.withId("client-1")
+                .clientId("client")
+                .clientAuthenticationMethod(ClientAuthenticationMethod.CLIENT_SECRET_BASIC)
+                .authorizationGrantType(AuthorizationGrantType.REFRESH_TOKEN)
+                .build());
+        when(repository.save(any())).thenAnswer(call -> {
+            MongoOAuth2Authorization entity = call.getArgument(0);
+            store.put(entity.getId(), entity);
+            return entity;
+        });
+        when(repository.findById(anyString()))
+                .thenAnswer(call -> Optional.ofNullable(store.get(call.<String>getArgument(0))));
+        when(repository.findByRefreshTokenValue(anyString())).thenAnswer(call -> store.values().stream()
+                .filter(e -> call.<String>getArgument(0).equals(e.getRefreshTokenValue()))
+                .findFirst());
+        when(repository.findByRetiredRefreshTokensValue(anyString())).thenAnswer(call -> store.values().stream()
+                .filter(e -> retiredValues().contains(call.<String>getArgument(0)))
+                .findFirst());
+        when(repository.findByAccessTokenValue(anyString())).thenReturn(Optional.empty());
+        when(repository.findByAuthorizationCodeValue(anyString())).thenReturn(Optional.empty());
+        when(repository.findByState(anyString())).thenReturn(Optional.empty());
+        when(repository.findAllByPrincipalNameIgnoreCaseAndRefreshTokenValueNotNull(PRINCIPAL))
+                .thenAnswer(call -> new ArrayList<>(store.values()));
+
+        service = new MongoAuthorizationService(repository, clients);
+        ReflectionTestUtils.setField(service, "refreshTokenGraceSeconds", GRACE_SECONDS);
+        store.put(ID, authorizationWithRefreshToken("token-0"));
+    }
+
+    private static MongoOAuth2Authorization authorizationWithRefreshToken(String token) {
+        Instant now = Instant.now();
+        return MongoOAuth2Authorization.builder()
+                .id(ID)
+                .registeredClientId("client-1")
+                .principalName(PRINCIPAL)
+                .authorizationGrantType(AuthorizationGrantType.AUTHORIZATION_CODE.getValue())
+                .refreshTokenValue(token)
+                .refreshTokenIssuedAt(now)
+                .refreshTokenExpiresAt(now.plus(7, ChronoUnit.DAYS))
+                .refreshTokenMetadata(new HashMap<>())
+                .build();
+    }
+
+    private MongoOAuth2Authorization stored() {
+        return store.get(ID);
+    }
+
+    private List<String> retiredValues() {
+        List<RetiredRefreshToken> retired = stored().getRetiredRefreshTokens();
+        return retired == null ? List.of() : retired.stream().map(RetiredRefreshToken::value).toList();
+    }
+
+    private OAuth2Authorization find(String token) {
+        return service.findByToken(token, OAuth2TokenType.REFRESH_TOKEN);
+    }
+
+    private void rotate(String presented, String replacement) {
+        Instant now = Instant.now();
+        service.save(OAuth2Authorization.from(find(presented))
+                .token(new OAuth2RefreshToken(replacement, now, now.plus(7, ChronoUnit.DAYS)))
+                .build());
+    }
+
+    private void retire(String value, long secondsAgo) {
+        stored().setRetiredRefreshTokens(new ArrayList<>(List.of(
+                new RetiredRefreshToken(value, Instant.now().minusSeconds(secondsAgo)))));
+    }
+
+    @Test
+    void rejectsARetiredTokenOnceTheGraceWindowHasPassed() {
+        retire("expired", GRACE_SECONDS + 1);
+
+        assertThat(find("expired")).isNull();
+    }
+
+    @Test
+    void acceptsARetiredTokenJustInsideTheGraceWindow() {
+        retire("recent", GRACE_SECONDS - 5);
+
+        assertThat(find("recent")).isNotNull();
+    }
+
+    @Test
+    void rejectsATokenThatWasNeverIssued() {
+        assertThat(find("never-issued")).isNull();
+    }
+
+    @Test
+    void aRetryThatRotatesAgainKeepsEveryRetiredTokenUsable() {
+        rotate("token-0", "token-1");
+        rotate("token-0", "token-2");
+
+        assertThat(stored().getRefreshTokenValue()).isEqualTo("token-2");
+        assertThat(retiredValues()).containsExactlyInAnyOrder("token-0", "token-1");
+        assertThat(find("token-0").getRefreshToken().isActive()).isTrue();
+        assertThat(find("token-1").getRefreshToken().isActive()).isTrue();
+    }
+
+    @Test
+    void aRetiredTokenCarriesItsOwnValueSoSpringRevocationCanLookItUp() {
+        rotate("token-0", "token-1");
+
+        assertThat(find("token-0").getToken("token-0")).isNotNull();
+    }
+
+    @Test
+    void aRetiredTokenKeepsTheCurrentExpiry() {
+        rotate("token-0", "token-1");
+
+        assertThat(find("token-0").getRefreshToken().getToken().getExpiresAt())
+                .isEqualTo(stored().getRefreshTokenExpiresAt());
+    }
+
+    @Test
+    void revokingTheSessionAlsoEndsItsRetiredTokens() {
+        rotate("token-0", "token-1");
+
+        service.revokeAllForPrincipal(PRINCIPAL);
+
+        assertThat(find("token-0").getRefreshToken().isActive()).isFalse();
+    }
+
+    @Test
+    void revokingThroughARetiredTokenLeavesNoActiveTokenBehind() {
+        rotate("token-0", "token-1");
+        OAuth2Authorization byRetired = find("token-0");
+
+        service.save(OAuth2Authorization.from(byRetired)
+                .token(byRetired.getToken("token-0").getToken(),
+                        metadata -> metadata.put(OAuth2Authorization.Token.INVALIDATED_METADATA_NAME, true))
+                .build());
+
+        assertThat(find("token-0").getRefreshToken().isActive()).isFalse();
+        assertThat(find("token-1").getRefreshToken().isActive()).isFalse();
+        assertThat(retiredValues()).containsExactly("token-1");
+    }
+
+    @Test
+    void savingWithoutRotatingRetiresNothingAndKeepsWhatIsRetired() {
+        retire("recent", 1);
+
+        service.save(find("token-0"));
+
+        assertThat(retiredValues()).containsExactly("recent");
+    }
+
+    @Test
+    void savingDropsRetiredTokensPastTheirWindow() {
+        retire("expired", GRACE_SECONDS + 1);
+
+        rotate("token-0", "token-1");
+
+        assertThat(retiredValues()).containsExactly("token-0");
+    }
+
+    @Test
+    void savingWithNothingRetiredStoresNoList() {
+        service.save(find("token-0"));
+
+        assertThat(stored().getRetiredRefreshTokens()).isNull();
+    }
+
+    @Test
+    void savingANewAuthorizationStoresNoList() {
+        OAuth2Authorization existing = find("token-0");
+        store.clear();
+
+        service.save(existing);
+
+        assertThat(stored().getRetiredRefreshTokens()).isNull();
+    }
+
+    @Test
+    void anUntypedLookupAlsoResolvesARetiredToken() {
+        rotate("token-0", "token-1");
+
+        assertThat(service.findByToken("token-0", null)).isNotNull();
+    }
+}