From 5b7d900424bc5bcc47247b196c356ce55338682b Mon Sep 17 00:00:00 2001 From: Joe Grandja Date: Tue, 22 Mar 2022 06:17:21 -0400 Subject: [PATCH] Optimize InMemoryOAuth2AuthorizationService Closes gh-654 --- .../InMemoryOAuth2AuthorizationService.java | 65 +++++++++++++++++-- ...tionCodeRequestAuthenticationProvider.java | 2 - ...MemoryOAuth2AuthorizationServiceTests.java | 39 ++++++++++- 3 files changed, 99 insertions(+), 7 deletions(-) diff --git a/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationService.java b/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationService.java index 46b4fd6e..763b2a26 100644 --- a/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationService.java +++ b/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationService.java @@ -1,5 +1,5 @@ /* - * Copyright 2020-2021 the original author or authors. + * Copyright 2020-2022 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,6 +17,7 @@ package org.springframework.security.oauth2.server.authorization; import java.util.Arrays; import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -41,8 +42,29 @@ import org.springframework.util.Assert; * @see OAuth2AuthorizationService */ public final class InMemoryOAuth2AuthorizationService implements OAuth2AuthorizationService { + private int maxInitializedAuthorizations = 100; + + /* + * Stores "initialized" (uncompleted) authorizations, where an access token has not yet been granted. + * This state occurs with the authorization_code grant flow during the user consent step OR + * when the code is returned in the authorization response but the access token request is not yet initiated. + */ + private Map initializedAuthorizations = + Collections.synchronizedMap(new MaxSizeHashMap<>(this.maxInitializedAuthorizations)); + + /* + * Stores "completed" authorizations, where an access token has been granted. + */ private final Map authorizations = new ConcurrentHashMap<>(); + /* + * Constructor used for testing only. + */ + InMemoryOAuth2AuthorizationService(int maxInitializedAuthorizations) { + this.maxInitializedAuthorizations = maxInitializedAuthorizations; + this.initializedAuthorizations = Collections.synchronizedMap(new MaxSizeHashMap<>(this.maxInitializedAuthorizations)); + } + /** * Constructs an {@code InMemoryOAuth2AuthorizationService}. */ @@ -77,20 +99,31 @@ public final class InMemoryOAuth2AuthorizationService implements OAuth2Authoriza @Override public void save(OAuth2Authorization authorization) { Assert.notNull(authorization, "authorization cannot be null"); - this.authorizations.put(authorization.getId(), authorization); + if (isComplete(authorization)) { + this.authorizations.put(authorization.getId(), authorization); + } else { + this.initializedAuthorizations.put(authorization.getId(), authorization); + } } @Override public void remove(OAuth2Authorization authorization) { Assert.notNull(authorization, "authorization cannot be null"); - this.authorizations.remove(authorization.getId(), authorization); + if (isComplete(authorization)) { + this.authorizations.remove(authorization.getId(), authorization); + } else { + this.initializedAuthorizations.remove(authorization.getId(), authorization); + } } @Nullable @Override public OAuth2Authorization findById(String id) { Assert.hasText(id, "id cannot be empty"); - return this.authorizations.get(id); + OAuth2Authorization authorization = this.authorizations.get(id); + return authorization != null ? + authorization : + this.initializedAuthorizations.get(id); } @Nullable @@ -102,9 +135,18 @@ public final class InMemoryOAuth2AuthorizationService implements OAuth2Authoriza return authorization; } } + for (OAuth2Authorization authorization : this.initializedAuthorizations.values()) { + if (hasToken(authorization, token, tokenType)) { + return authorization; + } + } return null; } + private static boolean isComplete(OAuth2Authorization authorization) { + return authorization.getAccessToken() != null; + } + private static boolean hasToken(OAuth2Authorization authorization, String token, @Nullable OAuth2TokenType tokenType) { if (tokenType == null) { return matchesState(authorization, token) || @@ -144,4 +186,19 @@ public final class InMemoryOAuth2AuthorizationService implements OAuth2Authoriza authorization.getToken(OAuth2RefreshToken.class); return refreshToken != null && refreshToken.getToken().getTokenValue().equals(token); } + + private static final class MaxSizeHashMap extends LinkedHashMap { + private final int maxSize; + + private MaxSizeHashMap(int maxSize) { + this.maxSize = maxSize; + } + + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + return size() > this.maxSize; + } + + } + } diff --git a/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProvider.java b/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProvider.java index 9350d713..b5ddca7b 100644 --- a/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProvider.java +++ b/oauth2-authorization-server/src/main/java/org/springframework/security/oauth2/server/authorization/authentication/OAuth2AuthorizationCodeRequestAuthenticationProvider.java @@ -266,8 +266,6 @@ public final class OAuth2AuthorizationCodeRequestAuthenticationProvider implemen .build(); this.authorizationService.save(authorization); - // TODO Need to remove 'in-flight' authorization if consent step is not completed (e.g. approved or cancelled) - Set currentAuthorizedScopes = currentAuthorizationConsent != null ? currentAuthorizationConsent.getScopes() : null; diff --git a/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationServiceTests.java b/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationServiceTests.java index eecd67da..2a090602 100644 --- a/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationServiceTests.java +++ b/oauth2-authorization-server/src/test/java/org/springframework/security/oauth2/server/authorization/InMemoryOAuth2AuthorizationServiceTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2020-2021 the original author or authors. + * Copyright 2020-2022 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -132,6 +132,43 @@ public class InMemoryOAuth2AuthorizationServiceTests { assertThat(authorization).isNotEqualTo(originalAuthorization); } + @Test + public void saveWhenInitializedAuthorizationsReachMaxThenOldestRemoved() { + int maxInitializedAuthorizations = 5; + InMemoryOAuth2AuthorizationService authorizationService = + new InMemoryOAuth2AuthorizationService(maxInitializedAuthorizations); + + OAuth2Authorization initialAuthorization = OAuth2Authorization.withRegisteredClient(REGISTERED_CLIENT) + .id(ID + "-initial") + .principalName(PRINCIPAL_NAME) + .authorizationGrantType(AUTHORIZATION_GRANT_TYPE) + .attribute(OAuth2ParameterNames.STATE, "state-initial") + .build(); + authorizationService.save(initialAuthorization); + + OAuth2Authorization authorization = authorizationService.findById(initialAuthorization.getId()); + assertThat(authorization).isEqualTo(initialAuthorization); + authorization = authorizationService.findByToken( + initialAuthorization.getAttribute(OAuth2ParameterNames.STATE), STATE_TOKEN_TYPE); + assertThat(authorization).isEqualTo(initialAuthorization); + + for (int i = 0; i < maxInitializedAuthorizations; i++) { + authorization = OAuth2Authorization.withRegisteredClient(REGISTERED_CLIENT) + .id(ID + "-" + i) + .principalName(PRINCIPAL_NAME) + .authorizationGrantType(AUTHORIZATION_GRANT_TYPE) + .attribute(OAuth2ParameterNames.STATE, "state-" + i) + .build(); + authorizationService.save(authorization); + } + + authorization = authorizationService.findById(initialAuthorization.getId()); + assertThat(authorization).isNull(); + authorization = authorizationService.findByToken( + initialAuthorization.getAttribute(OAuth2ParameterNames.STATE), STATE_TOKEN_TYPE); + assertThat(authorization).isNull(); + } + @Test public void removeWhenAuthorizationNullThenThrowIllegalArgumentException() { assertThatThrownBy(() -> this.authorizationService.remove(null))