From f5da63118e36287a7cefa9adb5da83b6617608b0 Mon Sep 17 00:00:00 2001
From: Josh Cummings <josh.cummings@gmail.com>
Date: Tue, 11 Jun 2019 21:44:51 -0600
Subject: [PATCH] Add MultiTenantAuthenticationManagerResolver

A class with a number of handy request-based implementations of
AuthenticationManagerResolver targeted at common multi-tenancy
scenarios.

Fixes: gh-6976
---
 ...h2ResourceServerSecurityConfiguration.java |  15 +-
 ...tiTenantAuthenticationManagerResolver.java | 174 ++++++++++++++++++
 ...antAuthenticationManagerResolverTests.java | 159 ++++++++++++++++
 3 files changed, 338 insertions(+), 10 deletions(-)
 create mode 100644 web/src/main/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolver.java
 create mode 100644 web/src/test/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolverTests.java

diff --git a/samples/boot/oauth2resourceserver-multitenancy/src/main/java/sample/OAuth2ResourceServerSecurityConfiguration.java b/samples/boot/oauth2resourceserver-multitenancy/src/main/java/sample/OAuth2ResourceServerSecurityConfiguration.java
index 6983cbf355..b8bd55b467 100644
--- a/samples/boot/oauth2resourceserver-multitenancy/src/main/java/sample/OAuth2ResourceServerSecurityConfiguration.java
+++ b/samples/boot/oauth2resourceserver-multitenancy/src/main/java/sample/OAuth2ResourceServerSecurityConfiguration.java
@@ -17,7 +17,6 @@ package sample;
 
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Optional;
 import javax.servlet.http.HttpServletRequest;
 
 import org.springframework.beans.factory.annotation.Value;
@@ -27,12 +26,14 @@ import org.springframework.security.authentication.AuthenticationManagerResolver
 import org.springframework.security.config.annotation.web.builders.HttpSecurity;
 import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity;
 import org.springframework.security.config.annotation.web.configuration.WebSecurityConfigurerAdapter;
-import org.springframework.security.oauth2.server.resource.introspection.NimbusOAuth2TokenIntrospectionClient;
-import org.springframework.security.oauth2.server.resource.introspection.OAuth2TokenIntrospectionClient;
 import org.springframework.security.oauth2.jwt.JwtDecoder;
 import org.springframework.security.oauth2.jwt.NimbusJwtDecoder;
 import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationProvider;
 import org.springframework.security.oauth2.server.resource.authentication.OAuth2IntrospectionAuthenticationProvider;
+import org.springframework.security.oauth2.server.resource.introspection.NimbusOAuth2TokenIntrospectionClient;
+import org.springframework.security.oauth2.server.resource.introspection.OAuth2TokenIntrospectionClient;
+
+import static org.springframework.security.web.authentication.MultiTenantAuthenticationManagerResolver.resolveFromPath;
 
 /**
  * @author Josh Cummings
@@ -64,13 +65,7 @@ public class OAuth2ResourceServerSecurityConfiguration extends WebSecurityConfig
 		Map<String, AuthenticationManager> authenticationManagers = new HashMap<>();
 		authenticationManagers.put("tenantOne", jwt());
 		authenticationManagers.put("tenantTwo", opaque());
-		return request -> {
-			String[] pathParts = request.getRequestURI().split("/");
-			String tenantId = pathParts.length > 0 ? pathParts[1] : null;
-			return Optional.ofNullable(tenantId)
-					.map(authenticationManagers::get)
-					.orElseThrow(() -> new IllegalArgumentException("unknown tenant"));
-		};
+		return resolveFromPath(authenticationManagers::get);
 	}
 
 	AuthenticationManager jwt() {
diff --git a/web/src/main/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolver.java b/web/src/main/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolver.java
new file mode 100644
index 0000000000..4056d35b9f
--- /dev/null
+++ b/web/src/main/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolver.java
@@ -0,0 +1,174 @@
+/*
+ * Copyright 2002-2019 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
+ *
+ *      https://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.web.authentication;
+
+import java.util.Optional;
+import javax.servlet.http.HttpServletRequest;
+
+import org.springframework.core.convert.converter.Converter;
+import org.springframework.security.authentication.AuthenticationManager;
+import org.springframework.security.authentication.AuthenticationManagerResolver;
+import org.springframework.util.Assert;
+import org.springframework.web.util.UriComponents;
+import org.springframework.web.util.UriComponentsBuilder;
+
+/**
+ * An implementation of {@link AuthenticationManagerResolver} that separates the tasks of
+ * extracting the request's tenant identifier and looking up an {@link AuthenticationManager}
+ * by that tenant identifier.
+ *
+ * @author Josh Cummings
+ * @since 5.2
+ * @see AuthenticationManagerResolver
+ */
+public final class MultiTenantAuthenticationManagerResolver<T> implements AuthenticationManagerResolver<HttpServletRequest> {
+
+	private final Converter<HttpServletRequest, AuthenticationManager> authenticationManagerResolver;
+
+	/**
+	 * Constructs a {@link MultiTenantAuthenticationManagerResolver} with the provided parameters
+	 *
+	 * @param tenantResolver
+	 * @param authenticationManagerResolver
+	 */
+	public MultiTenantAuthenticationManagerResolver
+			(Converter<HttpServletRequest, T> tenantResolver,
+					Converter<T, AuthenticationManager> authenticationManagerResolver) {
+
+		Assert.notNull(tenantResolver, "tenantResolver cannot be null");
+		Assert.notNull(authenticationManagerResolver, "authenticationManagerResolver cannot be null");
+
+		this.authenticationManagerResolver = request -> {
+			Optional<T> context = Optional.ofNullable(tenantResolver.convert(request));
+			return context.map(authenticationManagerResolver::convert)
+					.orElseThrow(() -> new IllegalArgumentException
+							("Could not resolve AuthenticationManager by reference " + context.orElse(null)));
+		};
+	}
+
+	@Override
+	public AuthenticationManager resolve(HttpServletRequest context) {
+		return this.authenticationManagerResolver.convert(context);
+	}
+
+	/**
+	 * Creates an {@link AuthenticationManagerResolver} that will use a hostname's first label as
+	 * the resolution key for the underlying {@link AuthenticationManagerResolver}.
+	 *
+	 * For example, you might have a set of {@link AuthenticationManager}s defined like so:
+	 *
+	 * <pre>
+	 * 	Map<String, AuthenticationManager> authenticationManagers = new HashMap<>();
+	 *  authenticationManagers.put("tenantOne", managerOne());
+	 *  authenticationManagers.put("tenantTwo", managerTwo());
+	 * </pre>
+	 *
+	 * And that your system serves hostnames like <pre>https://tenantOne.example.org</pre>.
+	 *
+	 * Then, you could create an {@link AuthenticationManagerResolver} that uses the "tenantOne" value from
+	 * the hostname to resolve Tenant One's {@link AuthenticationManager} like so:
+	 *
+	 * <pre>
+	 *	AuthenticationManagerResolver<HttpServletRequest> resolver =
+	 *			resolveFromSubdomain(authenticationManagers::get);
+	 * </pre>
+	 *
+	 * {@link HttpServletRequest}
+	 * @param resolver A {@link String}-resolving {@link AuthenticationManagerResolver}
+	 * @return A hostname-resolving {@link AuthenticationManagerResolver}
+	 */
+	public static AuthenticationManagerResolver<HttpServletRequest>
+			resolveFromSubdomain(Converter<String, AuthenticationManager> resolver) {
+
+		return new MultiTenantAuthenticationManagerResolver<>(request ->
+				Optional.ofNullable(request.getServerName())
+						.map(host -> host.split("\\."))
+						.filter(segments -> segments.length > 0)
+						.map(segments -> segments[0]).orElse(null), resolver);
+	}
+
+	/**
+	 * Creates an {@link AuthenticationManagerResolver} that will use a request path's first segment as
+	 * the resolution key for the underlying {@link AuthenticationManagerResolver}.
+	 *
+	 * For example, you might have a set of {@link AuthenticationManager}s defined like so:
+	 *
+	 * <pre>
+	 * 	Map<String, AuthenticationManager> authenticationManagers = new HashMap<>();
+	 *  authenticationManagers.put("tenantOne", managerOne());
+	 *  authenticationManagers.put("tenantTwo", managerTwo());
+	 * </pre>
+	 *
+	 * And that your system serves requests like <pre>https://example.org/tenantOne</pre>.
+	 *
+	 * Then, you could create an {@link AuthenticationManagerResolver} that uses the "tenantOne" value from
+	 * the request to resolve Tenant One's {@link AuthenticationManager} like so:
+	 *
+	 * <pre>
+	 *	AuthenticationManagerResolver<HttpServletRequest> resolver =
+	 *			resolveFromPath(authenticationManagers::get);
+	 * </pre>
+	 *
+	 * {@link HttpServletRequest}
+	 * @param resolver A {@link String}-resolving {@link AuthenticationManagerResolver}
+	 * @return A path-resolving {@link AuthenticationManagerResolver}
+	 */
+	public static AuthenticationManagerResolver<HttpServletRequest>
+			resolveFromPath(Converter<String, AuthenticationManager> resolver) {
+
+		return new MultiTenantAuthenticationManagerResolver<>(request ->
+			Optional.ofNullable(request.getRequestURI())
+					.map(UriComponentsBuilder::fromUriString)
+					.map(UriComponentsBuilder::build)
+					.map(UriComponents::getPathSegments)
+					.filter(segments -> !segments.isEmpty())
+					.map(segments -> segments.get(0)).orElse(null), resolver);
+	}
+
+	/**
+	 * Creates an {@link AuthenticationManagerResolver} that will use a request headers's value as
+	 * the resolution key for the underlying {@link AuthenticationManagerResolver}.
+	 *
+	 * For example, you might have a set of {@link AuthenticationManager}s defined like so:
+	 *
+	 * <pre>
+	 * 	Map<String, AuthenticationManager> authenticationManagers = new HashMap<>();
+	 *  authenticationManagers.put("tenantOne", managerOne());
+	 *  authenticationManagers.put("tenantTwo", managerTwo());
+	 * </pre>
+	 *
+	 * And that your system serves requests with a header like <pre>X-Tenant-Id: tenantOne</pre>.
+	 *
+	 * Then, you could create an {@link AuthenticationManagerResolver} that uses the "tenantOne" value from
+	 * the request to resolve Tenant One's {@link AuthenticationManager} like so:
+	 *
+	 * <pre>
+	 *	AuthenticationManagerResolver<HttpServletRequest> resolver =
+	 *			resolveFromHeader("X-Tenant-Id", authenticationManagers::get);
+	 * </pre>
+	 *
+	 * {@link HttpServletRequest}
+	 * @param resolver A {@link String}-resolving {@link AuthenticationManagerResolver}
+	 * @return A header-resolving {@link AuthenticationManagerResolver}
+	 */
+	public static AuthenticationManagerResolver<HttpServletRequest>
+			resolveFromHeader(String headerName, Converter<String, AuthenticationManager> resolver) {
+
+		return new MultiTenantAuthenticationManagerResolver<>
+				(request -> request.getHeader(headerName), resolver);
+	}
+}
diff --git a/web/src/test/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolverTests.java b/web/src/test/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolverTests.java
new file mode 100644
index 0000000000..902e6b72be
--- /dev/null
+++ b/web/src/test/java/org/springframework/security/web/authentication/MultiTenantAuthenticationManagerResolverTests.java
@@ -0,0 +1,159 @@
+/*
+ * Copyright 2002-2019 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
+ *
+ *      https://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.web.authentication;
+
+import java.util.Collections;
+import java.util.Map;
+import javax.servlet.http.HttpServletRequest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import org.springframework.core.convert.converter.Converter;
+import org.springframework.security.authentication.AuthenticationManager;
+import org.springframework.security.authentication.AuthenticationManagerResolver;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.springframework.security.web.authentication.MultiTenantAuthenticationManagerResolver.resolveFromSubdomain;
+import static org.springframework.security.web.authentication.MultiTenantAuthenticationManagerResolver.resolveFromPath;
+import static org.springframework.security.web.authentication.MultiTenantAuthenticationManagerResolver.resolveFromHeader;
+
+/**
+ * Tests for {@link MultiTenantAuthenticationManagerResolver}
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class MultiTenantAuthenticationManagerResolverTests {
+	private static final String TENANT = "tenant";
+
+	@Mock
+	AuthenticationManager authenticationManager;
+
+	@Mock
+	HttpServletRequest request;
+
+	Map<String, AuthenticationManager> authenticationManagers;
+
+	@Before
+	public void setup() {
+		this.authenticationManagers = Collections.singletonMap(TENANT, this.authenticationManager);
+	}
+
+	@Test
+	public void resolveFromSubdomainWhenGivenResolverThenReturnsSubdomainParsingResolver() {
+		AuthenticationManagerResolver<HttpServletRequest> fromSubdomain =
+				resolveFromSubdomain(this.authenticationManagers::get);
+
+		when(this.request.getServerName()).thenReturn(TENANT + ".example.org");
+
+		AuthenticationManager authenticationManager = fromSubdomain.resolve(this.request);
+		assertThat(authenticationManager).isEqualTo(this.authenticationManager);
+
+		when(this.request.getServerName()).thenReturn("wrong.example.org");
+
+		assertThatCode(() -> fromSubdomain.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+
+		when(this.request.getServerName()).thenReturn("example");
+
+		assertThatCode(() -> fromSubdomain.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	public void resolveFromPathWhenGivenResolverThenReturnsPathParsingResolver() {
+		AuthenticationManagerResolver<HttpServletRequest> fromPath =
+				resolveFromPath(this.authenticationManagers::get);
+
+		when(this.request.getRequestURI()).thenReturn("/" + TENANT + "/otherthings");
+
+		AuthenticationManager authenticationManager = fromPath.resolve(this.request);
+		assertThat(authenticationManager).isEqualTo(this.authenticationManager);
+
+		when(this.request.getRequestURI()).thenReturn("/otherthings");
+
+		assertThatCode(() -> fromPath.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+
+		when(this.request.getRequestURI()).thenReturn("/");
+
+		assertThatCode(() -> fromPath.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	public void resolveFromHeaderWhenGivenResolverTheReturnsHeaderParsingResolver() {
+		AuthenticationManagerResolver<HttpServletRequest> fromHeader =
+				resolveFromHeader("X-Tenant-Id", this.authenticationManagers::get);
+
+		when(this.request.getHeader("X-Tenant-Id")).thenReturn(TENANT);
+
+		AuthenticationManager authenticationManager = fromHeader.resolve(this.request);
+		assertThat(authenticationManager).isEqualTo(this.authenticationManager);
+
+		when(this.request.getHeader("X-Tenant-Id")).thenReturn("wrong");
+
+		assertThatCode(() -> fromHeader.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+
+		when(this.request.getHeader("X-Tenant-Id")).thenReturn(null);
+
+		assertThatCode(() -> fromHeader.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	public void resolveWhenGivenTenantResolverThenResolves() {
+		AuthenticationManagerResolver<HttpServletRequest> byRequestConverter =
+				new MultiTenantAuthenticationManagerResolver<>(HttpServletRequest::getQueryString,
+						this.authenticationManagers::get);
+
+		when(this.request.getQueryString()).thenReturn(TENANT);
+
+		AuthenticationManager authenticationManager = byRequestConverter.resolve(this.request);
+		assertThat(authenticationManager).isEqualTo(this.authenticationManager);
+
+		when(this.request.getQueryString()).thenReturn("wrong");
+
+		assertThatCode(() -> byRequestConverter.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+
+		when(this.request.getQueryString()).thenReturn(null);
+
+		assertThatCode(() -> byRequestConverter.resolve(this.request))
+				.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	public void constructorWhenUsingNullTenantResolverThenException() {
+		assertThatCode(() -> new MultiTenantAuthenticationManagerResolver
+				(null, mock(Converter.class)))
+				.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	public void constructorWhenUsingNullAuthenticationManagerResolverThenException() {
+		assertThatCode(() -> new MultiTenantAuthenticationManagerResolver
+				(mock(Converter.class), null))
+				.isInstanceOf(IllegalArgumentException.class);
+	}
+}