Customize the response of token endpoint and the token claims in Spring-Authorization-Server

It’s common to customize the token claims of the access token and also the token response of the token endpoint, when using Spring-Authorization-Server supporting oauth2 flow. https://github.com/spring-projects/spring-authorization-server

One example is to add “userId” to the token response and the refresh token expiration.

The example to customize the access token(although sounds more like id token) could be to add authorization related information, if the API client doesn’t want to call user info endpoint to obtain it. For example, more complicated “aud”, or custom field to represent user attribute.

Customizing token claims is very simple and it’s well-documented in the official documentation.

Core Model / Components

    @Bean
    fun jwtCustomizer(): OAuth2TokenCustomizer<JwtEncodingContext> {
        return OAuth2TokenCustomizer<JwtEncodingContext> { context: JwtEncodingContext ->
            if (context.tokenType.value == OAuth2ParameterNames.ACCESS_TOKEN) {
                context.claims.claim("test", "test")
            } else if (context.tokenType.value == OidcParameterNames.ID_TOKEN) {
                context.claims.claim("IdTokenTest", "IdTokenTest")
            }
        }
    }

Customizing the token response is a bit more tricky. I got an idea from the following discussion. Basically it adds the extra attributes to the .additionalParameters(additionalParameters) and rebuilds an access token response.

Simplify customizing the access token response · Issue #925 ·…

class CustomTokenResponseHandler: AuthenticationSuccessHandler {

    private val accessTokenHttpResponseConverter: HttpMessageConverter<OAuth2AccessTokenResponse> =
        OAuth2AccessTokenResponseHttpMessageConverter()

    override fun onAuthenticationSuccess(
        request: HttpServletRequest,
        response: HttpServletResponse,
        authentication: Authentication
    ) {
        val accessTokenAuthentication = authentication as OAuth2AccessTokenAuthenticationToken

        val accessToken = accessTokenAuthentication.accessToken
        val refreshToken = accessTokenAuthentication.refreshToken
        val additionalParameters = mapOf("userId" to "userId123")

        val builder = OAuth2AccessTokenResponse.withToken(accessToken.tokenValue)
            .tokenType(accessToken.tokenType)
            .scopes(accessToken.scopes)
        if (accessToken.issuedAt != null && accessToken.expiresAt != null) {
            builder.expiresIn(ChronoUnit.SECONDS.between(accessToken.issuedAt, accessToken.expiresAt))
        }
        if (refreshToken != null) {
            builder.refreshToken(refreshToken.tokenValue)
        }
        builder.additionalParameters(additionalParameters)


        val accessTokenResponse = builder.build()
        val httpResponse = ServletServerHttpResponse(response)
        this.accessTokenHttpResponseConverter.write(accessTokenResponse, null, httpResponse)
    }
}

You also need to put this custom success handler to the security config.

@Autowired
private lateinit var customTokenResponseHandler: CustomTokenResponseHandler
        
http.getConfigurer(OAuth2AuthorizationServerConfigurer::class.java)
            .tokenEndpoint {
                it.accessTokenResponseHandler(
                    customTokenResponseHandler
                )
            }