Optimize fetching tokenKeys for reactive actuators

Closes gh-10899

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/cloudfoundry/reactive/ReactiveTokenValidator.java

@@ -23,6 +23,8 @@ import java.security.PublicKey;
 import java.security.Signature;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.X509EncodedKeySpec;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 
 import reactor.core.publisher.Mono;
@@ -41,6 +43,8 @@ class ReactiveTokenValidator {
 
 	private final ReactiveCloudFoundrySecurityService securityService;
 
+	private Map<String, String> cachedTokenKeys = new ConcurrentHashMap<>();
+
 	ReactiveTokenValidator(ReactiveCloudFoundrySecurityService securityService) {
 		this.securityService = securityService;
 	}
@@ -67,11 +71,17 @@ class ReactiveTokenValidator {
 
 	private Mono<Void> validateKeyIdAndSignature(Token token) {
 		String keyId = token.getKeyId();
-		return this.securityService.fetchTokenKeys()
+		return Mono.just(this.cachedTokenKeys)
 				.filter((tokenKeys) -> tokenKeys.containsKey(keyId))
-				.switchIfEmpty(Mono.error(
+				.switchIfEmpty(this.securityService.fetchTokenKeys()
-						new CloudFoundryAuthorizationException(Reason.INVALID_KEY_ID,
+						.doOnSuccess(fetchedTokenKeys -> {
-								"Key Id present in token header does not match")))
+							this.cachedTokenKeys.clear();
+							this.cachedTokenKeys.putAll(fetchedTokenKeys);
+						})
+						.filter((tokenKeys) -> tokenKeys.containsKey(keyId))
+						.switchIfEmpty((Mono.error(
+								new CloudFoundryAuthorizationException(Reason.INVALID_KEY_ID,
+										"Key Id present in token header does not match")))))
 				.filter((tokenKeys) -> hasValidSignature(token, tokenKeys.get(keyId)))
 				.switchIfEmpty(Mono.error(new CloudFoundryAuthorizationException(
 						Reason.INVALID_SIGNATURE, "RSA Signature did not match content")))

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/cloudfoundry/reactive/ReactiveTokenValidatorTests.java

@@ -26,6 +26,7 @@ import java.security.Signature;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.PKCS8EncodedKeySpec;
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.Map;
 
 import org.apache.commons.codec.binary.Base64;
@@ -35,10 +36,12 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
+import reactor.test.publisher.PublisherProbe;
 
 import org.springframework.boot.actuate.autoconfigure.cloudfoundry.CloudFoundryAuthorizationException;
 import org.springframework.boot.actuate.autoconfigure.cloudfoundry.CloudFoundryAuthorizationException.Reason;
 import org.springframework.boot.actuate.autoconfigure.cloudfoundry.Token;
+import org.springframework.test.util.ReflectionTestUtils;
 import org.springframework.util.Base64Utils;
 import org.springframework.util.StreamUtils;
 
@@ -77,26 +80,84 @@ public class ReactiveTokenValidatorTests {
 			+ "r3F7aM9YpErzeYLrl0GhQr9BVJxOvXcVd4kmY+XkiCcrkyS1cnghnllh+LCwQu1s\n"
 			+ "YwIDAQAB\n-----END PUBLIC KEY-----";
 
-	private static final Map<String, String> INVALID_KEYS = Collections
+	private static final Map<String, String> INVALID_KEYS = new LinkedHashMap<>();
-			.singletonMap("invalid-key", INVALID_KEY);
 
-	private static final Map<String, String> VALID_KEYS = Collections
+	private static final Map<String, String> VALID_KEYS = new LinkedHashMap<>();
-			.singletonMap("valid-key", VALID_KEY);
 
 	@Before
 	public void setup() throws Exception {
 		MockitoAnnotations.initMocks(this);
+		VALID_KEYS.put("valid-key", VALID_KEY);
+		INVALID_KEYS.put("invalid-key", INVALID_KEY);
 		this.tokenValidator = new ReactiveTokenValidator(this.securityService);
 	}
 
 	@Test
-	public void validateTokenWhenKidValidationFailsShouldThrowException()
+	public void validateTokenWhenKidValidationFailsTwiceShouldThrowException() throws Exception {
-			throws Exception {
+		PublisherProbe<Map<String, String>> fetchTokenKeys = PublisherProbe.of(Mono.just(VALID_KEYS));
-		given(this.securityService.fetchTokenKeys()).willReturn(Mono.just(INVALID_KEYS));
+		ReflectionTestUtils.setField(this.tokenValidator, "cachedTokenKeys", VALID_KEYS);
+		given(this.securityService.fetchTokenKeys()).willReturn(fetchTokenKeys.mono());
+		given(this.securityService.getUaaUrl())
+				.willReturn(Mono.just("http://localhost:8080/uaa"));
+		String header = "{\"alg\": \"RS256\",  \"kid\": \"invalid-key\",\"typ\": \"JWT\"}";
+		String claims = "{\"exp\": 2147483647, \"iss\": \"http://localhost:8080/uaa/oauth/token\", \"scope\": [\"actuator.read\"]}";
+		StepVerifier
+				.create(this.tokenValidator.validate(
+						new Token(getSignedToken(header.getBytes(), claims.getBytes()))))
+				.consumeErrorWith((ex) -> {
+					assertThat(ex).isExactlyInstanceOf(
+							CloudFoundryAuthorizationException.class);
+					assertThat(((CloudFoundryAuthorizationException) ex).getReason())
+							.isEqualTo(Reason.INVALID_KEY_ID);
+				}).verify();
+		Object cachedTokenKeys = ReflectionTestUtils.getField(this.tokenValidator, "cachedTokenKeys");
+		assertThat(cachedTokenKeys).isEqualTo(VALID_KEYS);
+		fetchTokenKeys.assertWasSubscribed();
+	}
+
+	@Test
+	public void validateTokenWhenKidValidationSucceedsInTheSecondAttempt() throws Exception {
+		PublisherProbe<Map<String, String>> fetchTokenKeys = PublisherProbe.of(Mono.just(VALID_KEYS));
+		ReflectionTestUtils.setField(this.tokenValidator, "cachedTokenKeys", INVALID_KEYS);
+		given(this.securityService.fetchTokenKeys()).willReturn(fetchTokenKeys.mono());
+		given(this.securityService.getUaaUrl())
+				.willReturn(Mono.just("http://localhost:8080/uaa"));
+		String header = "{\"alg\": \"RS256\",  \"kid\": \"valid-key\",\"typ\": \"JWT\"}";
+		String claims = "{\"exp\": 2147483647, \"iss\": \"http://localhost:8080/uaa/oauth/token\", \"scope\": [\"actuator.read\"]}";
+		StepVerifier
+				.create(this.tokenValidator.validate(
+						new Token(getSignedToken(header.getBytes(), claims.getBytes()))))
+				.verifyComplete();
+		Object cachedTokenKeys = ReflectionTestUtils.getField(this.tokenValidator, "cachedTokenKeys");
+		assertThat(cachedTokenKeys).isEqualTo(VALID_KEYS);
+		fetchTokenKeys.assertWasSubscribed();
+	}
+
+	@Test
+	public void validateTokenWhenCacheIsEmptyShouldFetchTokenKeys() throws Exception {
+		PublisherProbe<Map<String, String>> fetchTokenKeys = PublisherProbe.of(Mono.just(VALID_KEYS));
+		given(this.securityService.fetchTokenKeys()).willReturn(fetchTokenKeys.mono());
 		given(this.securityService.getUaaUrl())
 				.willReturn(Mono.just("http://localhost:8080/uaa"));
 		String header = "{\"alg\": \"RS256\",  \"kid\": \"valid-key\",\"typ\": \"JWT\"}";
 		String claims = "{\"exp\": 2147483647, \"iss\": \"http://localhost:8080/uaa/oauth/token\", \"scope\": [\"actuator.read\"]}";
+		StepVerifier
+				.create(this.tokenValidator.validate(
+						new Token(getSignedToken(header.getBytes(), claims.getBytes()))))
+				.verifyComplete();
+		Object cachedTokenKeys = ReflectionTestUtils.getField(this.tokenValidator, "cachedTokenKeys");
+		assertThat(cachedTokenKeys).isEqualTo(VALID_KEYS);
+		fetchTokenKeys.assertWasSubscribed();
+	}
+
+	@Test
+	public void validateTokenWhenCacheEmptyAndInvalidKeyShouldThrowException() throws Exception {
+		PublisherProbe<Map<String, String>> fetchTokenKeys = PublisherProbe.of(Mono.just(VALID_KEYS));
+		given(this.securityService.fetchTokenKeys()).willReturn(fetchTokenKeys.mono());
+		given(this.securityService.getUaaUrl())
+				.willReturn(Mono.just("http://localhost:8080/uaa"));
+		String header = "{\"alg\": \"RS256\",  \"kid\": \"invalid-key\",\"typ\": \"JWT\"}";
+		String claims = "{\"exp\": 2147483647, \"iss\": \"http://localhost:8080/uaa/oauth/token\", \"scope\": [\"actuator.read\"]}";
 		StepVerifier
 				.create(this.tokenValidator.validate(
 						new Token(getSignedToken(header.getBytes(), claims.getBytes()))))
@@ -106,19 +167,25 @@ public class ReactiveTokenValidatorTests {
 					assertThat(((CloudFoundryAuthorizationException) ex).getReason())
 							.isEqualTo(Reason.INVALID_KEY_ID);
 				}).verify();
+		Object cachedTokenKeys = ReflectionTestUtils.getField(this.tokenValidator, "cachedTokenKeys");
+		assertThat(cachedTokenKeys).isEqualTo(VALID_KEYS);
+		fetchTokenKeys.assertWasSubscribed();
 	}
 
 	@Test
-	public void validateTokenWhenKidValidationSucceeds() throws Exception {
+	public void validateTokenWhenCacheValidShouldNotFetchTokenKeys() throws Exception {
-		given(this.securityService.fetchTokenKeys()).willReturn(Mono.just(VALID_KEYS));
+		PublisherProbe<Map<String, String>> fetchTokenKeys = PublisherProbe.empty();
+		ReflectionTestUtils.setField(this.tokenValidator, "cachedTokenKeys", VALID_KEYS);
+		given(this.securityService.fetchTokenKeys()).willReturn(fetchTokenKeys.mono());
 		given(this.securityService.getUaaUrl())
 				.willReturn(Mono.just("http://localhost:8080/uaa"));
-		String header = "{ \"alg\": \"RS256\",  \"kid\": \"valid-key\",\"typ\": \"JWT\"}";
+		String header = "{\"alg\": \"RS256\",  \"kid\": \"valid-key\",\"typ\": \"JWT\"}";
-		String claims = "{ \"exp\": 2147483647, \"iss\": \"http://localhost:8080/uaa/oauth/token\", \"scope\": [\"actuator.read\"]}";
+		String claims = "{\"exp\": 2147483647, \"iss\": \"http://localhost:8080/uaa/oauth/token\", \"scope\": [\"actuator.read\"]}";
 		StepVerifier
 				.create(this.tokenValidator.validate(
 						new Token(getSignedToken(header.getBytes(), claims.getBytes()))))
 				.verifyComplete();
+		fetchTokenKeys.assertWasNotSubscribed();
 	}
 
 	@Test