Update RefreshOidcIdTokenHandler to improve decoupling

Signed-off-by: Hao <[email protected]>

Author: yhao3

config/src/main/java/org/springframework/security/config/annotation/web/configurers/oauth2/client/OAuth2LoginConfigurer.java

@@ -395,8 +395,13 @@ public void init(B http) throws Exception {
 			}
 			http.authenticationProvider(this.postProcess(oidcAuthorizationCodeAuthenticationProvider));
 
-			RefreshOidcIdTokenHandler refreshOidcIdTokenHandler = new RefreshOidcIdTokenHandler(
-					oidcAuthorizationCodeAuthenticationProvider);
+			RefreshOidcIdTokenHandler refreshOidcIdTokenHandler = new RefreshOidcIdTokenHandler();
+			if (this.getSecurityContextHolderStrategy() != null) {
+				refreshOidcIdTokenHandler.setSecurityContextHolderStrategy(this.getSecurityContextHolderStrategy());
+			}
+			if (jwtDecoderFactory != null) {
+				refreshOidcIdTokenHandler.setJwtDecoderFactory(jwtDecoderFactory);
+			}
 			registerDelegateApplicationListener(refreshOidcIdTokenHandler);
 		}
 		else {

oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/oidc/authentication/OidcAuthorizationCodeAuthenticationProvider.java

@@ -232,7 +232,7 @@ public boolean supports(Class<?> authentication) {
 		return OAuth2LoginAuthenticationToken.class.isAssignableFrom(authentication);
 	}
 
-	protected OidcIdToken createOidcToken(ClientRegistration clientRegistration,
+	private OidcIdToken createOidcToken(ClientRegistration clientRegistration,
 			OAuth2AccessTokenResponse accessTokenResponse) {
 		JwtDecoder jwtDecoder = this.jwtDecoderFactory.createDecoder(clientRegistration);
 		Jwt jwt = getJwt(accessTokenResponse, jwtDecoder);

oauth2/oauth2-client/src/main/java/org/springframework/security/oauth2/client/oidc/authentication/RefreshOidcIdTokenHandler.java

@@ -16,46 +16,98 @@
 
 package org.springframework.security.oauth2.client.oidc.authentication;
 
+import java.util.Map;
+
 import org.springframework.context.ApplicationListener;
 import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolderStrategy;
 import org.springframework.security.oauth2.client.OAuth2AuthorizedClient;
 import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
 import org.springframework.security.oauth2.client.event.OAuth2TokenRefreshedEvent;
+import org.springframework.security.oauth2.client.registration.ClientRegistration;
+import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
+import org.springframework.security.oauth2.core.OAuth2Error;
 import org.springframework.security.oauth2.core.endpoint.OAuth2AccessTokenResponse;
 import org.springframework.security.oauth2.core.oidc.OidcIdToken;
 import org.springframework.security.oauth2.core.oidc.StandardClaimNames;
+import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
 import org.springframework.security.oauth2.core.oidc.user.DefaultOidcUser;
 import org.springframework.security.oauth2.core.oidc.user.OidcUser;
+import org.springframework.security.oauth2.jwt.Jwt;
+import org.springframework.security.oauth2.jwt.JwtDecoder;
+import org.springframework.security.oauth2.jwt.JwtDecoderFactory;
+import org.springframework.security.oauth2.jwt.JwtException;
+import org.springframework.util.Assert;
 
 /**
  * An {@link ApplicationListener} that listens for {@link OAuth2TokenRefreshedEvent}s
  */
 public class RefreshOidcIdTokenHandler implements ApplicationListener<OAuth2TokenRefreshedEvent> {
 
-	private final OidcAuthorizationCodeAuthenticationProvider oidcAuthorizationCodeAuthenticationProvider;
+	private static final String INVALID_ID_TOKEN_ERROR_CODE = "invalid_id_token";
 
-	public RefreshOidcIdTokenHandler(
-			OidcAuthorizationCodeAuthenticationProvider oidcAuthorizationCodeAuthenticationProvider) {
-		this.oidcAuthorizationCodeAuthenticationProvider = oidcAuthorizationCodeAuthenticationProvider;
-	}
+	private SecurityContextHolderStrategy securityContextHolderStrategy = SecurityContextHolder
+		.getContextHolderStrategy();
+
+	private JwtDecoderFactory<ClientRegistration> jwtDecoderFactory = new OidcIdTokenDecoderFactory();
 
 	@Override
 	public void onApplicationEvent(OAuth2TokenRefreshedEvent event) {
 		OAuth2AuthorizedClient authorizedClient = event.getAuthorizedClient();
 		OAuth2AccessTokenResponse accessTokenResponse = event.getAccessTokenResponse();
-		OidcIdToken refreshedOidcToken = this.oidcAuthorizationCodeAuthenticationProvider
-			.createOidcToken(authorizedClient.getClientRegistration(), accessTokenResponse);
-		Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
+		ClientRegistration clientRegistration = authorizedClient.getClientRegistration();
+		OidcIdToken refreshedOidcToken = createOidcToken(clientRegistration, accessTokenResponse);
+		Authentication authentication = this.securityContextHolderStrategy.getContext().getAuthentication();
 		if (authentication instanceof OAuth2AuthenticationToken oauth2AuthenticationToken) {
 			if (authentication.getPrincipal() instanceof DefaultOidcUser defaultOidcUser) {
 				OidcUser oidcUser = new DefaultOidcUser(defaultOidcUser.getAuthorities(), refreshedOidcToken,
 						defaultOidcUser.getUserInfo(), StandardClaimNames.SUB);
-				SecurityContextHolder.getContext()
-					.setAuthentication(new OAuth2AuthenticationToken(oidcUser, oidcUser.getAuthorities(),
-							oauth2AuthenticationToken.getAuthorizedClientRegistrationId()));
+				SecurityContext context = this.securityContextHolderStrategy.createEmptyContext();
+				context.setAuthentication(new OAuth2AuthenticationToken(oidcUser, oidcUser.getAuthorities(),
+						oauth2AuthenticationToken.getAuthorizedClientRegistrationId()));
+				this.securityContextHolderStrategy.setContext(context);
 			}
 		}
 	}
 
+	/**
+	 * Sets the {@link SecurityContextHolderStrategy} to use. The default action is to use
+	 * the {@link SecurityContextHolderStrategy} stored in {@link SecurityContextHolder}.
+	 */
+	public void setSecurityContextHolderStrategy(SecurityContextHolderStrategy securityContextHolderStrategy) {
+		this.securityContextHolderStrategy = securityContextHolderStrategy;
+	}
+
+	/**
+	 * Sets the {@link JwtDecoderFactory} used for {@link OidcIdToken} signature
+	 * verification. The factory returns a {@link JwtDecoder} associated to the provided
+	 * {@link ClientRegistration}.
+	 * @param jwtDecoderFactory the {@link JwtDecoderFactory} used for {@link OidcIdToken}
+	 * signature verification
+	 */
+	public final void setJwtDecoderFactory(JwtDecoderFactory<ClientRegistration> jwtDecoderFactory) {
+		Assert.notNull(jwtDecoderFactory, "jwtDecoderFactory cannot be null");
+		this.jwtDecoderFactory = jwtDecoderFactory;
+	}
+
+	private OidcIdToken createOidcToken(ClientRegistration clientRegistration,
+			OAuth2AccessTokenResponse accessTokenResponse) {
+		JwtDecoder jwtDecoder = this.jwtDecoderFactory.createDecoder(clientRegistration);
+		Jwt jwt = getJwt(accessTokenResponse, jwtDecoder);
+		return new OidcIdToken(jwt.getTokenValue(), jwt.getIssuedAt(), jwt.getExpiresAt(), jwt.getClaims());
+	}
+
+	private Jwt getJwt(OAuth2AccessTokenResponse accessTokenResponse, JwtDecoder jwtDecoder) {
+		try {
+			Map<String, Object> parameters = accessTokenResponse.getAdditionalParameters();
+			return jwtDecoder.decode((String) parameters.get(OidcParameterNames.ID_TOKEN));
+		}
+		catch (JwtException ex) {
+			OAuth2Error invalidIdTokenError = new OAuth2Error(INVALID_ID_TOKEN_ERROR_CODE, ex.getMessage(), null);
+			throw new OAuth2AuthenticationException(invalidIdTokenError, invalidIdTokenError.toString(), ex);
+		}
+	}
+
 }