Text case to illustrate the context of spring-security#13915

SpringSecurityIssue13915Test.java

package com.example.spring.security;

import org.junit.jupiter.api.Test;
import org.springframework.security.authentication.AnonymousAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.authority.AuthorityUtils;
import org.springframework.security.oauth2.client.*;
import org.springframework.security.oauth2.client.endpoint.AbstractOAuth2AuthorizationGrantRequest;
import org.springframework.security.oauth2.client.registration.ClientRegistration;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.OAuth2AuthorizationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.endpoint.OAuth2AccessTokenResponse;
import reactor.core.publisher.Mono;

import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.util.Objects;
import java.util.TimeZone;
import java.util.function.Consumer;

import static org.junit.jupiter.api.Assertions.*;
import static org.springframework.security.oauth2.client.OAuth2AuthorizationContext.PASSWORD_ATTRIBUTE_NAME;
import static org.springframework.security.oauth2.client.OAuth2AuthorizationContext.USERNAME_ATTRIBUTE_NAME;
import static org.springframework.security.oauth2.core.OAuth2AccessToken.TokenType.BEARER;
import static org.springframework.security.oauth2.core.OAuth2ErrorCodes.INVALID_GRANT;

@SuppressWarnings({"deprecation", "SameParameterValue"})
public class SpringSecurityIssue13915Test {
    private static final ClientRegistration CLIENT_REGISTRATION = ClientRegistration.withRegistrationId("id")
            .authorizationGrantType(AuthorizationGrantType.PASSWORD)
            .clientId("156a4261-c1bc-47f1-b9ce-0b40b63cef1f")
            .tokenUri("https://example.com/token")
            .build();
    private static final String USERNAME = "user";
    private static final String PASSWORD = "password";
    private static final String REFRESH_TOKEN = "refresh";
    private static final long TOKEN_EXPIRES_IN_SECONDS = 60;
    private static final Authentication PRINCIPAL = new AnonymousAuthenticationToken(
            "anonymous", "anonymousUser", AuthorityUtils.createAuthorityList("ROLE_USER"));
    private static final OAuth2AuthorizationContext EMPTY_CONTEXT = OAuth2AuthorizationContext.withClientRegistration(CLIENT_REGISTRATION)
            .principal(PRINCIPAL)
            .attribute(USERNAME_ATTRIBUTE_NAME, USERNAME)
            .attribute(PASSWORD_ATTRIBUTE_NAME, PASSWORD)
            .build();

    /**
     * Test case where the server does not return a refresh_token
     */
    @Test
    public void testNoRefresh() {
        // Initialize the clock and the mock client
        MockClock mockClock = new MockClock();
        Instant start = mockClock.instant();

        MockClient mock = new MockClient("a", null);

        // 1a: create a password provider
        ReactiveOAuth2AuthorizedClientProvider passwordProvider = getPasswordProvider(mock, mockClock);
        // 1b: create a refresh provider
        ReactiveOAuth2AuthorizedClientProvider refreshProvider = getRefreshProvider(mock, mockClock);

        // 2: Password provider can authorize the client
        OAuth2AuthorizedClient client = passwordProvider.authorize(EMPTY_CONTEXT).block();
        assertNotNull(client);
        assertHasAccessToken("password-a", client);
        assertNull(client.getRefreshToken()); // And we don't get a refresh_token

        // Create a new context with the authorized client
        OAuth2AuthorizationContext authorizedContext = createAuthorizedContext(client);

        mock.nextAccessToken = "b";
        // 3: Before the token is expired...
        mockClock.setInstant(start.plusSeconds(TOKEN_EXPIRES_IN_SECONDS / 2));
        // ...the providers should return nothing
        assertNull(passwordProvider.authorize(authorizedContext).block(), "Should not request a new token yet");
        assertNull(refreshProvider.authorize(authorizedContext).block(), "We should not be able to refresh at all");

        mock.nextAccessToken = "c";
        // 4: After the token is expired...
        mockClock.setInstant(start.plusSeconds(TOKEN_EXPIRES_IN_SECONDS * 2));

        // ... the refresh provider does not work with a client that does not have a refresh_token ...
        assertNull(refreshProvider.authorize(authorizedContext).block(), "refreshProvider should not work");
        // ... but if there is no refresh token, the password provider does return a token
        assertHasAccessToken("password-c", passwordProvider.authorize(authorizedContext).block());
    }

    /**
     * Test where the server returns a refresh_token, but the provider is not set up
     * on the builder.
     *
     * @see <a href="https://github.com/spring-projects/spring-security/issues/8831">spring-security#8831</a>
     */
    @Test
    public void testRefresh() {
        // Initialize the clock and the mock client
        MockClock mockClock = new MockClock();
        Instant start = mockClock.instant();

        MockClient mock = new MockClient("a", REFRESH_TOKEN);

        // 1a: create a password provider
        ReactiveOAuth2AuthorizedClientProvider passwordProvider = getPasswordProvider(mock, mockClock);
        // 1b: create a refresh provider
        ReactiveOAuth2AuthorizedClientProvider refreshProvider = getRefreshProvider(mock, mockClock);

        // 2: Password provider can authorize the client
        OAuth2AuthorizedClient client = passwordProvider.authorize(EMPTY_CONTEXT).block();
        assertHasAccessToken("password-a", client);
        // ...and we get a refresh_token
        assertHasRefreshToken(REFRESH_TOKEN, client);

        // Create a new context with the authorized client
        OAuth2AuthorizationContext authorizedContext = createAuthorizedContext(client);

        mock.nextAccessToken = "b";
        // 3: Before the token is expired...
        mockClock.setInstant(start.plusSeconds(TOKEN_EXPIRES_IN_SECONDS / 2));
        // ...the providers should return nothing
        assertNull(passwordProvider.authorize(authorizedContext).block(), "Should not request a new token yet");
        assertNull(refreshProvider.authorize(authorizedContext).block(), "Should not have refreshed yet");

        mock.nextAccessToken = "c";
        // 4: After the token is expired...
        mockClock.setInstant(start.plusSeconds(TOKEN_EXPIRES_IN_SECONDS * 2));
        // ...the refresh provider should consider the token expired and refresh it...
        assertHasAccessToken("refresh_token-c", refreshProvider.authorize(authorizedContext).block());
        // ... but this may fail!
        assertThrows(ClientAuthorizationException.class, () -> {
            mock.fail = AuthorizationGrantType.REFRESH_TOKEN;
            refreshProvider.authorize(authorizedContext).block();
        });
        mock.fail = null;

        // But the raw password provider does not return anything, which I think is bugged
        assertHasAccessToken("password-c", passwordProvider.authorize(authorizedContext).block());
    }

    /**
     * This function is the core of the {@link #testPasswordRefresh}, {@link #testRefreshPassword},
     * {@link #testPasswordRefreshNoRefreshToken}, and {@link #testRefreshPasswordNoRefreshToken} tests.
     * These combine password and refreshToken on the builder in different orders. Both have the same
     * behavior because of the current implementation, but that also means that both fail if the
     * refresh_token is rejected (e.g. because it is a JWT with an exp claim) forcing an external retry.
     *
     * @see <a href="https://github.com/spring-projects/spring-security/issues/13915">spring-security#13915</a>
     */
    private OAuth2AuthorizedClient testTokenExpiry(ReactiveOAuth2AuthorizedClientProvider provider, MockClient mock, MockClock mockClock) {
        Instant start = mockClock.instant();

        // 2: provider can authorize the client
        OAuth2AuthorizedClient client = provider.authorize(EMPTY_CONTEXT).block();
        assertHasAccessToken("password-a", client);
        // ...and we get a refresh_token
        // assertHasRefreshToken(REFRESH_TOKEN, client);

        // Create a new context with the authorized client
        OAuth2AuthorizationContext authorizedContext = createAuthorizedContext(client);

        mock.nextAccessToken = "b";
        // 3: Before the token is expired...
        mockClock.setInstant(start.plusSeconds(TOKEN_EXPIRES_IN_SECONDS / 2));
        // ...the provider should return nothing
        assertNull(provider.authorize(authorizedContext).block(), "Should not request a new token yet");

        mock.nextAccessToken = "c";
        // 4: After the token is expired...
        mockClock.setInstant(start.plusSeconds(TOKEN_EXPIRES_IN_SECONDS * 2));
        if (client != null && client.getRefreshToken() != null) {
            // ...the refresh provider should consider the token expired and try to refresh it...
            assertHasAccessToken("refresh_token-c", provider.authorize(authorizedContext).block());
        }
        // ... but this may fail!
        mock.fail = AuthorizationGrantType.REFRESH_TOKEN;
        OAuth2AuthorizedClient newClient = provider.authorize(authorizedContext).block();
        mock.fail = null;
        return newClient;
    }

    /**
     * Test if the server returns a refresh_token, and we combine password, then refresh providers
     */
    @Test
    public void testPasswordRefresh() {
        // Initialize the clock and the mock client
        MockClock mockClock = new MockClock();
        MockClient mock = new MockClient("a", REFRESH_TOKEN);

        // 1: create a refresh-then-password provider
        OAuth2AuthorizedClient client = testTokenExpiry(createPasswordRefreshProvider(mock, mockClock), mock, mockClock);
        assertHasAccessToken("refresh_token-c", client);
    }

    /**
     * Test if the server returns a refresh_token, and we combine refresh, then password providers
     */
    @Test
    public void testRefreshPassword() {
        // Initialize the clock and the mock client
        MockClock mockClock = new MockClock();
        MockClient mock = new MockClient("a", REFRESH_TOKEN);

        // 1: create a refresh-password provider
        OAuth2AuthorizedClient client = testTokenExpiry(createRefreshPasswordProvider(mock, mockClock), mock, mockClock);
        assertHasAccessToken("refresh_token-c", client);
    }

    /**
     * Test if the server does not return a refresh_token, and we combine password, then refresh providers
     */
    @Test
    public void testPasswordRefreshNoRefreshToken() {
        // Initialize the clock and the mock client
        MockClock mockClock = new MockClock();
        MockClient mock = new MockClient("a", null);

        // 1: create a refresh-password provider
        OAuth2AuthorizedClient client = testTokenExpiry(createPasswordRefreshProvider(mock, mockClock), mock, mockClock);
        assertHasAccessToken("password-c", client);
    }

    /**
     * Test if the server does not return a refresh_token, and we combine refresh, then password providers
     */
    @Test
    public void testRefreshPasswordNoRefreshToken() {
        // Initialize the clock and the mock client
        MockClock mockClock = new MockClock();
        MockClient mock = new MockClient("a", null);

        // 1: create a refresh-password provider
        OAuth2AuthorizedClient client = testTokenExpiry(createRefreshPasswordProvider(mock, mockClock), mock, mockClock);
        assertHasAccessToken("password-c", client);
    }

    /**
     * In general, the providers added to the builder are sensitive to the order they are added.
     */
    @Test
    public void testProviderOrder() {
        // If the first provider fails, the second one should be used
        ReactiveOAuth2AuthorizedClientProvider provider1 = ReactiveOAuth2AuthorizedClientProviderBuilder.builder()
                .provider(a -> Mono.error(new IllegalStateException("a")))
                .password(new MockClient("a", null).password(new MockClock()))
                .build();
        assertThrows(IllegalStateException.class, () -> provider1.authorize(EMPTY_CONTEXT).block());

        // If the first provider succeeds, the second one should not be used
        ReactiveOAuth2AuthorizedClientProvider provider2 = ReactiveOAuth2AuthorizedClientProviderBuilder.builder()
                .password(new MockClient("a", null).password(new MockClock()))
                .provider(a -> Mono.error(new IllegalStateException("a")))
                .build();
        assertHasAccessToken("password-a", provider2.authorize(EMPTY_CONTEXT).block());
    }

    private static void assertHasAccessToken(String expected, OAuth2AuthorizedClient client) {
        assertNotNull(client, "provider did not return an authorized client");
        assertNotNull(client.getAccessToken(), "authorization server did not return an access token");
        assertEquals(expected, client.getAccessToken().getTokenValue());
    }

    private static void assertHasRefreshToken(String expected, OAuth2AuthorizedClient client) {
        assertNotNull(client, "provider did not return an authorized client");
        assertNotNull(client.getRefreshToken(), "authorization server did not return a refresh token");
        assertEquals(expected, client.getRefreshToken().getTokenValue());
    }

    private static ReactiveOAuth2AuthorizedClientProvider createPasswordRefreshProvider(MockClient mock, MockClock mockClock) {
        return ReactiveOAuth2AuthorizedClientProviderBuilder.builder()
                .password(mock.password(mockClock))
                .refreshToken(mock.refreshToken(mockClock))
                .build();
    }

    private static ReactiveOAuth2AuthorizedClientProvider createRefreshPasswordProvider(MockClient mock, MockClock mockClock) {
        return ReactiveOAuth2AuthorizedClientProviderBuilder.builder()
                .refreshToken(mock.refreshToken(mockClock))
                .password(mock.password(mockClock))
                .build();
    }

    private static ReactiveOAuth2AuthorizedClientProvider getRefreshProvider(MockClient mock, MockClock mockClock) {
        return ReactiveOAuth2AuthorizedClientProviderBuilder.builder().refreshToken(mock.refreshToken(mockClock)).build();
    }

    private static ReactiveOAuth2AuthorizedClientProvider getPasswordProvider(MockClient mock, MockClock mockClock) {
        return ReactiveOAuth2AuthorizedClientProviderBuilder.builder().password(mock.password(mockClock)).build();
    }

    /**
     * Create an authorization context with the authorized client
     */
    private OAuth2AuthorizationContext createAuthorizedContext(OAuth2AuthorizedClient client) {
        return OAuth2AuthorizationContext.withAuthorizedClient(client)
                .principal(PRINCIPAL)
                .attribute(USERNAME_ATTRIBUTE_NAME, USERNAME)
                .attribute(PASSWORD_ATTRIBUTE_NAME, PASSWORD)
                .build();
    }

    /**
     * Mock Authorization Server Responses
     */
    private static class MockClient {
        String nextAccessToken;
        String nextRefreshToken;
        AuthorizationGrantType fail = null;

        MockClient(String nextAccessToken, String nextRefreshToken) {
            this.nextAccessToken = nextAccessToken;
            this.nextRefreshToken = nextRefreshToken;
        }

        Consumer<ReactiveOAuth2AuthorizedClientProviderBuilder.PasswordGrantBuilder> password(MockClock mockClock) {
            return c -> c.accessTokenResponseClient(this::tokenResponse).clock(mockClock).clockSkew(Duration.ZERO);
        }

        Consumer<ReactiveOAuth2AuthorizedClientProviderBuilder.RefreshTokenGrantBuilder> refreshToken(MockClock mockClock) {
            return c -> c.accessTokenResponseClient(this::tokenResponse).clock(mockClock).clockSkew(Duration.ZERO);
        }

        public Mono<OAuth2AccessTokenResponse> tokenResponse(AbstractOAuth2AuthorizationGrantRequest request) {
            if (Objects.equals(request.getGrantType(), fail)) {
                OAuth2Error error = new OAuth2Error(INVALID_GRANT,  fail.getValue() + " flow failed", null);
                return Mono.error(new OAuth2AuthorizationException(error));
            } else {
                return Mono.just(OAuth2AccessTokenResponse
                        .withToken(request.getGrantType().getValue() + "-" + nextAccessToken)
                        .tokenType(BEARER)
                        .expiresIn(TOKEN_EXPIRES_IN_SECONDS)
                        .refreshToken(nextRefreshToken).build());
            }
        }
    }

    /**
     * Simple Mock Clock Implementation
     */
    private static class MockClock extends Clock {
        private final ZoneId zoneId;
        private Instant instant;

        MockClock(ZoneId zoneId, Instant instant) {
            this.zoneId = zoneId;
            this.instant = instant;
        }

        MockClock() {
            this(TimeZone.getDefault().toZoneId(), Instant.now());
        }

        @Override
        public ZoneId getZone() {
            return zoneId;
        }

        @Override
        public Clock withZone(ZoneId zone) {
            return new MockClock();
        }

        @Override
        public Instant instant() {
            return instant;
        }

        void setInstant(Instant instant) {
            this.instant = instant;
        }
    }
}