Skip to content

Commit

Permalink
Add mock for TimeService
Browse files Browse the repository at this point in the history
  • Loading branch information
@strehle
strehle committed May 5, 2023
1 parent 3897112 commit 9569b5b
Show file tree
Hide file tree
Showing 10 changed files with 58 additions and 33 deletions.
11 changes: 7 additions & 4 deletions server/src/main/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenStore.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.cloudfoundry.identity.uaa.authentication.UaaAuthenticationDetails;
import org.cloudfoundry.identity.uaa.authentication.UaaPrincipal;
import org.cloudfoundry.identity.uaa.util.JsonUtils;
import org.cloudfoundry.identity.uaa.util.TimeService;
import org.cloudfoundry.identity.uaa.util.UaaStringUtils;
import org.cloudfoundry.identity.uaa.zone.IdentityZoneHolder;
import org.slf4j.Logger;
Expand Down Expand Up @@ -79,18 +80,20 @@ public class UaaTokenStore implements AuthorizationCodeServices {
private static final String SQL_CLEAN_STATEMENT = "delete from oauth_code where created < ? and expiresat = 0";

private final DataSource dataSource;
private final TimeService timeService;
private final Duration expirationTime;
private final RandomValueStringGenerator generator = new RandomValueStringGenerator(32);
private final RowMapper rowMapper = new TokenCodeRowMapper();

private AtomicReference<Instant> lastClean = new AtomicReference<>(Instant.EPOCH);

public UaaTokenStore(DataSource dataSource) {
this(dataSource, DEFAULT_EXPIRATION_TIME);
public UaaTokenStore(DataSource dataSource, TimeService timeService) {
this(dataSource, timeService, DEFAULT_EXPIRATION_TIME);
}

public UaaTokenStore(DataSource dataSource, Duration expirationTime) {
public UaaTokenStore(DataSource dataSource, TimeService timeService, Duration expirationTime) {
this.dataSource = dataSource;
this.timeService = timeService;
this.expirationTime = expirationTime;
}

Expand Down Expand Up @@ -213,7 +216,7 @@ else if (map.get(USER_AUTHENTICATION_UAA_PRINCIPAL)!=null) {
protected void performExpirationClean() {
Instant last = lastClean.get();
//check if we should expire again
Instant now = Instant.now();
Instant now = timeService.getCurrentInstant();
if (enoughTimeHasPassedSinceLastExpirationClean(last, now)) {
//avoid concurrent deletes from the same UAA - performance improvement
if (lastClean.compareAndSet(last, now)) {
Expand Down
7 changes: 5 additions & 2 deletions server/src/main/java/org/cloudfoundry/identity/uaa/util/TimeService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,15 @@
package org.cloudfoundry.identity.uaa.util;


import java.time.Instant;
import java.util.Date;

public interface TimeService {
default long getCurrentTimeMillis() {
return System.currentTimeMillis();
return getCurrentInstant().toEpochMilli();
}

default Date getCurrentDate() { return new Date(getCurrentTimeMillis()); }
default Date getCurrentDate() { return Date.from(getCurrentInstant()); }

default Instant getCurrentInstant() { return Instant.now(); }
}
9 changes: 6 additions & 3 deletions server/src/test/java/org/cloudfoundry/identity/uaa/oauth/ApprovalServiceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.client.BaseClientDetails;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import java.util.List;

Expand Down Expand Up @@ -50,6 +52,7 @@ public void ensureRequiredApprovals_happyCase() {
approval.setStatus(Approval.ApprovalStatus.APPROVED);
approval.setExpiresAt(new Date(approvalExpiry));
when(timeService.getCurrentTimeMillis()).thenReturn(approvalExpiry - 1L);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(approvalExpiry - 1L, ChronoUnit.MILLIS));
when(timeService.getCurrentDate()).thenCallRealMethod();

List<Approval> approvals = Lists.newArrayList(approval);
Expand All @@ -68,7 +71,7 @@ public void ensureRequiredApprovals_throwsWhenApprovalsExpired() {
approval.setScope("foo.read");
approval.setStatus(Approval.ApprovalStatus.APPROVED);
approval.setExpiresAt(new Date(approvalExpiry));
when(timeService.getCurrentTimeMillis()).thenReturn(approvalExpiry + 1L);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(approvalExpiry + 1L, ChronoUnit.MILLIS));
when(timeService.getCurrentDate()).thenCallRealMethod();

List<Approval> approvals = Lists.newArrayList(approval);
Expand Down Expand Up @@ -112,7 +115,7 @@ public void ensureRequiredApprovals_iteratesThroughAllApprovalsAndScopes() {
approval3.setStatus(Approval.ApprovalStatus.APPROVED);
approval3.setExpiresAt(new Date(approvalExpiry));

when(timeService.getCurrentTimeMillis()).thenReturn(approvalExpiry - 1L);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(approvalExpiry - 1L, ChronoUnit.MILLIS));
when(timeService.getCurrentDate()).thenCallRealMethod();

List<Approval> approvals = Lists.newArrayList(approval1, approval2, approval3);
Expand Down Expand Up @@ -140,7 +143,7 @@ public void ensureRequiredApprovals_throwsIfAnyRequestedScopesAreNotApproved() {
approval3.setStatus(Approval.ApprovalStatus.APPROVED);
approval3.setExpiresAt(new Date(approvalExpiry));

when(timeService.getCurrentTimeMillis()).thenReturn(approvalExpiry - 1L);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(approvalExpiry - 1L, ChronoUnit.MILLIS));
when(timeService.getCurrentDate()).thenCallRealMethod();

List<Approval> approvals = Lists.newArrayList(approval1, approval2, approval3);
Expand Down
6 changes: 4 additions & 2 deletions server/src/test/java/org/cloudfoundry/identity/uaa/oauth/CheckTokenEndpointTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@
import org.springframework.test.util.ReflectionTestUtils;
import org.springframework.web.HttpRequestMethodNotSupportedException;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
Expand Down Expand Up @@ -223,7 +225,7 @@ public void setUp(boolean opaque) throws Exception {

nowMillis = 10000L;
timeService = mock(TimeService.class);
when(timeService.getCurrentTimeMillis()).thenReturn(nowMillis);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(nowMillis, ChronoUnit.MILLIS));
when(timeService.getCurrentDate()).thenCallRealMethod();
userAuthorities = new ArrayList<>();
userAuthorities.add(new SimpleGrantedAuthority("read"));
Expand Down Expand Up @@ -942,7 +944,7 @@ public void testExpiredToken() throws Exception {
tokenServices.setClientDetailsService(clientDetailsService);
OAuth2AccessToken accessToken = tokenServices.createAccessToken(authentication);

when(timeService.getCurrentTimeMillis()).thenReturn(nowMillis + validitySeconds.longValue() * 1000 + 1L);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(nowMillis + validitySeconds.longValue() * 1000 + 1L, ChronoUnit.MILLIS));
endpoint.checkToken(accessToken.getValue(), Collections.emptyList(), request);
}

Expand Down
5 changes: 5 additions & 0 deletions ...er/src/test/java/org/cloudfoundry/identity/uaa/oauth/DeprecatedUaaTokenServicesTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.*;

import static java.util.Collections.*;
Expand Down Expand Up @@ -128,6 +130,7 @@ public void setUp() throws Exception {
tokenServices = tokenSupport.getUaaTokenServices();
tokenProvisioning = tokenSupport.getTokenProvisioning();
when(tokenSupport.timeService.getCurrentTimeMillis()).thenReturn(1000L);
when(tokenSupport.timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(1000L, ChronoUnit.MILLIS));
}

@After
Expand Down Expand Up @@ -249,6 +252,7 @@ public void refreshAccessToken_buildsIdToken_withRolesAndAttributesAndACR() thro

TimeService timeService = mock(TimeService.class);
when(timeService.getCurrentTimeMillis()).thenReturn(1000L);
when(timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(1000L, ChronoUnit.MILLIS));
when(timeService.getCurrentDate()).thenCallRealMethod();
RefreshTokenCreator refreshTokenCreator = mock(RefreshTokenCreator.class);
ApprovalService approvalService = mock(ApprovalService.class);
Expand Down Expand Up @@ -1823,6 +1827,7 @@ public void testLoadAuthenticationWithAnExpiredToken() {
assertThat(accessToken, validFor(is(1)));

when(tokenSupport.timeService.getCurrentTimeMillis()).thenReturn(2001L);
when(tokenSupport.timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(2001L, ChronoUnit.MILLIS));
try {
tokenServices.loadAuthentication(accessToken.getValue());
fail("Expected Exception was not thrown");
Expand Down
4 changes: 3 additions & 1 deletion server/src/test/java/org/cloudfoundry/identity/uaa/oauth/RefreshRotationTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.security.oauth2.provider.client.BaseClientDetails;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
Expand Down Expand Up @@ -55,7 +57,7 @@ void setUp() throws Exception {
persistToken.setExpiration(expiration);

tokenServices = tokenSupport.getUaaTokenServices();
when(tokenSupport.timeService.getCurrentTimeMillis()).thenReturn(1000L);
when(tokenSupport.timeService.getCurrentInstant()).thenReturn(Instant.EPOCH.plus(1000L, ChronoUnit.MILLIS));
new IdentityZoneManagerImpl().getCurrentIdentityZone().getConfig().getTokenPolicy().setRefreshTokenFormat(TokenConstants.TokenFormat.OPAQUE.getStringValue());
}

Expand Down
1 change: 1 addition & 0 deletions server/src/test/java/org/cloudfoundry/identity/uaa/oauth/TokenTestSupport.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,7 @@ public TokenTestSupport(UaaTokenEnhancer tokenEnhancer, KeyInfoService keyInfo)
requestFactory = new DefaultOAuth2RequestFactory(clientDetailsService);
timeService = mock(TimeService.class);
approvalService = new ApprovalService(timeService, approvalStore);
when(timeService.getCurrentInstant()).thenCallRealMethod();
when(timeService.getCurrentDate()).thenCallRealMethod();
TokenEndpointBuilder tokenEndpointBuilder = new TokenEndpointBuilder(DEFAULT_ISSUER);
keyInfoService = keyInfo != null ? keyInfo : new KeyInfoService(DEFAULT_ISSUER);
Expand Down
41 changes: 22 additions & 19 deletions server/src/test/java/org/cloudfoundry/identity/uaa/oauth/UaaTokenStoreTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
import org.cloudfoundry.identity.uaa.authentication.UaaAuthenticationDetails;
import org.cloudfoundry.identity.uaa.authentication.UaaPrincipal;
import org.cloudfoundry.identity.uaa.constants.OriginKeys;
import org.cloudfoundry.identity.uaa.util.TimeService;
import org.cloudfoundry.identity.uaa.util.TimeServiceImpl;
import org.cloudfoundry.identity.uaa.util.UaaStringUtils;
import org.cloudfoundry.identity.uaa.zone.IdentityZone;
import org.junit.jupiter.api.BeforeEach;
Expand Down Expand Up @@ -34,10 +36,8 @@
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.time.Duration;
import java.time.Instant;
import java.time.temporal.Temporal;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
Expand All @@ -56,7 +56,6 @@
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.atMost;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
Expand Down Expand Up @@ -86,7 +85,7 @@ void setUp() {
List<GrantedAuthority> userAuthorities = Collections.singletonList(new SimpleGrantedAuthority(
"openid"));

store = new UaaTokenStore(dataSource);
store = new UaaTokenStore(dataSource, givenMockedTime());
legacyCodeServices = new JdbcAuthorizationCodeServices(dataSource);
BaseClientDetails client = new BaseClientDetails("clientid", null, "openid", "client_credentials,password", "oauth.login", null);
Map<String, String> parameters = new HashMap<>();
Expand Down Expand Up @@ -198,9 +197,11 @@ void retrieveToken() {

@Test
void retrieveExpiredToken() {
TimeService timeMock = givenMockedTime();
store = new UaaTokenStore(dataSource, timeMock);
String code = store.createAuthorizationCode(clientAuthentication);
assertThat(jdbcTemplate.queryForObject("SELECT count(*) FROM oauth_code WHERE code = ?", new Object[]{code}, Integer.class), is(1));
jdbcTemplate.update("update oauth_code set expiresat = 1");
doReturn(Instant.now().plus(UaaTokenStore.DEFAULT_EXPIRATION_TIME)).when(timeMock).getCurrentInstant();
assertThrows(InvalidGrantException.class, () -> store.consumeAuthorizationCode(code));
}

Expand All @@ -214,14 +215,15 @@ void retrieveNonExistentToken() {
@Test
void cleanUpExpiredTokensBasedOnExpiresField() {
int count = 10;
store = new UaaTokenStore(dataSource, givenMockedExpiration());
TimeService timeMock = givenMockedTime();
store = new UaaTokenStore(dataSource, timeMock);
String lastCode = null;
for (int i = 0; i < count; i++) {
lastCode = store.createAuthorizationCode(clientAuthentication);
}
assertThat(jdbcTemplate.queryForObject("SELECT count(*) FROM oauth_code", Integer.class), is(count));

jdbcTemplate.update("UPDATE oauth_code SET expiresat = ?", System.currentTimeMillis() - 60000);
doReturn(Instant.now().plus(UaaTokenStore.LEGACY_CODE_EXPIRATION_TIME)).when(timeMock).getCurrentInstant();

final String finalLastCode = lastCode;
assertThrows(InvalidGrantException.class, () -> store.consumeAuthorizationCode(finalLastCode));
Expand All @@ -232,15 +234,16 @@ void cleanUpExpiredTokensBasedOnExpiresField() {
void cleanUpLegacyCodesCodesWithoutExpiresAtAfter3Days() {
int count = 10;
long oneday = 1000 * 60 * 60 * 24;
store = new UaaTokenStore(dataSource, givenMockedExpiration());
TimeService timeMock = givenMockedTime();
store = new UaaTokenStore(dataSource, timeMock);
for (int i = 0; i < count; i++) {
legacyCodeServices.createAuthorizationCode(clientAuthentication);
}
assertThat(jdbcTemplate.queryForObject("SELECT count(*) FROM oauth_code", Integer.class), is(count));
jdbcTemplate.update("UPDATE oauth_code SET created = ?", new Timestamp(System.currentTimeMillis() - (2 * oneday)));
doReturn(Instant.now().plus(Duration.ofDays(2))).when(timeMock).getCurrentInstant();
assertThrows(InvalidGrantException.class, () -> store.consumeAuthorizationCode("non-existent"));
assertThat(jdbcTemplate.queryForObject("SELECT count(*) FROM oauth_code", Integer.class), is(count));
jdbcTemplate.update("UPDATE oauth_code SET created = ?", new Timestamp(System.currentTimeMillis() - (4 * oneday)));
doReturn(Instant.now().plus(Duration.ofDays(4))).when(timeMock).getCurrentInstant();
assertThrows(InvalidGrantException.class, () -> store.consumeAuthorizationCode("non-existent"));
assertThat(jdbcTemplate.queryForObject("SELECT count(*) FROM oauth_code", Integer.class), is(0));
}
Expand Down Expand Up @@ -289,7 +292,7 @@ void cleanUpUnusedOldTokensMySQLInAnotherTimezone(
throw new RuntimeException("Unknown DB profile:" + db);
}

store = new UaaTokenStore(sameConnectionDataSource);
store = new UaaTokenStore(sameConnectionDataSource, givenMockedTime());
legacyCodeServices = new JdbcAuthorizationCodeServices(sameConnectionDataSource);
int count = 10;
String lastCode = null;
Expand All @@ -304,7 +307,7 @@ void cleanUpUnusedOldTokensMySQLInAnotherTimezone(
}
assertThat(template.queryForObject("SELECT count(*) FROM oauth_code", Integer.class), is(count - 1));
} finally {
store = new UaaTokenStore(dataSource);
store = new UaaTokenStore(dataSource, givenMockedTime());
legacyCodeServices = new JdbcAuthorizationCodeServices(dataSource);
}
}
Expand All @@ -318,7 +321,7 @@ void cleanUpExpiredTokensDeadlockLoser() throws Exception {

SameConnectionDataSource sameConnectionDataSource = new SameConnectionDataSource(expirationLoser);

store = new UaaTokenStore(sameConnectionDataSource, Duration.ofMillis(1));
store = new UaaTokenStore(sameConnectionDataSource, givenMockedTime(), Duration.ofMillis(1));
int count = 10;
for (int i = 0; i < count; i++) {
String code = store.createAuthorizationCode(clientAuthentication);
Expand All @@ -328,7 +331,7 @@ void cleanUpExpiredTokensDeadlockLoser() throws Exception {
}
}
} finally {
store = new UaaTokenStore(dataSource);
store = new UaaTokenStore(dataSource, givenMockedTime());
}
}

Expand All @@ -346,7 +349,7 @@ void testCountingTheExecutedSqlDeleteStatements() throws SQLException {
// Given, mocked data source to count how often it is used, call performExpirationClean 10 times.
DataSource mockedDataSource = mock(DataSource.class);
Instant before = Instant.now();
store = new UaaTokenStore(mockedDataSource);
store = new UaaTokenStore(mockedDataSource, givenMockedTime());
// When
for (int i = 0; i < 10; i++) {
try {
Expand Down Expand Up @@ -482,10 +485,10 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
}
}

private static Duration givenMockedExpiration() {
Duration durationMock = mock(Duration.class);
doReturn(Instant.now().plus(UaaTokenStore.DEFAULT_EXPIRATION_TIME)).when(durationMock).addTo(any(Temporal.class));
return durationMock;
private static TimeService givenMockedTime() {
TimeServiceImpl timeService = mock(TimeServiceImpl.class);
doReturn(Instant.now()).when(timeService).getCurrentInstant();
return timeService;
}

private static final byte[] UAA_AUTHENTICATION_DATA_OLD_STYLE = new byte[]{123, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 114, 101, 115, 112, 111, 110, 115, 101, 84, 121, 112, 101, 115, 34, 58, 91, 93, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 114, 101, 115, 111, 117, 114, 99, 101, 73, 100, 115, 34, 58, 91, 93, 44, 34, 117, 115, 101, 114, 65, 117, 116, 104, 101, 110, 116, 105, 99, 97, 116, 105, 111, 110, 46, 117, 97, 97, 80, 114, 105, 110, 99, 105, 112, 97, 108, 34, 58, 34, 123, 92, 34, 105, 100, 92, 34, 58, 92, 34, 117, 115, 101, 114, 105, 100, 92, 34, 44, 92, 34, 110, 97, 109, 101, 92, 34, 58, 92, 34, 117, 115, 101, 114, 110, 97, 109, 101, 92, 34, 44, 92, 34, 101, 109, 97, 105, 108, 92, 34, 58, 92, 34, 117, 115, 101, 114, 110, 97, 109, 101, 64, 116, 101, 115, 116, 46, 111, 114, 103, 92, 34, 44, 92, 34, 111, 114, 105, 103, 105, 110, 92, 34, 58, 92, 34, 117, 97, 97, 92, 34, 44, 92, 34, 101, 120, 116, 101, 114, 110, 97, 108, 73, 100, 92, 34, 58, 110, 117, 108, 108, 44, 92, 34, 122, 111, 110, 101, 73, 100, 92, 34, 58, 92, 34, 117, 97, 97, 92, 34, 125, 34, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 114, 101, 113, 117, 101, 115, 116, 80, 97, 114, 97, 109, 101, 116, 101, 114, 115, 34, 58, 123, 34, 103, 114, 97, 110, 116, 95, 116, 121, 112, 101, 34, 58, 34, 112, 97, 115, 115, 119, 111, 114, 100, 34, 44, 34, 99, 108, 105, 101, 110, 116, 95, 105, 100, 34, 58, 34, 99, 108, 105, 101, 110, 116, 105, 100, 34, 44, 34, 115, 99, 111, 112, 101, 34, 58, 34, 111, 112, 101, 110, 105, 100, 34, 125, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 114, 101, 100, 105, 114, 101, 99, 116, 85, 114, 105, 34, 58, 110, 117, 108, 108, 44, 34, 117, 115, 101, 114, 65, 117, 116, 104, 101, 110, 116, 105, 99, 97, 116, 105, 111, 110, 46, 97, 117, 116, 104, 111, 114, 105, 116, 105, 101, 115, 34, 58, 91, 34, 111, 112, 101, 110, 105, 100, 34, 93, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 97, 117, 116, 104, 111, 114, 105, 116, 105, 101, 115, 34, 58, 91, 34, 111, 97, 117, 116, 104, 46, 108, 111, 103, 105, 110, 34, 93, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 99, 108, 105, 101, 110, 116, 73, 100, 34, 58, 34, 99, 108, 105, 101, 110, 116, 105, 100, 34, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 97, 112, 112, 114, 111, 118, 101, 100, 34, 58, 116, 114, 117, 101, 44, 34, 111, 97, 117, 116, 104, 50, 82, 101, 113, 117, 101, 115, 116, 46, 115, 99, 111, 112, 101, 34, 58, 91, 34, 111, 112, 101, 110, 105, 100, 34, 93, 125};
Expand Down
Loading

0 comments on commit 9569b5b

Please sign in to comment.