Allow setting Clock in OAuth2TokenGenerator implementations

Closes gh-18017

oauth2/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/token/JwtGenerator.java

@@ -16,6 +16,7 @@
 
 package org.springframework.security.oauth2.server.authorization.token;
 
+import java.time.Clock;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.Collections;
@@ -65,6 +66,8 @@ public final class JwtGenerator implements OAuth2TokenGenerator<Jwt> {
 
 	private OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer;
 
+	private Clock clock = Clock.systemUTC();
+
 	/**
 	 * Constructs a {@code JwtGenerator} using the provided parameters.
 	 * @param jwtEncoder the jwt encoder
@@ -95,7 +98,7 @@ public final class JwtGenerator implements OAuth2TokenGenerator<Jwt> {
 		}
 		RegisteredClient registeredClient = context.getRegisteredClient();
 
-		Instant issuedAt = Instant.now();
+		Instant issuedAt = this.clock.instant();
 		Instant expiresAt;
 		JwsAlgorithm jwsAlgorithm = SignatureAlgorithm.RS256;
 		if (OidcParameterNames.ID_TOKEN.equals(context.getTokenType().getValue())) {
@@ -208,4 +211,15 @@ public final class JwtGenerator implements OAuth2TokenGenerator<Jwt> {
 		this.jwtCustomizer = jwtCustomizer;
 	}
 
+	/**
+	 * Sets the {@link Clock} used when obtaining the current instant via
+	 * {@link Clock#instant()}.
+	 * @param clock the {@link Clock} used when obtaining the current instant via
+	 * {@link Clock#instant()}
+	 */
+	public void setClock(Clock clock) {
+		Assert.notNull(clock, "clock cannot be null");
+		this.clock = clock;
+	}
+
 }

oauth2/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/token/OAuth2AccessTokenGenerator.java

@@ -16,6 +16,7 @@
 
 package org.springframework.security.oauth2.server.authorization.token;
 
+import java.time.Clock;
 import java.time.Instant;
 import java.util.Base64;
 import java.util.Collections;
@@ -56,6 +57,8 @@ public final class OAuth2AccessTokenGenerator implements OAuth2TokenGenerator<OA
 
 	private OAuth2TokenCustomizer<OAuth2TokenClaimsContext> accessTokenCustomizer;
 
+	private Clock clock = Clock.systemUTC();
+
 	@Nullable
 	@Override
 	public OAuth2AccessToken generate(OAuth2TokenContext context) {
@@ -72,7 +75,7 @@ public final class OAuth2AccessTokenGenerator implements OAuth2TokenGenerator<OA
 		}
 		RegisteredClient registeredClient = context.getRegisteredClient();
 
-		Instant issuedAt = Instant.now();
+		Instant issuedAt = this.clock.instant();
 		Instant expiresAt = issuedAt.plus(registeredClient.getTokenSettings().getAccessTokenTimeToLive());
 
 		// @formatter:off
@@ -140,6 +143,17 @@ public final class OAuth2AccessTokenGenerator implements OAuth2TokenGenerator<OA
 		this.accessTokenCustomizer = accessTokenCustomizer;
 	}
 
+	/**
+	 * Sets the {@link Clock} used when obtaining the current instant via
+	 * {@link Clock#instant()}.
+	 * @param clock the {@link Clock} used when obtaining the current instant via
+	 * {@link Clock#instant()}
+	 */
+	public void setClock(Clock clock) {
+		Assert.notNull(clock, "clock cannot be null");
+		this.clock = clock;
+	}
+
 	private static final class OAuth2AccessTokenClaims extends OAuth2AccessToken implements ClaimAccessor {
 
 		private final Map<String, Object> claims;

oauth2/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/token/OAuth2RefreshTokenGenerator.java

@@ -16,6 +16,7 @@
 
 package org.springframework.security.oauth2.server.authorization.token;
 
+import java.time.Clock;
 import java.time.Instant;
 import java.util.Base64;
 
@@ -27,6 +28,7 @@ import org.springframework.security.oauth2.core.ClientAuthenticationMethod;
 import org.springframework.security.oauth2.core.OAuth2RefreshToken;
 import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
 import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
+import org.springframework.util.Assert;
 
 /**
  * An {@link OAuth2TokenGenerator} that generates an {@link OAuth2RefreshToken}.
@@ -41,6 +43,8 @@ public final class OAuth2RefreshTokenGenerator implements OAuth2TokenGenerator<O
 	private final StringKeyGenerator refreshTokenGenerator = new Base64StringKeyGenerator(
 			Base64.getUrlEncoder().withoutPadding(), 96);
 
+	private Clock clock = Clock.systemUTC();
+
 	@Nullable
 	@Override
 	public OAuth2RefreshToken generate(OAuth2TokenContext context) {
@@ -52,11 +56,22 @@ public final class OAuth2RefreshTokenGenerator implements OAuth2TokenGenerator<O
 			return null;
 		}
 
-		Instant issuedAt = Instant.now();
+		Instant issuedAt = this.clock.instant();
 		Instant expiresAt = issuedAt.plus(context.getRegisteredClient().getTokenSettings().getRefreshTokenTimeToLive());
 		return new OAuth2RefreshToken(this.refreshTokenGenerator.generateKey(), issuedAt, expiresAt);
 	}
 
+	/**
+	 * Sets the {@link Clock} used when obtaining the current instant via
+	 * {@link Clock#instant()}.
+	 * @param clock the {@link Clock} used when obtaining the current instant via
+	 * {@link Clock#instant()}
+	 */
+	public void setClock(Clock clock) {
+		Assert.notNull(clock, "clock cannot be null");
+		this.clock = clock;
+	}
+
 	private static boolean isPublicClientForAuthorizationCodeGrant(OAuth2TokenContext context) {
 		// @formatter:off
 		if (AuthorizationGrantType.AUTHORIZATION_CODE.equals(context.getAuthorizationGrantType()) &&

oauth2/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/token/JwtGeneratorTests.java

@@ -17,6 +17,8 @@
 package org.springframework.security.oauth2.server.authorization.token;
 
 import java.security.Principal;
+import java.time.Clock;
+import java.time.Duration;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.Date;
@@ -104,6 +106,12 @@ public class JwtGeneratorTests {
 			.withMessage("jwtCustomizer cannot be null");
 	}
 
+	@Test
+	public void setClockWhenNullThenThrowIllegalArgumentException() {
+		assertThatExceptionOfType(IllegalArgumentException.class).isThrownBy(() -> this.jwtGenerator.setClock(null))
+			.withMessage("clock cannot be null");
+	}
+
 	@Test
 	public void generateWhenUnsupportedTokenTypeThenReturnNull() {
 		// @formatter:off
@@ -158,7 +166,10 @@ public class JwtGeneratorTests {
 				.build();
 		// @formatter:on
 
-		assertGeneratedTokenType(tokenContext);
+		Clock clock = Clock.offset(Clock.systemUTC(), Duration.ofMinutes(5));
+		this.jwtGenerator.setClock(clock);
+
+		assertGeneratedTokenType(tokenContext, clock);
 	}
 
 	@Test
@@ -282,6 +293,10 @@ public class JwtGeneratorTests {
 	}
 
 	private void assertGeneratedTokenType(OAuth2TokenContext tokenContext) {
+		assertGeneratedTokenType(tokenContext, Clock.systemUTC());
+	}
+
+	private void assertGeneratedTokenType(OAuth2TokenContext tokenContext, Clock clock) {
 		this.jwtGenerator.generate(tokenContext);
 
 		ArgumentCaptor<JwtEncodingContext> jwtEncodingContextCaptor = ArgumentCaptor.forClass(JwtEncodingContext.class);
@@ -318,7 +333,7 @@ public class JwtGeneratorTests {
 		assertThat(jwtClaimsSet.getSubject()).isEqualTo(tokenContext.getAuthorization().getPrincipalName());
 		assertThat(jwtClaimsSet.getAudience()).containsExactly(tokenContext.getRegisteredClient().getClientId());
 
-		Instant issuedAt = Instant.now();
+		Instant issuedAt = clock.instant();
 		Instant expiresAt;
 		if (tokenContext.getTokenType().equals(OAuth2TokenType.ACCESS_TOKEN)) {
 			expiresAt = issuedAt.plus(tokenContext.getRegisteredClient().getTokenSettings().getAccessTokenTimeToLive());

oauth2/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/token/OAuth2AccessTokenGeneratorTests.java

@@ -17,6 +17,8 @@
 package org.springframework.security.oauth2.server.authorization.token;
 
 import java.security.Principal;
+import java.time.Clock;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Collections;
 import java.util.Set;
@@ -81,6 +83,13 @@ public class OAuth2AccessTokenGeneratorTests {
 			.withMessage("accessTokenCustomizer cannot be null");
 	}
 
+	@Test
+	public void setClockWhenNullThenThrowIllegalArgumentException() {
+		assertThatExceptionOfType(IllegalArgumentException.class)
+			.isThrownBy(() -> this.accessTokenGenerator.setClock(null))
+			.withMessage("clock cannot be null");
+	}
+
 	@Test
 	public void generateWhenUnsupportedTokenTypeThenReturnNull() {
 		// @formatter:off
@@ -150,10 +159,13 @@ public class OAuth2AccessTokenGeneratorTests {
 				.build();
 		// @formatter:on
 
+		Clock clock = Clock.offset(Clock.systemUTC(), Duration.ofMinutes(5));
+		this.accessTokenGenerator.setClock(clock);
+
 		OAuth2AccessToken accessToken = this.accessTokenGenerator.generate(tokenContext);
 		assertThat(accessToken).isNotNull();
 
-		Instant issuedAt = Instant.now();
+		Instant issuedAt = clock.instant();
 		Instant expiresAt = issuedAt
 			.plus(tokenContext.getRegisteredClient().getTokenSettings().getAccessTokenTimeToLive());
 		assertThat(accessToken.getIssuedAt()).isBetween(issuedAt.minusSeconds(1), issuedAt.plusSeconds(1));

oauth2/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/token/OAuth2RefreshTokenGeneratorTests.java

@@ -16,6 +16,8 @@
 
 package org.springframework.security.oauth2.server.authorization.token;
 
+import java.time.Clock;
+import java.time.Duration;
 import java.time.Instant;
 
 import org.junit.jupiter.api.Test;
@@ -26,6 +28,7 @@ import org.springframework.security.oauth2.server.authorization.client.Registere
 import org.springframework.security.oauth2.server.authorization.client.TestRegisteredClients;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 
 /**
  * Tests for {@link OAuth2RefreshTokenGenerator}.
@@ -36,6 +39,12 @@ public class OAuth2RefreshTokenGeneratorTests {
 
 	private final OAuth2RefreshTokenGenerator tokenGenerator = new OAuth2RefreshTokenGenerator();
 
+	@Test
+	public void setClockWhenNullThenThrowIllegalArgumentException() {
+		assertThatExceptionOfType(IllegalArgumentException.class).isThrownBy(() -> this.tokenGenerator.setClock(null))
+			.withMessage("clock cannot be null");
+	}
+
 	@Test
 	public void generateWhenUnsupportedTokenTypeThenReturnNull() {
 		// @formatter:off
@@ -58,10 +67,13 @@ public class OAuth2RefreshTokenGeneratorTests {
 				.build();
 		// @formatter:on
 
+		Clock clock = Clock.offset(Clock.systemUTC(), Duration.ofMinutes(5));
+		this.tokenGenerator.setClock(clock);
+
 		OAuth2RefreshToken refreshToken = this.tokenGenerator.generate(tokenContext);
 		assertThat(refreshToken).isNotNull();
 
-		Instant issuedAt = Instant.now();
+		Instant issuedAt = clock.instant();
 		Instant expiresAt = issuedAt
 			.plus(tokenContext.getRegisteredClient().getTokenSettings().getRefreshTokenTimeToLive());
 		assertThat(refreshToken.getIssuedAt()).isBetween(issuedAt.minusSeconds(1), issuedAt.plusSeconds(1));