Skip to content

Commit 1dd79c3

Browse files
rwinchrwinch
committed
Add JdbcOneTimeTokenService
Closes gh-15735
2 parents f5991ae + f002fed commit 1dd79c3

File tree

7 files changed

+561
-1
lines changed

7 files changed

+561
-1
lines changed

‎core/src/main/java/org/springframework/security/aot/hint/OneTimeTokenRuntimeHints.java

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright 2002-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.security.aot.hint;
18+
19+
import org.springframework.aot.hint.RuntimeHints;
20+
import org.springframework.aot.hint.RuntimeHintsRegistrar;
21+
import org.springframework.jdbc.core.JdbcOperations;
22+
import org.springframework.security.authentication.ott.OneTimeToken;
23+
import org.springframework.security.authentication.ott.OneTimeTokenService;
24+
25+
/**
26+
*
27+
* A JDBC implementation of an {@link OneTimeTokenService} that uses a
28+
* {@link JdbcOperations} for {@link OneTimeToken} persistence.
29+
*
30+
* @author Max Batischev
31+
* @since 6.4
32+
*/
33+
class OneTimeTokenRuntimeHints implements RuntimeHintsRegistrar {
34+
35+
@Override
36+
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
37+
hints.resources().registerPattern("org/springframework/security/core/ott/jdbc/one-time-tokens-schema.sql");
38+
}
39+
40+
}

‎core/src/main/java/org/springframework/security/authentication/ott/JdbcOneTimeTokenService.java

Lines changed: 262 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
/*
2+
* Copyright 2002-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.security.authentication.ott;
18+
19+
import java.sql.ResultSet;
20+
import java.sql.SQLException;
21+
import java.sql.Timestamp;
22+
import java.sql.Types;
23+
import java.time.Clock;
24+
import java.time.Duration;
25+
import java.time.Instant;
26+
import java.util.ArrayList;
27+
import java.util.List;
28+
import java.util.UUID;
29+
import java.util.function.Function;
30+
31+
import org.apache.commons.logging.Log;
32+
import org.apache.commons.logging.LogFactory;
33+
34+
import org.springframework.beans.factory.DisposableBean;
35+
import org.springframework.beans.factory.InitializingBean;
36+
import org.springframework.jdbc.core.ArgumentPreparedStatementSetter;
37+
import org.springframework.jdbc.core.JdbcOperations;
38+
import org.springframework.jdbc.core.PreparedStatementSetter;
39+
import org.springframework.jdbc.core.RowMapper;
40+
import org.springframework.jdbc.core.SqlParameterValue;
41+
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
42+
import org.springframework.scheduling.support.CronTrigger;
43+
import org.springframework.util.Assert;
44+
import org.springframework.util.CollectionUtils;
45+
46+
/**
47+
*
48+
* A JDBC implementation of an {@link OneTimeTokenService} that uses a
49+
* {@link JdbcOperations} for {@link OneTimeToken} persistence.
50+
*
51+
* <p>
52+
* <b>NOTE:</b> This {@code JdbcOneTimeTokenService} depends on the table definition
53+
* described in
54+
* "classpath:org/springframework/security/core/ott/jdbc/one-time-tokens-schema.sql" and
55+
* therefore MUST be defined in the database schema.
56+
*
57+
* @author Max Batischev
58+
* @since 6.4
59+
*/
60+
public final class JdbcOneTimeTokenService implements OneTimeTokenService, DisposableBean, InitializingBean {
61+
62+
private final Log logger = LogFactory.getLog(getClass());
63+
64+
private final JdbcOperations jdbcOperations;
65+
66+
private Function<OneTimeToken, List<SqlParameterValue>> oneTimeTokenParametersMapper = new OneTimeTokenParametersMapper();
67+
68+
private RowMapper<OneTimeToken> oneTimeTokenRowMapper = new OneTimeTokenRowMapper();
69+
70+
private Clock clock = Clock.systemUTC();
71+
72+
private ThreadPoolTaskScheduler taskScheduler;
73+
74+
private static final String DEFAULT_CLEANUP_CRON = "@hourly";
75+
76+
private static final String TABLE_NAME = "one_time_tokens";
77+
78+
// @formatter:off
79+
private static final String COLUMN_NAMES = "token_value, "
80+
+ "username, "
81+
+ "expires_at";
82+
// @formatter:on
83+
84+
// @formatter:off
85+
private static final String SAVE_AUTHORIZED_CLIENT_SQL = "INSERT INTO " + TABLE_NAME
86+
+ " (" + COLUMN_NAMES + ") VALUES (?, ?, ?)";
87+
// @formatter:on
88+
89+
private static final String FILTER = "token_value = ?";
90+
91+
private static final String DELETE_ONE_TIME_TOKEN_SQL = "DELETE FROM " + TABLE_NAME + " WHERE " + FILTER;
92+
93+
// @formatter:off
94+
private static final String SELECT_ONE_TIME_TOKEN_SQL = "SELECT " + COLUMN_NAMES
95+
+ " FROM " + TABLE_NAME
96+
+ " WHERE " + FILTER;
97+
// @formatter:on
98+
99+
// @formatter:off
100+
private static final String DELETE_SESSIONS_BY_EXPIRY_TIME_QUERY = "DELETE FROM "
101+
+ TABLE_NAME
102+
+ " WHERE expires_at < ?";
103+
// @formatter:on
104+
105+
/**
106+
* Constructs a {@code JdbcOneTimeTokenService} using the provide parameters.
107+
* @param jdbcOperations the JDBC operations
108+
*/
109+
public JdbcOneTimeTokenService(JdbcOperations jdbcOperations) {
110+
Assert.notNull(jdbcOperations, "jdbcOperations cannot be null");
111+
this.jdbcOperations = jdbcOperations;
112+
this.taskScheduler = createTaskScheduler(DEFAULT_CLEANUP_CRON);
113+
}
114+
115+
/**
116+
* Sets the chron expression used for cleaning up expired sessions. The default is to
117+
* run hourly.
118+
*
119+
* For more advanced use cases the cleanupCron may be set to null which will disable
120+
* the built-in cleanup. Users can then invoke {@link #cleanupExpiredTokens()} using
121+
* custom logic.
122+
* @param cleanupCron the chron expression passed to {@link CronTrigger} used for
123+
* determining how frequent to perform cleanup. The default is "@hourly".
124+
* @see CronTrigger
125+
* @see #cleanupExpiredTokens()
126+
*/
127+
public void setCleanupCron(String cleanupCron) {
128+
this.taskScheduler = createTaskScheduler(cleanupCron);
129+
}
130+
131+
@Override
132+
public OneTimeToken generate(GenerateOneTimeTokenRequest request) {
133+
Assert.notNull(request, "generateOneTimeTokenRequest cannot be null");
134+
String token = UUID.randomUUID().toString();
135+
Instant fiveMinutesFromNow = this.clock.instant().plus(Duration.ofMinutes(5));
136+
OneTimeToken oneTimeToken = new DefaultOneTimeToken(token, request.getUsername(), fiveMinutesFromNow);
137+
insertOneTimeToken(oneTimeToken);
138+
return oneTimeToken;
139+
}
140+
141+
private void insertOneTimeToken(OneTimeToken oneTimeToken) {
142+
List<SqlParameterValue> parameters = this.oneTimeTokenParametersMapper.apply(oneTimeToken);
143+
PreparedStatementSetter pss = new ArgumentPreparedStatementSetter(parameters.toArray());
144+
this.jdbcOperations.update(SAVE_AUTHORIZED_CLIENT_SQL, pss);
145+
}
146+
147+
@Override
148+
public OneTimeToken consume(OneTimeTokenAuthenticationToken authenticationToken) {
149+
Assert.notNull(authenticationToken, "authenticationToken cannot be null");
150+
151+
List<OneTimeToken> tokens = selectOneTimeToken(authenticationToken);
152+
if (CollectionUtils.isEmpty(tokens)) {
153+
return null;
154+
}
155+
OneTimeToken token = tokens.get(0);
156+
deleteOneTimeToken(token);
157+
if (isExpired(token)) {
158+
return null;
159+
}
160+
return token;
161+
}
162+
163+
private boolean isExpired(OneTimeToken ott) {
164+
return this.clock.instant().isAfter(ott.getExpiresAt());
165+
}
166+
167+
private List<OneTimeToken> selectOneTimeToken(OneTimeTokenAuthenticationToken authenticationToken) {
168+
List<SqlParameterValue> parameters = List
169+
.of(new SqlParameterValue(Types.VARCHAR, authenticationToken.getTokenValue()));
170+
PreparedStatementSetter pss = new ArgumentPreparedStatementSetter(parameters.toArray());
171+
return this.jdbcOperations.query(SELECT_ONE_TIME_TOKEN_SQL, pss, this.oneTimeTokenRowMapper);
172+
}
173+
174+
private void deleteOneTimeToken(OneTimeToken oneTimeToken) {
175+
List<SqlParameterValue> parameters = List
176+
.of(new SqlParameterValue(Types.VARCHAR, oneTimeToken.getTokenValue()));
177+
PreparedStatementSetter pss = new ArgumentPreparedStatementSetter(parameters.toArray());
178+
this.jdbcOperations.update(DELETE_ONE_TIME_TOKEN_SQL, pss);
179+
}
180+
181+
private ThreadPoolTaskScheduler createTaskScheduler(String cleanupCron) {
182+
if (cleanupCron == null) {
183+
return null;
184+
}
185+
ThreadPoolTaskScheduler taskScheduler = new ThreadPoolTaskScheduler();
186+
taskScheduler.setThreadNamePrefix("spring-one-time-tokens-");
187+
taskScheduler.initialize();
188+
taskScheduler.schedule(this::cleanupExpiredTokens, new CronTrigger(cleanupCron));
189+
return taskScheduler;
190+
}
191+
192+
public void cleanupExpiredTokens() {
193+
List<SqlParameterValue> parameters = List.of(new SqlParameterValue(Types.TIMESTAMP, Instant.now()));
194+
PreparedStatementSetter pss = new ArgumentPreparedStatementSetter(parameters.toArray());
195+
int deletedCount = this.jdbcOperations.update(DELETE_SESSIONS_BY_EXPIRY_TIME_QUERY, pss);
196+
if (this.logger.isDebugEnabled()) {
197+
this.logger.debug("Cleaned up " + deletedCount + " expired tokens");
198+
}
199+
}
200+
201+
@Override
202+
public void afterPropertiesSet() throws Exception {
203+
this.taskScheduler.afterPropertiesSet();
204+
}
205+
206+
@Override
207+
public void destroy() throws Exception {
208+
if (this.taskScheduler != null) {
209+
this.taskScheduler.shutdown();
210+
}
211+
}
212+
213+
/**
214+
* Sets the {@link Clock} used when generating one-time token and checking token
215+
* expiry.
216+
* @param clock the clock
217+
*/
218+
public void setClock(Clock clock) {
219+
Assert.notNull(clock, "clock cannot be null");
220+
this.clock = clock;
221+
}
222+
223+
/**
224+
* The default {@code Function} that maps {@link OneTimeToken} to a {@code List} of
225+
* {@link SqlParameterValue}.
226+
*
227+
* @author Max Batischev
228+
* @since 6.4
229+
*/
230+
private static class OneTimeTokenParametersMapper implements Function<OneTimeToken, List<SqlParameterValue>> {
231+
232+
@Override
233+
public List<SqlParameterValue> apply(OneTimeToken oneTimeToken) {
234+
List<SqlParameterValue> parameters = new ArrayList<>();
235+
parameters.add(new SqlParameterValue(Types.VARCHAR, oneTimeToken.getTokenValue()));
236+
parameters.add(new SqlParameterValue(Types.VARCHAR, oneTimeToken.getUsername()));
237+
parameters.add(new SqlParameterValue(Types.TIMESTAMP, Timestamp.from(oneTimeToken.getExpiresAt())));
238+
return parameters;
239+
}
240+
241+
}
242+
243+
/**
244+
* The default {@link RowMapper} that maps the current row in
245+
* {@code java.sql.ResultSet} to {@link OneTimeToken}.
246+
*
247+
* @author Max Batischev
248+
* @since 6.4
249+
*/
250+
private static class OneTimeTokenRowMapper implements RowMapper<OneTimeToken> {
251+
252+
@Override
253+
public OneTimeToken mapRow(ResultSet rs, int rowNum) throws SQLException {
254+
String tokenValue = rs.getString("token_value");
255+
String userName = rs.getString("username");
256+
Instant expiresAt = rs.getTimestamp("expires_at").toInstant();
257+
return new DefaultOneTimeToken(tokenValue, userName, expiresAt);
258+
}
259+
260+
}
261+
262+
}

‎core/src/main/resources/META-INF/spring/aot.factories

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
org.springframework.aot.hint.RuntimeHintsRegistrar=\
2-
org.springframework.security.aot.hint.CoreSecurityRuntimeHints
2+
org.springframework.security.aot.hint.CoreSecurityRuntimeHints,\
3+
org.springframework.security.aot.hint.OneTimeTokenRuntimeHints
4+
35
org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor=\
46
org.springframework.security.aot.hint.SecurityHintsAotProcessor

‎core/src/main/resources/org/springframework/security/core/ott/jdbc/one-time-tokens-schema.sql

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
create table one_time_tokens(
2+
token_value varchar(36) not null primary key,
3+
username varchar_ignorecase(50) not null,
4+
expires_at timestamp not null
5+
);

‎core/src/test/java/org/springframework/security/aot/hint/OneTimeTokenRuntimeHintsTests.java

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright 2002-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.security.aot.hint;
18+
19+
import java.util.stream.Stream;
20+
21+
import org.junit.jupiter.api.BeforeEach;
22+
import org.junit.jupiter.params.ParameterizedTest;
23+
import org.junit.jupiter.params.provider.MethodSource;
24+
25+
import org.springframework.aot.hint.RuntimeHints;
26+
import org.springframework.aot.hint.RuntimeHintsRegistrar;
27+
import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
28+
import org.springframework.core.io.support.SpringFactoriesLoader;
29+
import org.springframework.util.ClassUtils;
30+
31+
import static org.assertj.core.api.Assertions.assertThat;
32+
33+
/**
34+
* Tests for {@link OneTimeTokenRuntimeHints}
35+
*
36+
* @author Max Batischev
37+
*/
38+
class OneTimeTokenRuntimeHintsTests {
39+
40+
private final RuntimeHints hints = new RuntimeHints();
41+
42+
@BeforeEach
43+
void setup() {
44+
SpringFactoriesLoader.forResourceLocation("META-INF/spring/aot.factories")
45+
.load(RuntimeHintsRegistrar.class)
46+
.forEach((registrar) -> registrar.registerHints(this.hints, ClassUtils.getDefaultClassLoader()));
47+
}
48+
49+
@ParameterizedTest
50+
@MethodSource("getOneTimeTokensSqlFiles")
51+
void oneTimeTokensSqlFilesHasHints(String schemaFile) {
52+
assertThat(RuntimeHintsPredicates.resource().forResource(schemaFile)).accepts(this.hints);
53+
}
54+
55+
private static Stream<String> getOneTimeTokensSqlFiles() {
56+
return Stream.of("org/springframework/security/core/ott/jdbc/one-time-tokens-schema.sql");
57+
}
58+
59+
}

0 commit comments

Comments
 (0)