Spring Authorization Server customization

@Configuration
@EnableWebSecurity
public class SecurityConfig {

  // ... the rest of the security config ...

  @Bean
  public OAuth2TokenGenerator<OAuth2Token> tokenGenerator(JWKSource<SecurityContext> jwkSource) {
    // The default JWT generator, for access tokens
    var jwtEncoder = new NimbusJwtEncoder(jwkSource);
    var jwtGenerator = new JwtGenerator(jwtEncoder);

    // Our custom refresh token generator
    var refreshTokenGenerator = new CustomRefreshTokenGenerator();

    // We do not include the opaque token generator, OAuth2AccessTokenGenerator, because we're only
    // doing JWT in this example
    return new DelegatingOAuth2TokenGenerator(jwtGenerator, refreshTokenGenerator);
  }

  /**
   * Custom refresh token generator, most of the code is lifted from OAuth2RefreshTokenGenerator.
   */
  static class CustomRefreshTokenGenerator implements OAuth2TokenGenerator<OAuth2RefreshToken> {
    @Nullable
    @Override
    public OAuth2RefreshToken generate(OAuth2TokenContext context) {
      // 📋 BEGIN copy-paste from OAuth2RefreshTokenGenerator
      // Some formatting altered to fit on the blog
      if (!OAuth2TokenType.REFRESH_TOKEN.equals(context.getTokenType())) {
        return null;
      }
      if (isPublicClientForAuthorizationCodeGrant(context)) {
        // Do not issue refresh token to public client
        return null;
      }

      Instant issuedAt = Instant.now();
      var configuredTtl = context.getRegisteredClient().getTokenSettings().getRefreshTokenTimeToLive();
      Instant expiresAt = issuedAt.plus();
      // 📋 END copy-paste

      //🚨 Here is where you set your custom value
      return new OAuth2RefreshToken("jaffars-custom-token", issuedAt, expiresAt);
    }

    // 📋 Pure copy-paste from OAuth2RefreshTokenGenerator
    private static boolean isPublicClientForAuthorizationCodeGrant(OAuth2TokenContext context) {
      if (AuthorizationGrantType.AUTHORIZATION_CODE.equals(context.getAuthorizationGrantType()) &&
           (
            context.getAuthorizationGrant().getPrincipal()
            instanceof OAuth2ClientAuthenticationToken clientPrincipal
           )
         ) {
        return clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE);
      }
      return false;
    }

  }

}

@Configuration
@EnableWebSecurity
public class SecurityConfiguration {

  @Bean
  public SecurityFilterChain securityFilterChain(HttpSecurity http) throws Exception {
    var authServerConfigurer = new OAuth2AuthorizationServerConfigurer();
    authServerConfigurer
        .tokenEndpoint(token -> token.accessTokenResponseHandler(new CustomTokenResponseHandler()));
        // other configurations ...
    http
        .with(authServerConfigurer, Customizer.withDefaults())
        // all your filter chain configuration
        .build();
  }

  /**
   * Code lifted from OAuth2TokenEndpointFilter#sendAccessTokenResponse.
   */
  static class CustomTokenResponseHandler implements AuthenticationSuccessHandler {

    private final OAuth2AccessTokenResponseHttpMessageConverter accessTokenHttpResponseConverter
      = new OAuth2AccessTokenResponseHttpMessageConverter();

    @Override
    public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response,
        Authentication authentication) throws IOException, ServletException {
      // 📋 BEGIN copy-paste from OAuth2TokenEndpointFilter.sendAccessTokenResponse(...)
      // Some formatting altered to fit on the blog
      OAuth2AccessTokenAuthenticationToken accessTokenAuthentication 
        = (OAuth2AccessTokenAuthenticationToken) authentication;

      OAuth2AccessToken accessToken = accessTokenAuthentication.getAccessToken();
      OAuth2RefreshToken refreshToken = accessTokenAuthentication.getRefreshToken();
      Map<String, Object> additionalParameters = accessTokenAuthentication.getAdditionalParameters();

      OAuth2AccessTokenResponse.Builder builder
        = OAuth2AccessTokenResponse.withToken(accessToken.getTokenValue())
          .tokenType(accessToken.getTokenType())
          .scopes(accessToken.getScopes());
      if (accessToken.getIssuedAt() != null && accessToken.getExpiresAt() != null) {
        var exp = ChronoUnit.SECONDS.between(accessToken.getIssuedAt(), accessToken.getExpiresAt());
        builder.expiresIn(exp);
      }
      if (refreshToken != null) {
        builder.refreshToken(refreshToken.getTokenValue());
      }
      // 📋 END copy-paste

      // 🚨 Here we create a map that contains our new parameters, as well as the
      // original additional parameters. This is where we add a custom claim.
      var otherParams = new HashMap<String, Object>();
      otherParams.put("my_parameter", "and here we add a custom value");
      if (!CollectionUtils.isEmpty(additionalParameters)) {
        // Those are the additional parameters from the original authentication
        // process.
        // Importantly, it contains the `id_token` value.
        otherParams.putAll(additionalParameters);
      }
      builder.additionalParameters(otherParams);

      // 📋 Some more copy-paste
      OAuth2AccessTokenResponse accessTokenResponse = builder.build();
      ServletServerHttpResponse httpResponse = new ServletServerHttpResponse(response);
      this.accessTokenHttpResponseConverter.write(accessTokenResponse, null, httpResponse);
    }

  }

}