Customization of TokenEndpoint in Sprin OAuth2

Question

I would like to provide a custom implmentation of the TokenEndpoint class in Spring framework.

Ive copied over the TokenEndpoint class of spring and have made my changes to the required places. But when the applications starts, I'm always getting the error

Caused by: java.lang.IllegalStateException: TokenGranter must be provided

I have provided an implementation for TokenGranter in my OAuthConfig, but spring is not picking up that

@Override
public void configure(AuthorizationServerEndpointsConfigurer endpoints) throws Exception {
    endpoints.pathMapping("/oauth/token", "/oauth/token/v1")
             .tokenServices(tokenServices())
             .tokenGranter(tokenGranter())
            .authenticationManager(authenticationManager).tokenStore(tokenStore())
            .tokenEnhancer(tokenEnhancer()).accessTokenConverter(accessTokenConverter());
}

@Bean
@Primary
public TokenGranter tokenGranter() {
    TokenGranter tokenGranter = null;
    if (tokenGranter == null) {
        tokenGranter = new TokenGranter() {
            private CompositeTokenGranter delegate;

            @Override
            public OAuth2AccessToken grant(String grantType, TokenRequest tokenRequest) {
                if (delegate == null) {
                    delegate = new CompositeTokenGranter(getDefaultTokenGranters());
                }
                return delegate.grant(grantType, tokenRequest);
            }
        };
    }
    return tokenGranter;
}

I even tried to provide this implementation, in my custom TokenEndpoint class. For now, the implementation of custom TokenEndpoint is exactly the same as Spring's TokenEndpoint.

OAuth2AccessToken token = getTokenGranter().grant(tokenRequest.getGrantType(), tokenRequest);

private List<TokenGranter> getDefaultTokenGranters() {
    ClientDetailsService clientDetails = clientDetailsService();
    AuthorizationServerTokenServices tokenServices = tokenServices();
    AuthorizationCodeServices authorizationCodeServices = authorizationCodeServices();
    OAuth2RequestFactory requestFactory = requestFactory();

    List<TokenGranter> tokenGranters = new ArrayList<TokenGranter>();
    tokenGranters.add(new AuthorizationCodeTokenGranter(tokenServices, authorizationCodeServices, clientDetails,
            requestFactory));
    tokenGranters.add(new RefreshTokenGranter(tokenServices, clientDetails, requestFactory));
    ImplicitTokenGranter implicit = new ImplicitTokenGranter(tokenServices, clientDetails, requestFactory);
    tokenGranters.add(implicit);
    tokenGranters.add(new ClientCredentialsTokenGranter(tokenServices, clientDetails, requestFactory));
    if (authenticationManager != null) {
        tokenGranters.add(new ResourceOwnerPasswordTokenGranter(authenticationManager, tokenServices, clientDetails,
                requestFactory));
    }
    return tokenGranters;
}

private DefaultTokenServices createDefaultTokenServices() {
    DefaultTokenServices tokenServices = new DefaultTokenServices();
    tokenServices.setTokenStore(tokenStore());
    tokenServices.setSupportRefreshToken(true);
    tokenServices.setReuseRefreshToken(true);
    tokenServices.setClientDetailsService(clientDetailsService());
    tokenServices.setTokenEnhancer(tokenEnhancer());
    addUserDetailsService(tokenServices, new CustomDetailsService());
    return tokenServices;
}

private ClientDetailsService clientDetailsService() {
    ClientDetailsService clientDetailsService = null;
    clientDetailsService = new InMemoryClientDetailsService();
    addUserDetailsService(createDefaultTokenServices(), new CustomDetailsService());
    return clientDetailsService;
}

private void addUserDetailsService(DefaultTokenServices tokenServices, UserDetailsService userDetailsService) {
    if (userDetailsService != null) {
        PreAuthenticatedAuthenticationProvider provider = new PreAuthenticatedAuthenticationProvider();
        provider.setPreAuthenticatedUserDetailsService(new UserDetailsByNameServiceWrapper<PreAuthenticatedAuthenticationToken>(
                userDetailsService));
        tokenServices
                .setAuthenticationManager(new ProviderManager(Arrays.<AuthenticationProvider> asList(provider)));
    }
}

private AuthorizationCodeServices authorizationCodeServices() {
    AuthorizationCodeServices authorizationCodeServices = new InMemoryAuthorizationCodeServices();
    return authorizationCodeServices;
}

private OAuth2RequestFactory requestFactory() {
    OAuth2RequestFactory requestFactory = new DefaultOAuth2RequestFactory(clientDetailsService());
    return requestFactory;
}

@Bean
public JwtTokenStore tokenStore() {
    JwtTokenStore jwtTokenStore = new JwtTokenStore(accessTokenConverter());
    return jwtTokenStore;
}

@Bean
@Primary
public AuthorizationServerTokenServices tokenServices() {
    final DefaultTokenServices defaultTokenServices = new DefaultTokenServices();
    defaultTokenServices.setAccessTokenValiditySeconds(-1);
    defaultTokenServices.setTokenStore(tokenStore());
    return defaultTokenServices;
}

@Bean
public TokenEnhancer tokenEnhancer() {
    return new CustomTokenEnhancer();
}

@Bean
public JwtAccessTokenConverter accessTokenConverter() {
    final JwtAccessTokenConverter converter = new JwtAccessTokenConverter() {
        @Override
        public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
            return accessToken;
        }
    };
    return converter;
}

Ive been trying to figure this out for a couple of days, but without any luck. So any help would be much appreciated.

Answer 1

I know the question is quite old, but I encountered the same problem and didn't manage to find a complete guide on customizing TokenEndpoint. I wasn't be able to use TokenEnhancer, because I needed to change headers of the response. So, this is the version worked for me.

You define your overwritten controller as usual:

@RequestMapping(value = "/oauth/token")
public class CustomTokenEndpoint extends TokenEndpoint {
  @PostMapping
  public ResponseEntity<OAuth2AccessToken> postAccessToken(
      Principal principal,
      @RequestParam Map<String, String> parameters
  ) throws HttpRequestMethodNotSupportedException {

    ResponseEntity<OAuth2AccessToken> defaultResponse = super.postAccessToken(principal, parameters);

    // do some work

    return defaultResponse;
  }
}

And you need to create your own TokenEndpoint bean:

  @Bean
  @Primary
  public TokenEndpoint tokenEndpoint(AuthorizationServerEndpointsConfiguration conf) {
    TokenEndpoint tokenEndpoint = new CustomTokenEndpoint();
    tokenEndpoint.setClientDetailsService(conf.getEndpointsConfigurer().getClientDetailsService());
    tokenEndpoint.setProviderExceptionHandler(conf.getEndpointsConfigurer().getExceptionTranslator());
    tokenEndpoint.setTokenGranter(conf.getEndpointsConfigurer().getTokenGranter());
    tokenEndpoint.setOAuth2RequestFactory(conf.getEndpointsConfigurer().getOAuth2RequestFactory());
    tokenEndpoint.setOAuth2RequestValidator(conf.getEndpointsConfigurer().getOAuth2RequestValidator());
    tokenEndpoint.setAllowedRequestMethods(conf.getEndpointsConfigurer().getAllowedTokenEndpointRequestMethods());
    return tokenEndpoint;
  }

And here's the kicker. You need to allow overwriting spring beans in your application.properties:

spring.main.allow-bean-definition-overriding: true

Hope this helps someone

Answer 2

Why do you need to implement TokenEndpoint again?

You can create a TokenGranter bean and inject it to default endpoints.

Where is getDefaultTokenGranters() method?

It looks like you have an incomplete copy of AuthorizationServerEndpointsConfigurer source code.

Update:

If you want to customize the token response ,use TokenEnhancer.

for example:

public class CustomTokenEnhancer implements TokenEnhancer {

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        OurUser user = (OurUser) authentication.getPrincipal();
        final Map<String, Object> additionalInfo = new HashMap<>();
        Map<String, Object> userDetails = new HashMap<>();
        userDetails.put(USERID, user.getId().getId());
        userDetails.put(NAME, user.getName());
        userDetails.put(MOBILE, user.getMobile());
        userDetails.put(EMAIL, user.getEmail());
        additionalInfo.put(USERINFO, userDetails);
        // Set additional information in token for retriving in #org.springframework.security.oauth2.provider.endpoint.CheckTokenEndpoint
        ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(additionalInfo);
        return accessToken;
    }

}

in OAuth2 Config:

    @Override
    public void configure(AuthorizationServerEndpointsConfigurer endpoints) throws Exception {
        super.configure(endpoints);
        endpoints.
                .....
                // Include additional information to OAuth2 Access token with custom token enhancer
                .tokenEnhancer(tokenEnhancer());
    }

    @Bean
    public TokenEnhancer tokenEnhancer() {
        return new CustomTokenEnhancer();
    }

https://stackoverflow.com/a/28512607/4377110