Add support OneTimeToken customization

Closes spring-projectsgh-16939

Signed-off-by: Max Batischev <[email protected]>

Author: franticticktick

Diff for: core/src/main/java/org/springframework/security/authentication/ott/GenerateOneTimeTokenRequest.java

@@ -17,13 +17,15 @@
 package org.springframework.security.authentication.ott;
 
 import java.time.Duration;
+import java.util.UUID;
 
 import org.springframework.util.Assert;
 
 /**
  * Class to store information related to an One-Time Token authentication request
  *
  * @author Marcus da Coregio
+ * @author Max Batiscev
  * @since 6.4
  */
 public class GenerateOneTimeTokenRequest {
@@ -34,6 +36,8 @@ public class GenerateOneTimeTokenRequest {
 
 	private final Duration expiresIn;
 
+	private final String tokenValue;
+
 	public GenerateOneTimeTokenRequest(String username) {
 		this(username, DEFAULT_EXPIRES_IN);
 	}
@@ -43,6 +47,16 @@ public GenerateOneTimeTokenRequest(String username, Duration expiresIn) {
 		Assert.notNull(expiresIn, "expiresIn cannot be null");
 		this.username = username;
 		this.expiresIn = expiresIn;
+		this.tokenValue = UUID.randomUUID().toString();
+	}
+
+	public GenerateOneTimeTokenRequest(String username, Duration expiresIn, String tokenValue) {
+		Assert.hasText(username, "username cannot be empty");
+		Assert.hasText(tokenValue, "tokenValue cannot be empty");
+		Assert.notNull(expiresIn, "expiresIn cannot be null");
+		this.username = username;
+		this.expiresIn = expiresIn;
+		this.tokenValue = tokenValue;
 	}
 
 	public String getUsername() {
@@ -53,4 +67,8 @@ public Duration getExpiresIn() {
 		return this.expiresIn;
 	}
 
+	public String getTokenValue() {
+		return this.tokenValue;
+	}
+
 }

Diff for: core/src/main/java/org/springframework/security/authentication/ott/InMemoryOneTimeTokenService.java

@@ -32,6 +32,7 @@
  * there is more or equal than 100 tokens stored in the map.
  *
  * @author Marcus da Coregio
+ * @author Max Batischev
  * @since 6.4
  */
 public final class InMemoryOneTimeTokenService implements OneTimeTokenService {
@@ -43,10 +44,9 @@ public final class InMemoryOneTimeTokenService implements OneTimeTokenService {
 	@Override
 	@NonNull
 	public OneTimeToken generate(GenerateOneTimeTokenRequest request) {
-		String token = UUID.randomUUID().toString();
 		Instant expiresAt = this.clock.instant().plus(request.getExpiresIn());
-		OneTimeToken ott = new DefaultOneTimeToken(token, request.getUsername(), expiresAt);
-		this.oneTimeTokenByToken.put(token, ott);
+		OneTimeToken ott = new DefaultOneTimeToken(request.getTokenValue(), request.getUsername(), expiresAt);
+		this.oneTimeTokenByToken.put(request.getTokenValue(), ott);
 		cleanExpiredTokensIfNeeded();
 		return ott;
 	}

Diff for: core/src/main/java/org/springframework/security/authentication/ott/JdbcOneTimeTokenService.java

@@ -24,7 +24,6 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.UUID;
 import java.util.function.Function;
 
 import org.apache.commons.logging.Log;
@@ -130,9 +129,8 @@ public void setCleanupCron(String cleanupCron) {
 	@Override
 	public OneTimeToken generate(GenerateOneTimeTokenRequest request) {
 		Assert.notNull(request, "generateOneTimeTokenRequest cannot be null");
-		String token = UUID.randomUUID().toString();
 		Instant expiresAt = this.clock.instant().plus(request.getExpiresIn());
-		OneTimeToken oneTimeToken = new DefaultOneTimeToken(token, request.getUsername(), expiresAt);
+		OneTimeToken oneTimeToken = new DefaultOneTimeToken(request.getTokenValue(), request.getUsername(), expiresAt);
 		insertOneTimeToken(oneTimeToken);
 		return oneTimeToken;
 	}

Diff for: web/src/main/java/org/springframework/security/web/authentication/ott/DefaultGenerateOneTimeTokenRequestResolver.java

@@ -17,6 +17,8 @@
 package org.springframework.security.web.authentication.ott;
 
 import java.time.Duration;
+import java.util.UUID;
+import java.util.function.Supplier;
 
 import jakarta.servlet.http.HttpServletRequest;
 
@@ -37,13 +39,15 @@ public final class DefaultGenerateOneTimeTokenRequestResolver implements Generat
 
 	private Duration expiresIn = DEFAULT_EXPIRES_IN;
 
+	private Supplier<String> tokenValueFactory = () -> UUID.randomUUID().toString();
+
 	@Override
 	public GenerateOneTimeTokenRequest resolve(HttpServletRequest request) {
 		String username = request.getParameter("username");
 		if (!StringUtils.hasText(username)) {
 			return null;
 		}
-		return new GenerateOneTimeTokenRequest(username, this.expiresIn);
+		return new GenerateOneTimeTokenRequest(username, this.expiresIn, this.tokenValueFactory.get());
 	}
 
 	/**
@@ -55,4 +59,14 @@ public void setExpiresIn(Duration expiresIn) {
 		this.expiresIn = expiresIn;
 	}
 
+	/**
+	 * Sets factory for token value generation
+	 * @param tokenValueFactory factory for token value generation
+	 * @since 6.5
+	 */
+	public void setTokenValueFactory(Supplier<String> tokenValueFactory) {
+		Assert.notNull(tokenValueFactory, "tokenValueFactory cannot be null");
+		this.tokenValueFactory = tokenValueFactory;
+	}
+
 }

Diff for: web/src/test/java/org/springframework/security/web/authentication/ott/DefaultGenerateOneTimeTokenRequestResolverTests.java

@@ -64,4 +64,15 @@ void resolveWhenExpiresInSetThenResolvesGenerateRequest() {
 		assertThat(generateRequest.getExpiresIn()).isEqualTo(Duration.ofSeconds(600));
 	}
 
+	@Test
+	void resolveWhenTokenValueFactorySetThenResolvesGenerateRequest() {
+		MockHttpServletRequest request = new MockHttpServletRequest();
+		request.setParameter("username", "test");
+		this.requestResolver.setTokenValueFactory(() -> "tokenValue");
+
+		GenerateOneTimeTokenRequest generateRequest = this.requestResolver.resolve(request);
+
+		assertThat(generateRequest.getTokenValue()).isEqualTo("tokenValue");
+	}
+
 }