Specify clientRegistrationId in TokenRelayFilterFunctions

Closes spring-cloudgh-3541
sjohnr · Nov 12, 2024 · 6c3dff4 · 6c3dff4
1 parent 031e249
commit 6c3dff4
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/...n/java/org/springframework/cloud/gateway/server/mvc/filter/TokenRelayFilterFunctions.java b/...n/java/org/springframework/cloud/gateway/server/mvc/filter/TokenRelayFilterFunctions.java
@@ -19,6 +19,7 @@
 import java.security.Principal;
 
 import org.springframework.cloud.gateway.server.mvc.common.Shortcut;
+import org.springframework.security.core.Authentication;
 import org.springframework.security.oauth2.client.OAuth2AuthorizeRequest;
 import org.springframework.security.oauth2.client.OAuth2AuthorizedClient;
 import org.springframework.security.oauth2.client.OAuth2AuthorizedClientManager;
@@ -37,13 +38,21 @@ private TokenRelayFilterFunctions() {
 
 	@Shortcut
 	public static HandlerFilterFunction<ServerResponse, ServerResponse> tokenRelay() {
+		return tokenRelay(null);
+	}
+
+	public static HandlerFilterFunction<ServerResponse, ServerResponse> tokenRelay(String defaultClientRegistrationId) {
 		return (request, next) -> {
-			Principal principle = request.servletRequest().getUserPrincipal();
-			if (principle instanceof OAuth2AuthenticationToken token) {
-				String clientRegistrationId = token.getAuthorizedClientRegistrationId();
+			Authentication principal = (Authentication) request.servletRequest().getUserPrincipal();
+
+			String clientRegistrationId = defaultClientRegistrationId;
+			if (clientRegistrationId == null && principal instanceof OAuth2AuthenticationToken token) {
+				clientRegistrationId = token.getAuthorizedClientRegistrationId();
+			}
+			if (clientRegistrationId != null) {
 				OAuth2AuthorizeRequest authorizeRequest = OAuth2AuthorizeRequest
 					.withClientRegistrationId(clientRegistrationId)
-					.principal(token)
+					.principal(principal)
 					.build();
 				OAuth2AuthorizedClientManager clientManager = getApplicationContext(request)
 					.getBean(OAuth2AuthorizedClientManager.class);

diff --git a/...est/java/org/springframework/cloud/gateway/server/mvc/TokenRelayFilterFunctionsTests.java b/...est/java/org/springframework/cloud/gateway/server/mvc/TokenRelayFilterFunctionsTests.java
@@ -104,6 +104,29 @@ public void whenPrincipalExistsAuthorizationHeaderAdded() throws Exception {
 		});
 	}
 
+	@Test
+	public void whenDefaultClientRegistrationIdProvidedAuthorizationHeaderAdded() throws Exception {
+		OAuth2AccessToken accessToken = mock(OAuth2AccessToken.class);
+		when(accessToken.getTokenValue()).thenReturn("mytoken");
+
+		ClientRegistration clientRegistration = ClientRegistration.withRegistrationId("myregistrationid")
+			.authorizationGrantType(AuthorizationGrantType.CLIENT_CREDENTIALS)
+			.clientId("myclientid")
+			.tokenUri("mytokenuri")
+			.build();
+		OAuth2AuthorizedClient authorizedClient = new OAuth2AuthorizedClient(clientRegistration, "joe", accessToken);
+
+		when(authorizedClientManager.authorize(any(OAuth2AuthorizeRequest.class))).thenReturn(authorizedClient);
+
+		request.setUserPrincipal(new TestingAuthenticationToken("my", null));
+
+		filter = TokenRelayFilterFunctions.tokenRelay("myId");
+		filter.filter(ServerRequest.create(request, converters), req -> {
+			assertThat(req.headers().firstHeader(HttpHeaders.AUTHORIZATION)).isEqualTo("Bearer mytoken");
+			return null;
+		});
+	}
+
 	@Test
 	public void principalIsNotOAuth2AuthenticationToken() throws Exception {
 		request.setUserPrincipal(new TestingAuthenticationToken("my", null));