JwkTokenStore Supports Multiple JWK Set URLs

Fixes spring-atticgh-1050

Author: rwinch

spring-security-oauth2/pom.xml

@@ -193,6 +193,13 @@
             <optional>true</optional>
 		</dependency>
 
+		<dependency>
+			<groupId>com.squareup.okhttp3</groupId>
+			<artifactId>mockwebserver</artifactId>
+			<version>3.7.0</version>
+			<scope>test</scope>
+		</dependency>
+
 		<dependency>
 			<groupId>junit</groupId>
 			<artifactId>junit</artifactId>

spring-security-oauth2/src/main/java/org/springframework/security/oauth2/provider/token/store/jwk/JwkDefinitionSource.java

@@ -27,9 +27,7 @@
 import java.security.KeyFactory;
 import java.security.interfaces.RSAPublicKey;
 import java.security.spec.RSAPublicKeySpec;
-import java.util.LinkedHashMap;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -44,7 +42,7 @@
  * @author Joe Grandja
  */
 class JwkDefinitionSource {
-	private final URL jwkSetUrl;
+	private final List<URL> jwkSetUrls;
 	private final Map<String, JwkDefinitionHolder> jwkDefinitions = new ConcurrentHashMap<String, JwkDefinitionHolder>();
 	private static final JwkSetConverter jwkSetConverter = new JwkSetConverter();
 
@@ -54,10 +52,22 @@ class JwkDefinitionSource {
 	 * @param jwkSetUrl the JWK Set URL
 	 */
 	JwkDefinitionSource(String jwkSetUrl) {
-		try {
-			this.jwkSetUrl = new URL(jwkSetUrl);
-		} catch (MalformedURLException ex) {
-			throw new IllegalArgumentException("Invalid JWK Set URL: " + ex.getMessage(), ex);
+		this(Arrays.asList(jwkSetUrl));
+	}
+
+	/**
+	 * Creates a new instance using the provided URLs as the location for the JWK Sets.
+	 *
+	 * @param jwkSetUrls the JWK Set URLs
+	 */
+	JwkDefinitionSource(List<String> jwkSetUrls) {
+		this.jwkSetUrls = new ArrayList<URL>();
+		for(String jwkSetUrl : jwkSetUrls) {
+			try {
+				this.jwkSetUrls.add(new URL(jwkSetUrl));
+			} catch (MalformedURLException ex) {
+				throw new IllegalArgumentException("Invalid JWK Set URL: " + ex.getMessage(), ex);
+			}
 		}
 	}
 
@@ -90,7 +100,9 @@ JwkDefinition getDefinitionLoadIfNecessary(String keyId) {
 			return result;
 		}
 		this.jwkDefinitions.clear();
-		this.jwkDefinitions.putAll(loadJwkDefinitions(this.jwkSetUrl));
+		for(URL jwkSetUrl : jwkSetUrls) {
+			this.jwkDefinitions.putAll(loadJwkDefinitions(jwkSetUrl));
+		}
 		return this.getDefinition(keyId);
 	}
 

spring-security-oauth2/src/main/java/org/springframework/security/oauth2/provider/token/store/jwk/JwkTokenStore.java

@@ -23,7 +23,9 @@
 import org.springframework.security.oauth2.provider.token.store.JwtTokenStore;
 import org.springframework.util.Assert;
 
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 
 /**
  * A {@link TokenStore} implementation that provides support for verifying the
@@ -86,16 +88,23 @@
  * @author Joe Grandja
  */
 public final class JwkTokenStore implements TokenStore {
-	private final JwtTokenStore delegate;
+	private final TokenStore delegate;
 
 	/**
 	 * Creates a new instance using the provided URL as the location for the JWK Set.
 	 *
 	 * @param jwkSetUrl the JWK Set URL
 	 */
 	public JwkTokenStore(String jwkSetUrl) {
-		Assert.hasText(jwkSetUrl, "jwkSetUrl cannot be empty");
-		JwkDefinitionSource jwkDefinitionSource = new JwkDefinitionSource(jwkSetUrl);
+		this(Arrays.asList(jwkSetUrl));
+	}
+
+	/**
+	 * Creates a new instance using the provided URLs as the location for the JWK Sets.
+	 * @param jwkSetUrls the JWK Set URLs
+	 */
+	public JwkTokenStore(List<String> jwkSetUrls) {
+		JwkDefinitionSource jwkDefinitionSource = new JwkDefinitionSource(jwkSetUrls);
 		JwkVerifyingJwtAccessTokenConverter accessTokenConverter =
 				new JwkVerifyingJwtAccessTokenConverter(jwkDefinitionSource);
 		this.delegate = new JwtTokenStore(accessTokenConverter);

spring-security-oauth2/src/test/java/org/springframework/security/oauth2/provider/token/store/jwk/JwkDefinitionSourceITest.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2012-2017 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.security.oauth2.provider.token.store.jwk;
+
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.MockWebServer;
+import org.apache.http.HttpHeaders;
+import org.junit.Before;
+import org.junit.Test;
+import org.springframework.http.MediaType;
+
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author Rob Winch
+ */
+public class JwkDefinitionSourceITest {
+
+	private MockWebServer server;
+
+	private JwkDefinitionSource source;
+
+	@Before
+	public void setup() {
+		this.server = new MockWebServer();
+	}
+
+	@Test
+	public void getDefinitionLoadIfNecessaryWhenMultipleUrlsThenBothUrlsAreLoaded() {
+		this.server.enqueue(new MockResponse().setHeader(HttpHeaders.ACCEPT, MediaType.APPLICATION_JSON_VALUE).setBody("{\n" +
+				"    \"keys\": [\n" +
+				"        {\n" +
+				"            \"kid\": \"key-id-1\",\n" +
+				"            \"kty\": \"RSA\",\n" +
+				"            \"alg\": \"RS256\",\n" +
+				"            \"use\": \"sig\",\n" +
+				"            \"n\": \"rne3dowbQHcFCzg2ejWb6az5QNxWFiv6kRpd34VDzYNMhWeewfeEL5Pf5clE8Xh1KlllrDYSxtnzUQm-t9p92yEBASfV96ydTYG-ITfxfJzKtJUN-iIS5K9WGYXnDNS4eYZ_ygW-zBU_9NwFMXdwSTzRqHeJmLJrfbmmjoIuuWyfh2Ko52KzyidceR5SJxGeW0ckeyWka1lDf4cr7fv-s093Y_sd2wrNvg0-9IAkXotbxWWXcfMgXFyw0qHFT_5LrKmiwkY3HCaiV5NgEFJmC6fBIG2EOZG4rqjBoYV6LZwrfTMHknaeel9MOZesW6SR2bswtuuWN3DGq2zg0KamLw\",\n" +
+				"            \"e\": \"AQAB\"\n" +
+				"        }\n" +
+				"    ]\n" +
+				"}\n"));
+		this.server.enqueue(new MockResponse().setHeader(HttpHeaders.ACCEPT, MediaType.APPLICATION_JSON_VALUE).setBody("{\n" +
+				"    \"keys\": [\n" +
+				"        {\n" +
+				"            \"kid\": \"key-id-2\",\n" +
+				"            \"kty\": \"RSA\",\n" +
+				"            \"alg\": \"RS256\",\n" +
+				"            \"use\": \"sig\",\n" +
+				"            \"n\": \"t6Q8PWSi1dkJj9hTP8hNYFlvadM7DflW9mWepOJhJ66w7nyoK1gPNqFMSQRyO125Gp-TEkodhWr0iujjHVx7BcV0llS4w5ACGgPrcAd6ZcSR0-Iqom-QFcNP8Sjg086MwoqQU_LYywlAGZ21WSdS_PERyGFiNnj3QQlO8Yns5jCtLCRwLHL0Pb1fEv45AuRIuUfVcPySBWYnDyGxvjYGDSM-AqWS9zIQ2ZilgT-GqUmipg0XOC0Cc20rgLe2ymLHjpHciCKVAbY5-L32-lSeZO-Os6U15_aXrk9Gw8cPUaX1_I8sLGuSiVdt3C_Fn2PZ3Z8i744FPFGGcG1qs2Wz-Q\",\n" +
+				"            \"e\": \"AQAB\"\n" +
+				"        }\n" +
+				"    ]\n" +
+				"}\n"));
+		this.source = new JwkDefinitionSource(Arrays.asList(serverUrl("/jwk1"), serverUrl("/jkw2")));
+
+		String keyId1 = "key-id-1";
+		String keyId2 = "key-id-2";
+		JwkDefinition jwkDef1 = this.source.getDefinitionLoadIfNecessary(keyId1);
+		JwkDefinition jwkDef2 = this.source.getDefinitionLoadIfNecessary(keyId2);
+
+		assertEquals(jwkDef1.getKeyId(), keyId1);
+		assertEquals(jwkDef1.getAlgorithm(), JwkDefinition.CryptoAlgorithm.RS256);
+		assertEquals(jwkDef1.getPublicKeyUse(), JwkDefinition.PublicKeyUse.SIG);
+		assertEquals(jwkDef1.getKeyType(), JwkDefinition.KeyType.RSA);
+
+		assertEquals(jwkDef2.getKeyId(), keyId2);
+		assertEquals(jwkDef2.getAlgorithm(), JwkDefinition.CryptoAlgorithm.RS256);
+		assertEquals(jwkDef2.getPublicKeyUse(), JwkDefinition.PublicKeyUse.SIG);
+		assertEquals(jwkDef2.getKeyType(), JwkDefinition.KeyType.RSA);
+	}
+
+	private String serverUrl(String path) {
+		return this.server.url(path).toString();
+	}
+}

spring-security-oauth2/src/test/java/org/springframework/security/oauth2/provider/token/store/jwk/JwkDefinitionSourceTest.java

@@ -26,6 +26,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
+import java.util.Arrays;
 import java.util.Collections;
 
 import static org.mockito.Matchers.any;
@@ -46,6 +47,11 @@ public void constructorWhenInvalidJwkSetUrlThenThrowIllegalArgumentException() t
 		new JwkDefinitionSource(DEFAULT_JWK_SET_URL.substring(1));
 	}
 
+	@Test(expected = IllegalArgumentException.class)
+	public void constructorListWhenInvalidJwkSetUrlThenThrowIllegalArgumentException() throws Exception {
+		new JwkDefinitionSource(Arrays.asList(DEFAULT_JWK_SET_URL.substring(1)));
+	}
+
 	@Test
 	public void getDefinitionLoadIfNecessaryWhenKeyIdNotFoundThenLoadJwkDefinitions() throws Exception {
 		JwkDefinitionSource jwkDefinitionSource = spy(new JwkDefinitionSource(DEFAULT_JWK_SET_URL));