Extend Micronaut CustomJWTClaimsSetGenerator to provide all attributes

I have the following two classes which provide the JWT authentication mechanisem.

CustomDelegatingAuthenticationProvider

@Singleton
@Replaces(value = DelegatingAuthenticationProvider.class)
public class CustomDelegatingAuthenticationProvider extends DelegatingAuthenticationProvider {
    /**
     * @param userFetcher        Fetches users from persistence
     * @param passwordEncoder    Collaborator which checks if a raw password matches an encoded password
     * @param authoritiesFetcher Fetches authorities for a particular user
     */
    public CustomDelegatingAuthenticationProvider(UserFetcher userFetcher, PasswordEncoder passwordEncoder, AuthoritiesFetcher authoritiesFetcher) {
        super(userFetcher, passwordEncoder, authoritiesFetcher);
    }

    @Override
    protected Publisher<AuthenticationResponse> createSuccessfulAuthenticationResponse(AuthenticationRequest authenticationRequest, UserState userState) {

        if (userState instanceof UserMember) {
            UserMember user = (UserMember) userState;
            return Flowable
                    .fromPublisher(authoritiesFetcher.findAuthoritiesByUsername(user.getUsername()))
                    .map(authorities -> new HDSUser(user.getUsername(), authorities, user.getId()));
        }
        return super.createSuccessfulAuthenticationResponse(authenticationRequest, userState);
    }
}

CustomJWTClaimsSetGenerator

@Singleton
@Replaces(value = JWTClaimsSetGenerator.class)
public class CustomJWTClaimsSetGenerator extends JWTClaimsSetGenerator {

    CustomJWTClaimsSetGenerator(TokenConfiguration tokenConfiguration, @Nullable JwtIdGenerator jwtIdGenerator, @Nullable ClaimsAudienceProvider claimsAudienceProvider) {
        super(tokenConfiguration, jwtIdGenerator, claimsAudienceProvider);
    }

    protected void populateWithUserDetails(JWTClaimsSet.Builder builder, UserDetails userDetails) {
        super.populateWithUserDetails(builder, userDetails);

        if (userDetails instanceof HDSUser) {
            builder.claim("userId", ((HDSUser) userDetails).getId());
        }
    }

}

The default response to the client looks like this:

My question. How can I extend the class to return all user attributes? Besides username I want to have the user id.

UPDATE

HDS user class which gathers the DB id

@CompileStatic
public class HDSUser  extends UserDetails {

    private long id;

    public HDSUser(String username, Collection<String> roles, long id) {
        super(username, roles);
        this.id = id;
    }

    public long getId() {
        return id;
    }

    public void setId(long id) {
        this.id = id;
    }
}

Solution

To extend the returned data you need to extend (implement custom) TokenRenderer as well as a custom version of the AccessRefreshToken. Just as a simple example see the following code snipped which will extend the default access token payload with userId field.

First, create a custom AccessRefreshToken class with additional fields which are required.

@Introspected
@Getter
@Setter
public class CustomAccessRefreshToken extends BearerAccessRefreshToken {
    
    // the new field which will be in the response
    private String userId;

    public CustomAccessRefreshToken(String username,
                                Collection<String> roles,
                                Integer expiresIn,
                                String accessToken,
                                String refreshToken,
                                String tokenType
    ) {
        super(username, roles, expiresIn, accessToken, refreshToken, tokenType);
    }
  
}

Next, we will need a TokenRenderer which will be used by the underlying subsystem to generate our custom token.

@Singleton
@Replaces(value = BearerTokenRenderer.class)
public class CustomTokenRenderer implements TokenRenderer {
    private static final String BEARER_TOKEN_TYPE = HttpHeaderValues.AUTHORIZATION_PREFIX_BEARER;

    @Override
    public AccessRefreshToken render(Integer expiresIn, String accessToken, @Nullable String refreshToken) {
    return new AccessRefreshToken(accessToken, refreshToken, BEARER_TOKEN_TYPE, expiresIn);
    }

    @Override
    public AccessRefreshToken render(Authentication authentication, Integer expiresIn, String accessToken, @Nullable String refreshToken) {
        CustomAccessRefreshToken token =  new CustomAccessRefreshToken(authentication.getName(), authentication.getRoles(), expiresIn, accessToken, refreshToken, BEARER_TOKEN_TYPE);
        // here just take the user data from Authentication object or access any other service
        token.setUserId("Some user id");
        return token;
    }
}

That's it )) Just implement render() method the way you want and add as many custom fields as needed.

The response from the given example will look like

{
    "username": "sherlock",
    "userId": "Some user id",
    "access_token": "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJzaGVybG9jayIsIm5iZiI6MTYzNjk5MTgzMSwicm9sZXMiOltdLCJpc3MiOiJtaWNyb25hdXRndWlkZSIsImV4cCI6MTYzNjk5NTQzMSwiaWF0IjoxNjM2OTkxODMxfQ.Cat1CTsUZkCj-OHGafiefNm1snPsALoaNw9y2xwF5Pw",
    "token_type": "Bearer",
    "expires_in": 3600
}

If you are on the older version of the Micronaut v1.x the TokenRenderer will look like this.

@Singleton
@Replaces(value = BearerTokenRenderer.class)
public class CustomTokenRenderer implements TokenRenderer {
    private static final String BEARER_TOKEN_TYPE = HttpHeaderValues.AUTHORIZATION_PREFIX_BEARER;

    public AccessRefreshToken render(Integer expiresIn, String accessToken, String refreshToken) {
        return new AccessRefreshToken(accessToken, refreshToken, BEARER_TOKEN_TYPE, expiresIn);
    }

    public AccessRefreshToken render(UserDetails userDetails, Integer expiresIn, String accessToken, String refreshToken) {
        CustomAccessRefreshToken token =  new CustomAccessRefreshToken(userDetails.getUsername(), userDetails.getRoles(), expiresIn, accessToken, refreshToken, BEARER_TOKEN_TYPE);
        token.setUserId("Some user id, Some user id");
        return token;
    }
}