How to inject WebResponseExceptionTranslator in TokenEndPoint Spring oAuth2

Question

I am trying to customize the way the Spring Security oAuth TokenEndpoint class handles exceptions. I want to return a more customized JSON response than the DefaultWebResponseExceptionTranslator currently does. I cannot figure out how to inject a custom Translator via XML. Here are my code snippets:

spring-security.xml

<oauth:authorization-server
    client-details-service-ref="clientDetails" token-services-ref="tokenServices" 
    user-approval-handler-ref="userApprovalHandler" request-validator-ref="requestValidator" >
    <oauth:client-credentials />
    <oauth:password />
</oauth:authorization-server>

Here is the relevant code portion of AbstractServerBeanDefinitionParser

    BeanDefinitionBuilder tokenEndpointBean = BeanDefinitionBuilder.rootBeanDefinition(TokenEndpoint.class);
    tokenEndpointBean.addPropertyReference("clientDetailsService", clientDetailsRef);
    tokenEndpointBean.addPropertyReference("tokenGranter", tokenGranterRef);
    authorizationEndpointBean.addPropertyReference("oAuth2RequestValidator", oAuth2RequestValidatorRef);
    parserContext.getRegistry()
            .registerBeanDefinition("oauth2TokenEndpoint", tokenEndpointBean.getBeanDefinition());
    if (StringUtils.hasText(oAuth2RequestFactoryRef)) {
        tokenEndpointBean.addPropertyReference("oAuth2RequestFactory", oAuth2RequestFactoryRef);
    }
    if (StringUtils.hasText(oAuth2RequestValidatorRef)) {
        tokenEndpointBean.addPropertyReference("oAuth2RequestValidator", oAuth2RequestValidatorRef);
    }

TokenEndpoint extends AbstractEndpoint which is where the translator is defined.

@FrameworkEndpoint
public class TokenEndpoint extends AbstractEndpoint {

    @ExceptionHandler(Exception.class)
    public ResponseEntity<OAuth2Exception> handleException(Exception e) throws Exception {
        logger.info("Handling error: " + e.getClass().getSimpleName() + ", " + e.getMessage());
        return getExceptionTranslator().translate(e);
    }

Here is a snippet from AbtractEndpoint

public class AbstractEndpoint implements InitializingBean {

   protected final Log logger = LogFactory.getLog(getClass());

   private WebResponseExceptionTranslator providerExceptionHandler = new DefaultWebResponseExceptionTranslator();

   public void setProviderExceptionHandler(WebResponseExceptionTranslator providerExceptionHandler) {
       this.providerExceptionHandler = providerExceptionHandler;
   }

So my question is. How do I call the setProviderExceptionHandler to pass in my custom class via my spring-security.xml?

Answer 1

In order to customize spring-oauth's exceptions handler you have to define an instance of WebResponseExceptionTranslator:

@Bean
public WebResponseExceptionTranslator webResponseExceptionTranslator() {
     return new DefaultWebResponseExceptionTranslator() {

         @Override
         public ResponseEntity<OAuth2Exception> translate(Exception e) throws Exception {
             ResponseEntity<OAuth2Exception> responseEntity = super.translate(e);
             OAuth2Exception body = responseEntity.getBody();
             HttpHeaders headers = new HttpHeaders();
             headers.setAll(responseEntity.getHeaders().toSingleValueMap());
             // do something with header or response
             return new ResponseEntity<>(body, headers, responseEntity.getStatusCode());
         }
     };
 }

and pass it to exceptionTranslator property of configurer. I used annotation-based configuration so it looks like that:

@Configuration
@EnableAuthorizationServer
public class OAuth2ServerConfiguration extends AuthorizationServerConfigurerAdapter {
    // some config

    @Override
    public void configure(AuthorizationServerEndpointsConfigurer endpoints) throws Exception {
        endpoints
                .pathMapping("/oauth/authorize", "/authorize")
                .pathMapping("/oauth/error", "/error")
                .exceptionTranslator(webResponseExceptionTranslator())
        ;
    }
}

Just check which properties available for <oauth:authorization-server /> with xml-based configuration.