Support Custom Principal in Jwt Authentication Flow

Closes gh-6237

Signed-off-by: Andrey Litvitski <[email protected]>

Author: therepanic

oauth2/oauth2-resource-server/src/main/java/org/springframework/security/oauth2/server/resource/authentication/JwtAuthenticationConverter.java

@@ -21,6 +21,7 @@
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.security.authentication.AbstractAuthenticationToken;
 import org.springframework.security.core.GrantedAuthority;
+import org.springframework.security.oauth2.core.OAuth2AuthenticatedPrincipal;
 import org.springframework.security.oauth2.jwt.Jwt;
 import org.springframework.security.oauth2.jwt.JwtClaimNames;
 import org.springframework.util.Assert;
@@ -30,10 +31,12 @@
  * @author Josh Cummings
  * @author Evgeniy Cheban
  * @author Olivier Antoine
+ * @author Andrey Litvitski
  * @since 5.1
  */
 public class JwtAuthenticationConverter implements Converter<Jwt, AbstractAuthenticationToken> {
 
+	private Converter<Jwt, OAuth2AuthenticatedPrincipal> jwtPrincipalConverter;
 	private Converter<Jwt, Collection<GrantedAuthority>> jwtGrantedAuthoritiesConverter = new JwtGrantedAuthoritiesConverter();
 
 	private String principalClaimName = JwtClaimNames.SUB;
@@ -42,8 +45,26 @@ public class JwtAuthenticationConverter implements Converter<Jwt, AbstractAuthen
 	public final AbstractAuthenticationToken convert(Jwt jwt) {
 		Collection<GrantedAuthority> authorities = this.jwtGrantedAuthoritiesConverter.convert(jwt);
 
-		String principalClaimValue = jwt.getClaimAsString(this.principalClaimName);
-		return new JwtAuthenticationToken(jwt, authorities, principalClaimValue);
+		if (this.jwtPrincipalConverter == null) {
+			String principalClaimValue = jwt.getClaimAsString(this.principalClaimName);
+			return new JwtAuthenticationToken(jwt, authorities, principalClaimValue);
+		} else {
+			OAuth2AuthenticatedPrincipal principal = this.jwtPrincipalConverter.convert(jwt);
+			authorities.addAll(principal.getAuthorities());
+			return new JwtAuthenticationToken(jwt, principal, authorities);
+		}
+	}
+
+	/**
+	 * Sets the {@link Converter Converter&lt;Jwt, Collection&lt;OAuth2AuthenticatedPrincipal&gt;&gt;}
+	 * to use.
+	 * @param jwtPrincipalConverter The converter
+	 * @since 6.5.0
+	 */
+	public void setJwtPrincipalConverter(
+			Converter<Jwt, OAuth2AuthenticatedPrincipal> jwtPrincipalConverter) {
+		Assert.notNull(jwtPrincipalConverter, "jwtPrincipalConverter cannot be null");
+		this.jwtPrincipalConverter = jwtPrincipalConverter;
 	}
 
 	/**

oauth2/oauth2-resource-server/src/main/java/org/springframework/security/oauth2/server/resource/authentication/JwtAuthenticationToken.java

@@ -19,6 +19,7 @@
 import java.util.Collection;
 import java.util.Map;
 
+import org.springframework.security.core.AuthenticatedPrincipal;
 import org.springframework.security.core.GrantedAuthority;
 import org.springframework.security.core.SpringSecurityCoreVersion;
 import org.springframework.security.core.Transient;
@@ -29,6 +30,7 @@
  * {@link Jwt} {@code Authentication}.
  *
  * @author Joe Grandja
+ * @author Andrey Litvitski
  * @since 5.1
  * @see AbstractOAuth2TokenAuthenticationToken
  * @see Jwt
@@ -72,6 +74,22 @@ public JwtAuthenticationToken(Jwt jwt, Collection<? extends GrantedAuthority> au
 		this.name = name;
 	}
 
+	/**
+	 * Constructs a {@code JwtAuthenticationToken} using the provided parameters.
+	 * @param jwt the JWT
+	 * @param principal the principal
+	 * @param authorities the authorities assigned to the JWT
+	 */
+	public JwtAuthenticationToken(Jwt jwt, Object principal, Collection<? extends GrantedAuthority> authorities) {
+		super(jwt, principal, jwt, authorities);
+		this.setAuthenticated(true);
+		if (principal instanceof AuthenticatedPrincipal) {
+			this.name = ((AuthenticatedPrincipal) principal).getName();
+		} else {
+			this.name = jwt.getSubject();
+		}
+	}
+
 	@Override
 	public Map<String, Object> getTokenAttributes() {
 		return this.getToken().getClaims();

oauth2/oauth2-resource-server/src/main/java/org/springframework/security/oauth2/server/resource/authentication/JwtBearerTokenAuthenticationConverter.java

@@ -26,6 +26,7 @@
 import org.springframework.security.oauth2.core.OAuth2AccessToken;
 import org.springframework.security.oauth2.core.OAuth2AuthenticatedPrincipal;
 import org.springframework.security.oauth2.jwt.Jwt;
+import org.springframework.util.Assert;
 
 /**
  * A {@link Converter} that takes a {@link Jwt} and converts it into a
@@ -41,6 +42,7 @@
  * {@link BearerTokenAuthentication}.
  *
  * @author Josh Cummings
+ * @author Andrey Litvitski
  * @since 5.2
  */
 public final class JwtBearerTokenAuthenticationConverter implements Converter<Jwt, AbstractAuthenticationToken> {
@@ -58,4 +60,16 @@ public AbstractAuthenticationToken convert(Jwt jwt) {
 		return new BearerTokenAuthentication(principal, accessToken, authorities);
 	}
 
+	/**
+	 * Sets the {@link Converter Converter&lt;Jwt, Collection&lt;OAuth2AuthenticatedPrincipal&gt;&gt;}
+	 * to use.
+	 * @param jwtPrincipalConverter The converter
+	 * @since 6.5.0
+	 */
+	public void setJwtPrincipalConverter(
+			Converter<Jwt, OAuth2AuthenticatedPrincipal> jwtPrincipalConverter) {
+		Assert.notNull(jwtPrincipalConverter, "jwtPrincipalConverter cannot be null");
+		this.jwtAuthenticationConverter.setJwtPrincipalConverter(jwtPrincipalConverter);
+	}
+
 }