Customizing Jwt claims and headers needs to be more flexible

[jgrandja]

The NimbusJwsEncoder.jwtCustomizer (#173) needs to be re-designed as it does not easily provide all the context required for customization. For example, it is complicated to obtain the associated OAuth2Authorization and/or RegisteredClient in order to provide context for token customization.

[joshuatcasey]

Glad to see this, @jgrandja ! The sso4k8s team hopes to add some of the standard claims to the id_token pretty soon.

[jgrandja]

@joshuatcasey I'm getting close to a flexible design so hoping to have this merged next week. I'll be looking for your teams feedback on whether the 2nd iteration of the design provides the flexibility you need.

[nickmelis]

Excellent work @jgrandja, I too am looking forward to this. While we're here, can I raise a couple of issues I stumbled upon today?
I noticed that OAuth2TokenIssuerUtil has the issuer and timeToLive values hardcoded (with a big TODO next to it), which means there's currently no way to host my auth server on anything different than http://auth-server:9000. Is it something you're planning to tackle as part of this big refactoring?

[jgrandja]

@nickmelis

Is it something you're planning to tackle as part of this big refactoring?

I'm not sure if it will be part of the commit for this ticket. But it will definitely be part of the 0.1.0 (end of month).

[joshuatcasey]

@nickmelis in the meantime it's possible to use the setJwtCustomizer method to modify claims after the default claim values have been added.

Something like this works for us:

    @Bean
    public JwtEncoder jwtEncoder(final JWKSource jwkSource) {
        final NimbusJwsEncoder nimbusJwsEncoder = new NimbusJwsEncoder(jwkSource);
        nimbusJwsEncoder.setJwtCustomizer((header, claims) -> claims.issuer("our.url.here"));
        return nimbusJwsEncoder;
    }

[nickmelis]

Oh yes you're absolutely right @joshuatcasey, I didn't even think about it. Thanks!

[jgrandja]

@Kehrlann @joshuatcasey @jzheaux @anoopgarlapati @nickmelis

I've flushed out the new design for Jwt customizer and having a WIP in this branch.

Run the test OAuth2AuthorizationCodeGrantTests.requestWhenTokenRequestValidThenReturnAccessTokenResponse() to see how the OAuth2TokenCustomizer<JwtEncodingContext> is applied. It's configured here.

The updates have only been applied to authorization_code grant so you can step through the code in OAuth2AuthorizationCodeAuthenticationProvider.

I likely won't get back to this until Thursday. Feedback would be super helpful to ensure this new design provides the flexibility required for the various use cases.

[joshuatcasey]

I like the flexibility of the model. The context object can contain all the pertinent information.

One aspect of sso4k8s is that we anticipate that the authorization context (probably org.springframework.security.oauth2.server.authorization.OAuth2Authorization?) will need to contain all of the information about the user. We likely won't have the ability to load information about the user later when the token is requested.

If we could add the entire principal to the OAuth2Authorization in OAuth2AuthorizationEndpointFilter I think that would be enough for us to populate claims such as given_name, family_name, and roles. I'll play around with the following and push to a branch at https://github.com/pivotal/spring-authorization-server.

Maybe this:

		OAuth2Authorization.Builder builder = OAuth2Authorization.withRegisteredClient(registeredClient)
				.principal(principal)
				.attribute(OAuth2AuthorizationAttributeNames.AUTHORIZATION_REQUEST, authorizationRequest);

or this:

		OAuth2Authorization.Builder builder = OAuth2Authorization.withRegisteredClient(registeredClient)
				.principalName(principal.getName())
				.attribute(OAuth2AuthorizationAttributeNames.AUTHORIZATION_REQUEST, authorizationRequest);
				.attribute("Some-Attribute-Name", principal);

[joshuatcasey]

Would it be possible to add something to the context to state the grant_type and token type (access_token vs id_token), since the customizer is used for both access_token and id_token.

Note: the branch didn't use the customizer for client_credentials access_token, was that intentional?

[nickmelis]

This may be a stupid question. As part of my implementation of UserDetailsService.loadUserByUsername, I fetch user details from DB, including roles and authorities, and set them to the UserDetails being returned.
Is there any chance to access these information within the new token customizer?
Right now what I have to do is fetch the user first (in UserDetailsService), then in the token customizer, make a second DB query to fetch authorities and everything that needs to go into the JWT.
Hope it makes some sort of sense.

[jgrandja]

@joshuatcasey I have a major refactoring task for the OAuth2Authorization model before I release 0.1.0. I will enhance it to support your needs for:

add the entire principal to the OAuth2Authorization in OAuth2AuthorizationEndpointFilter

add something to the context to state the grant_type

Regarding:

the branch didn't use the customizer for client_credentials access_token, was that intentional?

The branch only demonstrated for the authorization_code flow. The other grants were not updated, however, the same logic will be applied after we finalize the new design.

[jgrandja]

@nickmelis There are no stupid questions :)

On the next update, you will no longer have to perform a 2nd DB query as the OAuth2Authorization will contain the principal.

See previous comment:

add the entire principal to the OAuth2Authorization in OAuth2AuthorizationEndpointFilter

[nickmelis]

@jgrandja thanks, really looking forward to it! Do you have an estimate for when this code will go in?

[jgrandja]

@joshuatcasey @nickmelis @fhanik @Kehrlann @anoopgarlapati A lot of work was put into the customizer and I think it's pretty flexible. Please take a look at it and provide feedback as soon as you can. I'm releasing 0.1.0 next Thursday Feb 11 so I want to make sure I get in any further enhancements required from your end.

Start by looking at these tests:

• OAuth2AuthorizationCodeGrantTests and associated config
• OidcTests and associated config

[nickmelis]

@jgrandja I can confirm it works a treat! I'm really pleased with it!
I can both set claims the way I used to do with the old jwtCustomizer and propagate Authentication authorities into JWT claims.
I will do some big refactoring to my app tomorrow so if I think anything is missing, I'll drop a message here. Thanks again!

[jgrandja]

Excellent @nickmelis ! Yes, please provide further feedback here if anything else comes up.

[nickmelis]

Hi @jgrandja, I noticed that JwtEncodingContextUtils still has RS256 hardcoded as algorithm header, which is currently preventing me from using EC keys. I raised the issue originally a while ago as part of #190, and I'm not clear as to whether there's a way to override/bypass this behaviour. Would be nice to get your advice on this.

[anoopgarlapati]

@jgrandja I spent some time customizing the OAuth2TokenCustomizer to suit some of the custom claims for my service and it worked great so far. I will let you know of any improvements as I continue to work with it.

[jgrandja]

@nickmelis Here is how you would override the alg header:

@Bean
OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer() {
	return context -> {
		// Override the default signing algorithm from RS256 to ES256
		// NOTE: The ES256 key must be available in the configured `JWKSource<SecurityContext>` `@Bean`
		context.getHeaders().jwsAlgorithm(SignatureAlgorithm.ES256);

		// TODO Further customizations

	};
}

[jgrandja]

@anoopgarlapati @Kehrlann Thank you for the feedback! I'm happy that it suits your requirements. But I also want to ensure that the API is intuitive to use as this is one of our primary goals. Can you comment on this further?

Was it intuitive to use? Does the naming of the API's make sense? e.g. OAuth2TokenCustomizer -> OAuth2TokenContext -> JwtEncodingContext

[Kehrlann]

1. For our current use-case, the OAuth2TokenContext isn't useful without the full JwtEncodingContext. It might
   be for other people though.
2. Our current use-case reaches deep into the context to grab the authorization scopes (see below for code example).
   But we do so without checking whether it's an authorization_code-type (or, in the future, implicit) or another grant
   type. For example, implicitly here in client_credentials there are no scopes, so it's null. It works 100% but that
   null-check feels a bit off. Any thoughts on this ?

private boolean shouldCustomize(JwtEncodingContext context) {
    if (!Objects.equals(new TokenType(OidcParameterNames.ID_TOKEN), context.getTokenType())) {
        return false;
    }

    final OAuth2Authorization authorization = context.get(OAuth2Authorization.class);
    if (authorization == null) return false;

    Set<String> scopes = authorization.getAttribute(OAuth2Authorization.AUTHORIZED_SCOPE_ATTRIBUTE_NAME);
    if (scopes == null) return false;

    return scopes.contains(OidcScopes.PROFILE);
}

[jgrandja]

@Kehrlann

implicitly here in client_credentials there are no scopes, so it's null

I'm confused with this statement as I interpret the shouldCustomize() logic as follows:

If the token being customized is an ID Token and the user has consented to the profile scope then return true, else false.

If my interpretation is correct then the grant_type should be authorization_code NOT client_credentials - and you can optimize the code as follows:

private boolean shouldCustomize(JwtEncodingContext context) {
	if (context.getTokenType().getValue().equals(OidcParameterNames.ID_TOKEN)) {
		OAuth2Authorization authorization = context.getAuthorization();
		Set<String> authorizedScopes = authorization.getAttribute(OAuth2Authorization.AUTHORIZED_SCOPE_ATTRIBUTE_NAME);
		return authorizedScopes.contains(OidcScopes.PROFILE);
	}
	return false;
}

NOTE: This code is null-safe and if it does throw a NPE then please log an issue as it will be a bug on our end.

the OAuth2TokenContext isn't useful without the full JwtEncodingContext. It might
be for other people though

The design intent of the OAuth2TokenContext is to allow for reuse in other OAuth 2.0 Token encoding contexts. It's not totally clear to me which other context(s) it might be used but I have a feeling it might come into play. If not, then it will be easy to merge OAuth2TokenContext into JwtEncodingContext.

[Kehrlann]

@jgrandja

Right, sorry that was really poorly phrased.

• You are correct, in my case we're looking at authorization_code grant
• Those null-checks are paranoid checks on our part, we shouldn't worry about it and just report it as a bug if it every blows up 👌

I think what I had in mind was related type-safety, to distinguish between grant types at the OAuth2Authorization level. Imagine you had a ClientCredentialsOAuth2Authorization, an AuthorizationCodeOAuth2Authorization, etc.

When getting the OAuth2Authorization object through context.getAuthorization(), as a user you could check its type, and have different flows based on different grant types. Today a user can tell the flow through .getAuthorizationGrantType(), but this does not provide type-level guarantees about what you'll get when you call .getAttribute().

Whereas if you had an AuthorizationCodeOAuth2Authorization, it could define .getAuthorizedScopes(), whereas ClientCredentialsOAuth2Authorization wouldn't define it. It makes the "attributes" more discoverable as I don't need to check through the code which attributes are set on each flow.

This kills the flexibility of getAttribute, so I'm unsure whether it's a Good Idea™ or not.

Hope that makes more sense!

[anoopgarlapati]

@jgrandja In our server we map authorized scopes to a custom claim called permissions and we do not currently have the scope claim in our access token. So, I was customizing the access token in the following manner

@Bean
public OAuth2TokenCustomizer<JwtEncodingContext> customizer() {
	return context -> {
		if (context.getTokenType() == OAuth2TokenType.ACCESS_TOKEN) {
			Set<String> authorizedScopes =
					context.getAuthorization().getAttribute(OAuth2Authorization.AUTHORIZED_SCOPE_ATTRIBUTE_NAME);
			if (!CollectionUtils.isEmpty(authorizedScopes)) {
				context.getClaims().claims(claims -> {
					claims.remove(OAuth2ParameterNames.SCOPE);
					claims.put("permissions", authorizedScopes);
				});
			}
		}
	};
}

When I execute the authorization_code flow with this customization, I do get the access token I desired with permissions claim containing the authorized scopes but the scope attribute in the response returns null. Here, I would expect the authorized scope attribute to be returned.

This is happening because the scopes in OAuth2AccessToken are set by extracting the scope claim from the Jwt.
There maybe other cases like ours where an authorization server would customize and remove the scope claim and this would return null scopes in the access token response. I suggest the scopes that are set here are retrieved using the OAuth2TokenContext like how it is set initially for the scope claim in the Jwt.

I would like to hear thoughts on this suggestion or other alternatives.

[jgrandja]

@anoopgarlapati Thanks again for staying on top of the feedback. I totally understand the issue you have. This is a bug and I will fix today.

[jgrandja]

@anoopgarlapati This is now resolved via 6ffda38

[jgrandja]

@Kehrlann

to distinguish between grant types at the OAuth2Authorization level. Imagine you had a ClientCredentialsOAuth2Authorization, an AuthorizationCodeOAuth2Authorization, etc.

I did consider this hierarchical design, however, it didn't make sense at this point. As I worked through #213, I found the only difference between an authorization_code and client_credentials OAuth2Authorization was the additional attributes OAuth2AuthorizationRequest and java.security.Principal (Resource Owner) for authorization_code. Other then that, the constructs are the same but differ in data. As we add new grant_type support, we may introduce a hierarchy if it makes sense. For now, this keeps it simple without introducing any unnecessary complexity.

Whereas if you had an AuthorizationCodeOAuth2Authorization, it could define .getAuthorizedScopes(), whereas ClientCredentialsOAuth2Authorization wouldn't define it

We still need authorizedScopes for client_credentials. This is typically auto-approved given there is no authorization consent phase.

Thanks for pointing this out as I discovered I missed storing the authorizedScopes in the client_credentials OAuth2Authorization. See c00226d

I also added JwtEncodingContext.getAuthorizedScopes() as a convenient accessor as it seems it would be useful. See ece5f2b

[anoopgarlapati]

@jgrandja Thanks for the updates!!
I was trying another customization and was unable to access the token endpoint request.
I have a use-case to add a custom claim based on custom request parameters specific to our authorization server and was able to do that for authorization_code flow (as the request parameters are usually in authorize request) but was unable to find the token request parameters.

@Bean
public OAuth2TokenCustomizer<JwtEncodingContext> customizer() {
	return context -> {
		if (context.getTokenType() == OAuth2TokenType.ACCESS_TOKEN) {
			OAuth2Authorization oAuth2Authorization = context.getAuthorization();
			if (oAuth2Authorization != null) {
				OAuth2AuthorizationRequest authorizationRequest =
						oAuth2Authorization.getAttribute(OAuth2AuthorizationRequest.class.getName());
				if (authorizationRequest != null) {
					// process custom request parameter and add the custom claim
					// similar object for token request is not available in the OAuth2TokenContext for customization
					// hence unable to apply this customization for client_credentials and refresh_token grants
				}
			}
		}
	};
}

The token requests in the three grants (authorization_code exchange request, client_credentials and refresh_token) are not accessible in the context for customizing Jwt claims.
I understand there is a release today and this seems like changes are required at lot of places to get the token request in the OAuth2TokenCustomizer, so can this be considered for next milestone at least?
I can come up with the design and work on it if needed.

[jgrandja]

@anoopgarlapati You beat me to it ! :)

This is the last change I'm getting in before I release today. The token request (additional) parameters will be stored here. And you can access it via JwtEncodingContext.getAuthorizationGrant(). I will let you know when it's merged.

[Kehrlann]

@jgrandja makes complete sense! And with the getAuthorizedScopes(), I have everything I need right here 🍾

[jgrandja]

@anoopgarlapati See #226

[giorgimoreira]

Good morning, would anyone have an example that how to use JWT customization(claims) after successful authentication? Thank you so much!

[joshuatcasey]

@giorgimoreira are you looking for a specific code example? There are a few in the tests, if you look for usages of OAuth2TokenCustomizer<JwtEncodingContext>. Just provide a @Bean that implements that interface and it will be autoconfigured.

Here's an example that relies on authentication, is this helpful?

    @Bean
    public OAuth2TokenCustomizer<JwtEncodingContext> buildCustomizer() {
        OAuth2TokenCustomizer<JwtEncodingContext> customizer = (context) -> {
            UsernamePasswordAuthenticationToken token = (UsernamePasswordAuthenticationToken) SecurityContextHolder.getContext().getAuthentication();

            if (token.isAuthenticated() && OAuth2TokenType.ACCESS_TOKEN.equals(context.getTokenType())) {
                context.getClaims().claim("user-authorities", token.getAuthorities()
                        .stream()
                        .map(GrantedAuthority::getAuthority)
                        .collect(Collectors.toList()));
            }

        };
        return customizer;
    }

[atjohn-csam]

Can anyone pls help me add some custom parameters to token response?(not inside the claim)
{
"access_token": **",
"scope": "articles.read openid",
"token_type": "Bearer",
"expires_in": 299
},
in this i want to add "param1":param

[Kehrlann]

Hey @atjohn-csam 👋

What you could do is add some additionalParameters into the OAuth2AccessTokenAuthenticationToken that is used in the token endpoint. Those are ultimately written to the token response in OAuth2AccessTokenResponseHttpMessageConverter.java

The cleanest extension point is probably when those authentication objects are created, but beware, they can be produced in three grant types, each with their own AuthenticationProvider:

• access_token (provider: OAuth2AuthorizationCodeAuthenticationProvider)
• refresh_token (provider: OAuth2RefreshTokenAuthenticationProvider)
• client_credentials (provider: OAuth2ClientCredentialsAuthenticationProvider)

So you would have to tweak all three of those. Potentially wrap them all in lightweight auth provider:

class AddStuffToAccessTokenProvider implements AuthenticationProvider {
	private AuthenticationProvider delegate; // this can be either of the above providers

	@Override
	public Authentication authenticate(Authentication authentication) throws AuthenticationException {
		OAuth2AccessTokenAuthenticationToken rawAuthentication = delegate.authenticate(authentication);
		Map<String, Object> newParameters = new HashMap<>(rawAuthentication.getAdditionalParameters());
		newParameters.put(..., ...);
		// technically I believe you can mutate the previous auth object...
		// but I would advise against it, it might not be possible in the future
		return new OAuth2AccessTokenAuthenticationToken(
				rawAuthentication.getRegisteredClient(),
				rawAuthentication.getPrincipal(),
				rawAuthentication.getAccessToken(),
				rawAuthentication.getRefreshToken(),
				newParameters);
	}

	@Override
	public boolean supports(Class<?> authentication) {
		return OAuth2ClientCredentialsAuthenticationToken.class.isAssignableFrom(authentication);
	}

}

Then you have to wire those up, e.g. at init time in the OAuth2AuthorizationServerConfigurer, or with a custom ObjectPostProcessor.

There are other extension point you could hack your way into, such as OAuth2TokenEndpointFilter#setAuthenticationSuccessHandler or others[1]

Hope it helps!

[1] I was thinking about OAuth2AccessTokenResponseHttpMessageConverter, but apparently it's hardcoded into the Filter. Maybe there is a PR to make it configurable 🤷. Maybe there's a PR here to make it configurable.

[atjohn-csam]

I can't thank you enough for your help @Kehrlann Kehrlann. I was able to get it work. It was a little tricky to get AddStuffToAccessTokenProvider handle three separate customizations(For the three grant types), but i was too excited to see this output:
{
"access_token": "jwt",
"scope": "articles.read openid",
"customparam": "hello",
"token_type": "Bearer",
"expires_in": 299
}

Thank you again, I hope Spring provides us an easier solution in the next releases. Cheers.

[Kehrlann]

Hey, no worries! Glad to see it worked for you.

FYI I had something like this in mind, but I have not tested it.

OAuth2AuthorizationServerConfigurer<HttpSecurity> authorizationServerConfigurer = new OAuth2AuthorizationServerConfigurer<>();
authorizationServerConfigurer
        .withObjectPostProcessor(new ObjectPostProcessor<AuthenticationProvider>() {
            @Override
            public <O extends AuthenticationProvider> O postProcess(O object) {
                if (
                        object instanceof OAuth2AuthorizationCodeAuthenticationProvider
                                || object instanceof OAuth2ClientCredentialsAuthenticationProvider
                                || object instanceof OAuth2RefreshTokenAuthenticationProvider
                ) {
                    return (O) new AddStuffToAccessTokenProvider(object);
                } else {
                    return object;
                }
            }
        });

Cheers!

[atjohn-csam]

Thank you again @Kehrlann Kehrlann.

I tried something similar,

@Component
    public class CustomBeanPostProcessor implements BeanPostProcessor {


        @Override
        public Object postProcessBeforeInitialization(Object bean, String beanName)
                throws BeansException {

            return bean;
        }

        @Override
        public Object postProcessAfterInitialization(Object bean, String beanName)
                throws BeansException {
            if (bean instanceof OAuth2AuthorizationCodeAuthenticationProvider|
                    bean instanceof OAuth2RefreshTokenAuthenticationProvider |
                    bean instanceof OAuth2ClientCredentialsAuthenticationProvider) {
                return new AddStuffToAccessTokenProvider(bean);
            }
            return bean;
        }

}

The problem was AddStuffToAccessTokenProvider was not able to find the delegate correctly. It always had OAuth2AuthorizationCodeAuthenticationProvider as delegate even though i tried a client_secret workflow. So I ended up creating three separate implantations of AddStuffToAccessTokenProvider, which will get initialized from the bean post processor based on the grant type conditions.

Thank you again, your solution works like a charm!

[Waylon-Firework]

Do anyone know why the authorities only include scope. We need user details role authorities

{
    "authorities": [
        {
            "authority": "SCOPE_message.read"
        }
    ],
}