From d0a4e49870e96bd5aa24b8212cb1c26a99d44577 Mon Sep 17 00:00:00 2001
From: Joe Grandja <jgrandja@pivotal.io>
Date: Wed, 25 Oct 2017 17:13:44 -0400
Subject: [PATCH] Map custom OAuth2User types using String

Fixes gh-4691
---
 .../configurers/oauth2/client/OAuth2LoginConfigurer.java | 6 +++---
 .../userinfo/CustomUserTypesOAuth2UserService.java       | 9 ++++-----
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/config/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/client/OAuth2LoginConfigurer.java b/config/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/client/OAuth2LoginConfigurer.java
index 2f0ca5c0fd..2d660dd1d0 100644
--- a/config/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/client/OAuth2LoginConfigurer.java
+++ b/config/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/client/OAuth2LoginConfigurer.java
@@ -190,7 +190,7 @@ public final class OAuth2LoginConfigurer<B extends HttpSecurityBuilder<B>> exten
 
 	public class UserInfoEndpointConfig {
 		private OAuth2UserService userService;
-		private Map<URI, Class<? extends OAuth2User>> customUserTypes = new HashMap<>();
+		private Map<String, Class<? extends OAuth2User>> customUserTypes = new HashMap<>();
 		private GrantedAuthoritiesMapper userAuthoritiesMapper;
 
 		private UserInfoEndpointConfig() {
@@ -202,9 +202,9 @@ public final class OAuth2LoginConfigurer<B extends HttpSecurityBuilder<B>> exten
 			return this;
 		}
 
-		public UserInfoEndpointConfig customUserType(Class<? extends OAuth2User> customUserType, URI userInfoUri) {
+		public UserInfoEndpointConfig customUserType(Class<? extends OAuth2User> customUserType, String userInfoUri) {
 			Assert.notNull(customUserType, "customUserType cannot be null");
-			Assert.notNull(userInfoUri, "userInfoUri cannot be null");
+			Assert.hasText(userInfoUri, "userInfoUri cannot be empty");
 			this.customUserTypes.put(userInfoUri, customUserType);
 			return this;
 		}
diff --git a/oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/authentication/userinfo/CustomUserTypesOAuth2UserService.java b/oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/authentication/userinfo/CustomUserTypesOAuth2UserService.java
index c23c041a4d..ff9bd21316 100644
--- a/oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/authentication/userinfo/CustomUserTypesOAuth2UserService.java
+++ b/oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/authentication/userinfo/CustomUserTypesOAuth2UserService.java
@@ -20,7 +20,6 @@ import org.springframework.security.oauth2.client.authentication.OAuth2Authentic
 import org.springframework.security.oauth2.core.user.OAuth2User;
 import org.springframework.util.Assert;
 
-import java.net.URI;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -29,7 +28,7 @@ import java.util.Map;
  * An implementation of an {@link OAuth2UserService} that supports custom {@link OAuth2User} types.
  * <p>
  * The custom user type(s) is supplied via the constructor,
- * using a <code>Map</code> of {@link OAuth2User} type <i>keyed</i> by <code>URI</code>,
+ * using a <code>Map</code> of {@link OAuth2User} type <i>keyed</i> by <code>String</code>,
  * representing the <i>UserInfo Endpoint</i> address.
  * <p>
  * This implementation uses a {@link UserInfoRetriever} to obtain the user attributes
@@ -42,17 +41,17 @@ import java.util.Map;
  * @see UserInfoRetriever
  */
 public class CustomUserTypesOAuth2UserService implements OAuth2UserService {
-	private final Map<URI, Class<? extends OAuth2User>> customUserTypes;
+	private final Map<String, Class<? extends OAuth2User>> customUserTypes;
 	private UserInfoRetriever userInfoRetriever = new NimbusUserInfoRetriever();
 
-	public CustomUserTypesOAuth2UserService(Map<URI, Class<? extends OAuth2User>> customUserTypes) {
+	public CustomUserTypesOAuth2UserService(Map<String, Class<? extends OAuth2User>> customUserTypes) {
 		Assert.notEmpty(customUserTypes, "customUserTypes cannot be empty");
 		this.customUserTypes = Collections.unmodifiableMap(new LinkedHashMap<>(customUserTypes));
 	}
 
 	@Override
 	public OAuth2User loadUser(AuthorizedClient authorizedClient) throws OAuth2AuthenticationException {
-		URI userInfoUri = URI.create(authorizedClient.getClientRegistration().getProviderDetails().getUserInfoEndpoint().getUri());
+		String userInfoUri = authorizedClient.getClientRegistration().getProviderDetails().getUserInfoEndpoint().getUri();
 		Class<? extends OAuth2User> customUserType;
 		if ((customUserType = this.customUserTypes.get(userInfoUri)) == null) {
 			return null;