Spring OAuth Server: Token Claim Customization

I wrote previously about the default configuration of Spring oauth-authorization-server. Now let's jump into how we can customize it to suit our requirements. Starting with this article, we will discuss how we can customize the JWT token claims with default configurations (though you can change them as per your requirement).

The default access_token claims are:

{
  "iss": "http://localhost:6060",
  "sub": "spring-test",
  "aud": "spring-test",
  "nbf": 1697183856,
  "exp": 1697184156,
  "iat": 1697183856
}

After customization with additional claims (roles, email, ssn and username), it looks like:

{
  "sub": "spring-test",
  "aud": "spring-test",
  "nbf": 1699198349,
  "roles": [
    "admin",
    "user"
  ],
  "iss": "http://localhost:6060",
  "exp": 1699198649,
  "iat": 1699198349,
  "client_id": "spring-test",
  "email": "test-user@d3softtech.com",
  "ssn": "197611119877",
  "username": "test-user"
}

Let's see how we can achieve that in the Spring Authorization Server.

Spring provides the OAuth2TokenCustomizer<T extends OAuth2TokenContext> interface (FunctionalInterface) to customize the OAuth2Token which can be used to customize any token issued by Spring OAuth Server.

@FunctionalInterface
public interface OAuth2TokenCustomizer<T extends OAuth2TokenContext> {

	/**
	 * Customize the OAuth 2.0 Token attributes.
	 *
	 * @param context the context containing the OAuth 2.0 Token attributes
	 */
	void customize(T context);

}

Therefore, to provide the customizer to the Spring context, define a bean using configuration. You can define one or more customizers to support different token flows.

Single Customizer

If there is a requirement to customize the token for a single flow, it can be defined with Customizer as a bean, like the one below for a client-credential (grant-type) token.

@Configuration
public class AuthorizationServerConfiguration {


      @Bean
  protected OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer() {
    return jwtContext -> {
      if (CLIENT_CREDENTIALS.equals(jwtContext.getAuthorizationGrantType()) && ACCESS_TOKEN.equals(
          jwtContext.getTokenType())) {
        OAuth2ClientCredentialsAuthenticationToken clientCredentialsAuthentication = jwtContext.getAuthorizationGrant();
        Map<String, Object> additionalParameters = clientCredentialsAuthentication.getAdditionalParameters();
        additionalParameters.forEach((key, value) -> jwtContext.getClaims().claim(key, value));
      }
    };
  }
  
}

First, it checks for the flow (client-credential, code, etc.) and then pulls the additional parameters from the request and adds them to the JwtContext. Once added to JwtContext, it will be added to JWT claims in response. 

Additional parameters in the request can be provided as the query param or as a body, such as:

Query Param

In the test (refer to AuthorizationServerTest.verifyTokenEndpoint_WithAdditionParamsAsQueryParam):

webTestClient.post()
            .uri(uriBuilder -> uriBuilder.path("/oauth2/token").queryParam("grant_type", "client_credentials")
                .queryParam("email", TEST_USER_EMAIL).queryParam("ssn", TEST_USER_SSN)
                .queryParam("username", TEST_USER_NAME).queryParam("roles", Set.of("admin", "user")).build())
            .headers(httpHeaders -> httpHeaders.setBasicAuth("spring-test", "test-secret")).exchange()
            .expectStatus().isOk()
            .expectBody()
            .jsonPath("$.access_token").value(this::verifyAccessToken)
            .jsonPath("$.token_type").isEqualTo("Bearer")
            .jsonPath("$.expires_in").isEqualTo(299);

In the example above, a POST request is used to invoke the /oauth2/token endpoint of the authorization server to get the access-token. The minimum parameters required by the authorization server are:

1. grant_type
2. client_id (as header)
3. client_secret (as header)

All the other parameters are additional parameters that you can provide to customize the access_token. As in the above example, we have added email, ssn, username and roles as additional parameters. 

Body Param

In the test (refer to AuthorizationServerTest.verifyTokenEndpoint):

        MultiValueMap<String, String> tokenRequestParams = new LinkedMultiValueMap<>();
        tokenRequestParams.add("grant_type", CLIENT_CREDENTIALS.getValue());
        tokenRequestParams.add("email", TEST_CLIENT_ID);
        tokenRequestParams.add("ssn", TEST_SECRET);
        tokenRequestParams.add("username", TEST_CLIENT_ID);
        tokenRequestParams.add("roles", TEST_SECRET);

        webTestClient.post()
            .uri(uriBuilder -> uriBuilder.path("/oauth2/token").build())
            .contentType(MediaType.APPLICATION_FORM_URLENCODED)
            .body(BodyInserters.fromFormData(tokenRequestParams))
            .headers(httpHeaders -> httpHeaders.setBasicAuth("spring-test", "test-secret"))
            .exchange()
            .expectStatus().isOk()
            .expectBody()
            .jsonPath("$.access_token").exists()
            .jsonPath("$.token_type").isEqualTo("Bearer")
            .jsonPath("$.expires_in").isEqualTo(299);

Parameters to the oauth2/token endpoint can be provided as a body to the POST request. In the above example, client_id and client_secret were passed as basic auth headers, and in this case, as body params.

Multiple Customizer

If there is a need to customize the token for multiple flows, we can take the approach of delegate customizer. The delegate customizer will delegate the request to all custom customizers defined, and therefore, the token will be customized by one or more who are responsible for the through-filter criteria defined in that customizer.

Let's take an example where we want to customize the token for client-credentials and code flow. To do so, we will first define a delegate customizer as:

@Component
public class OAuth2TokenCustomizerDelegate implements OAuth2TokenCustomizer<JwtEncodingContext> {

  private List<OAuth2TokenCustomizer<JwtEncodingContext>> oAuth2TokenCustomizers;

  public OAuth2TokenCustomizerDelegate() {
    oAuth2TokenCustomizers = List.of(
        new OAuth2AuthorizationCodeTokenCustomizer(),
        new OAuth2ClientCredentialsTokenCustomizer());
  }

  @Override
  public void customize(JwtEncodingContext context) {
    oAuth2TokenCustomizers.forEach(tokenCustomizer -> tokenCustomizer.customize(context));
  }
}

As the delegate customizer is defined as a component, it will be consumed by Spring as a bean and will be added to the application context as OAuth2TokenCustomizer. With every request for token creation, a request will be delegated to this customizer to customize.

Now we can define our own customizers that will customize the token according to our needs.

Client-Credentials Token Customizer

public class OAuth2ClientCredentialsTokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {

  @Override
  public void customize(JwtEncodingContext jwtContext) {
    if (CLIENT_CREDENTIALS.equals(jwtContext.getAuthorizationGrantType()) && ACCESS_TOKEN.equals(
        jwtContext.getTokenType())) {
      OAuth2ClientCredentialsAuthenticationToken clientCredentialsAuthentication = jwtContext.getAuthorizationGrant();
      Map<String, Object> additionalParameters = clientCredentialsAuthentication.getAdditionalParameters();
      additionalParameters.forEach((key, value) -> jwtContext.getClaims().claim(key, value));
    }
  }
}

OAuth2ClientCredentialsTokenCustomizer will be responsible for client-credential and grant-type (flow). It will check if the request needs to be handled or not by checking the grant-type and token-type.

Authorization-Code Token Customizer

public class OAuth2AuthorizationCodeTokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {

  @Override
  public void customize(JwtEncodingContext jwtContext) {
    if (AUTHORIZATION_CODE.equals(jwtContext.getAuthorizationGrantType()) && ACCESS_TOKEN.equals(
        jwtContext.getTokenType())) {
      OAuth2AuthorizationCodeAuthenticationToken oAuth2AuthorizationCodeAuthenticationToken = jwtContext.getAuthorizationGrant();
      Map<String, Object> additionalParameters = oAuth2AuthorizationCodeAuthenticationToken.getAdditionalParameters();
      additionalParameters.forEach((key, value) -> jwtContext.getClaims().claim(key, value));
    }
  }
}

OAuth2AuthorizationCodeTokenCustomizer will be responsible as the name suggested for authorization-code grant-type (code flow). 

Sample code

The sample code can be found here. 

The functional test class AuthorizationServerTest has steps on how to:

• Initiate the code flow with the authorized endpoint with the required parameters
• Authenticate the user
• Collect code after successful authentication
• Exchange code for tokens
• Introspect token
• Refresh token
• Revoke tokens
• Introspect post revocation

I hope this post will help you in customizing the tokens.