Provide configuration for authorization code generator

Closes gh-376

oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProvider.java

@@ -25,6 +25,7 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
 import org.springframework.security.authentication.AnonymousAuthenticationToken;
@@ -73,13 +74,16 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen
 	private static final String PKCE_ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc7636#section-4.4.1";
 	private static final Pattern LOOPBACK_ADDRESS_PATTERN =
 			Pattern.compile("^127(?:\\.[0-9]+){0,2}\\.[0-9]+$|^\\[(?:0*:)*?:?0*1]$");
+	private static final StringKeyGenerator DEFAULT_AUTHORIZATION_CODE_GENERATOR =
+			new Base64StringKeyGenerator(Base64.getUrlEncoder().withoutPadding(), 96);
+	private static final StringKeyGenerator DEFAULT_STATE_GENERATOR =
+			new Base64StringKeyGenerator(Base64.getUrlEncoder());
 	private static final Function<String, OAuth2AuthenticationValidator> DEFAULT_AUTHENTICATION_VALIDATOR_RESOLVER =
 			createDefaultAuthenticationValidatorResolver();
 	private final RegisteredClientRepository registeredClientRepository;
 	private final OAuth2AuthorizationService authorizationService;
 	private final OAuth2AuthorizationConsentService authorizationConsentService;
-	private final StringKeyGenerator codeGenerator = new Base64StringKeyGenerator(Base64.getUrlEncoder().withoutPadding(), 96);
-	private final StringKeyGenerator stateGenerator = new Base64StringKeyGenerator(Base64.getUrlEncoder());
+	private Supplier<String> authorizationCodeGenerator = DEFAULT_AUTHORIZATION_CODE_GENERATOR::generateKey;
 	private Function<String, OAuth2AuthenticationValidator> authenticationValidatorResolver = DEFAULT_AUTHENTICATION_VALIDATOR_RESOLVER;
 
 	/**
@@ -114,6 +118,16 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen
 		return OAuth2AuthorizationCodeRequestAuthenticationToken.class.isAssignableFrom(authentication);
 	}
 
+	/**
+	 * Sets the {@code Supplier<String>} that generates the value for the {@link OAuth2AuthorizationCode}.
+	 *
+	 * @param authorizationCodeGenerator the {@code Supplier<String>} that generates the value for the {@link OAuth2AuthorizationCode}
+	 */
+	public void setAuthorizationCodeGenerator(Supplier<String> authorizationCodeGenerator) {
+		Assert.notNull(authorizationCodeGenerator, "authorizationCodeGenerator cannot be null");
+		this.authorizationCodeGenerator = authorizationCodeGenerator;
+	}
+
 	/**
 	 * Sets the resolver that resolves an {@link OAuth2AuthenticationValidator} from the provided OAuth 2.0 Authorization Request parameter.
 	 *
@@ -199,7 +213,7 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen
 				registeredClient.getId(), principal.getName());
 
 		if (requireAuthorizationConsent(registeredClient, authorizationRequest, currentAuthorizationConsent)) {
-			String state = this.stateGenerator.generateKey();
+			String state = DEFAULT_STATE_GENERATOR.generateKey();
 			OAuth2Authorization authorization = authorizationBuilder(registeredClient, principal, authorizationRequest)
 					.attribute(OAuth2ParameterNames.STATE, state)
 					.build();
@@ -257,7 +271,7 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen
 	private OAuth2AuthorizationCode createAuthorizationCode() {
 		Instant issuedAt = Instant.now();
 		Instant expiresAt = issuedAt.plus(5, ChronoUnit.MINUTES);		// TODO Allow configuration for authorization code time-to-live
-		return new OAuth2AuthorizationCode(this.codeGenerator.generateKey(), issuedAt, expiresAt);
+		return new OAuth2AuthorizationCode(this.authorizationCodeGenerator.get(), issuedAt, expiresAt);
 	}
 
 	private Authentication authenticateAuthorizationConsent(Authentication authentication) throws AuthenticationException {

oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProviderTests.java

@@ -22,6 +22,7 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -56,6 +57,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -113,6 +115,13 @@ public class OAuth2AuthorizationCodeRequestAuthenticationProviderTests {
 		assertThat(this.authenticationProvider.supports(OAuth2AuthorizationCodeRequestAuthenticationToken.class)).isTrue();
 	}
 
+	@Test
+	public void setAuthorizationCodeGeneratorWhenNullThenThrowIllegalArgumentException() {
+		assertThatThrownBy(() -> this.authenticationProvider.setAuthorizationCodeGenerator(null))
+				.isInstanceOf(IllegalArgumentException.class)
+				.hasMessage("authorizationCodeGenerator cannot be null");
+	}
+
 	@Test
 	public void setAuthenticationValidatorResolverWhenNullThenThrowIllegalArgumentException() {
 		assertThatThrownBy(() -> this.authenticationProvider.setAuthenticationValidatorResolver(null))
@@ -487,6 +496,34 @@ public class OAuth2AuthorizationCodeRequestAuthenticationProviderTests {
 		assertAuthorizationCodeRequestWithAuthorizationCodeResult(registeredClient, authentication, authenticationResult);
 	}
 
+	@Test
+	public void authenticateWhenCustomAuthorizationCodeGeneratorThenUsed() {
+		RegisteredClient registeredClient = TestRegisteredClients.registeredClient().build();
+		when(this.registeredClientRepository.findByClientId(eq(registeredClient.getClientId())))
+				.thenReturn(registeredClient);
+
+		@SuppressWarnings("unchecked")
+		Supplier<String> authorizationCodeGenerator = spy(new Supplier<String>() {
+			@Override
+			public String get() {
+				return "custom-code";
+			}
+		});
+		this.authenticationProvider.setAuthorizationCodeGenerator(authorizationCodeGenerator);
+
+		OAuth2AuthorizationCodeRequestAuthenticationToken authentication =
+				authorizationCodeRequestAuthentication(registeredClient, this.principal)
+						.build();
+
+		OAuth2AuthorizationCodeRequestAuthenticationToken authenticationResult =
+				(OAuth2AuthorizationCodeRequestAuthenticationToken) this.authenticationProvider.authenticate(authentication);
+
+		assertAuthorizationCodeRequestWithAuthorizationCodeResult(registeredClient, authentication, authenticationResult);
+
+		verify(authorizationCodeGenerator).get();
+		assertThat(authenticationResult.getAuthorizationCode().getTokenValue()).isEqualTo(authorizationCodeGenerator.get());
+	}
+
 	@Test
 	public void authenticateWhenCustomAuthenticationValidatorResolverThenUsed() {
 		RegisteredClient registeredClient = TestRegisteredClients.registeredClient().build();