Allow to set custom BodyExtractor for OAuth 2.0 access token response

[rchigvintsev]

Recently I faced the issue when some OAuth 2.0 provider did not return token type in the access token response. Spring expects "token_type" parameter and throws exception when it is missing. The solution I came up was to change response body and set token type "Bearer" manually. For now it is not easy to do since org.springframework.security.oauth2.client.endpoint.WebClientReactiveAuthorizationCodeTokenResponseClient does not expose any of its logic related to response reading. I think ability to set custom BodyExtractor could facilitate this task.

[jgrandja]

@rchigvintsev Have you considered supplying a custom WebClient via WebClientReactiveAuthorizationCodeTokenResponseClient.setWebClient() configured as follows:

ExchangeFilterFunction tokenResponseFilter = ExchangeFilterFunction.ofResponseProcessor(response -> {
	ClientResponse.Builder builder = ClientResponse.from(response);

	// TODO Change response body and set token_type to Bearer

	return Mono.just(builder.build());
});

WebClient webClient = WebClient.builder()
		.filter(tokenResponseFilter)
		.build();

This configuration should allow you to post-process the response and return a modified response before the default OAuth2BodyExtractors.oauth2AccessTokenResponse() is applied - avoiding the error from occurring. I believe this should work for you?

[rchigvintsev]

@jgrandja
Yes, I implemented almost the same ExchangeFilterFunction and of course it works but using custom body extractor looks like a cleaner solution to me:

WebClientReactiveAuthorizationCodeTokenResponseClient 
        = new WebClientReactiveAuthorizationCodeTokenResponseClient();
client.setTokenResponseBodyExtractor(NaughtyOAuth2ProviderBodyExtractor())

public class NaughtyOAuth2ProviderBodyExtractor extends OAuth2AccessTokenResponseBodyExtractor {
    @Override
    protected TokenResponse parse(Map<String, Object> json) {
        Map<String, Object> newJson = new HashMap<>(json);
        json.put("token_type", "Bearer");
        super.parse(newJson);
    }
}

[jgrandja]

@rchigvintsev

using custom body extractor looks like a cleaner solution to me

IMO, it's pretty similar but it does help with reading/writing the response body which is convenient. However, there are quite a few configuration points for WebClient and it really should be configured directly with WebClient.

If we were to provide setTokenResponseBodyExtractor() than I could see a user requesting setTokenRequestBodyInserter() or setDefaultHeaders() etc. This would result in WebClientReactiveAuthorizationCodeTokenResponseClient simply delegating configuration to WebClient and not much more than that. It would simply pollute the API and not really bring much value. Does that make sense?

I'm reluctant on adding this since you can directly configure WebClient. This is the main reason we exposed setWebClient() in the first place.

[rchigvintsev]

@jgrandja
This makes sense, of course, although using a filter still seems to me like a diving to a lower API level.
Maybe it would be better to have some not related to WebClient method/interface for converting map of parameters from response body to OAuth2AccessTokenResponse.
Anyway thanks for considering this issue! I'm going to close it.

[jgrandja]

Thank you for understanding @rchigvintsev

[avinashkris9]

@rchigvintsev would you mind sharing how you have transformed response.

[rchigvintsev]

@avinashkris9 Yes, certainly. Here is my filter function with several additional methods:

@Override
public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
    return next.exchange(request)
            .flatMap(response -> {
                if (!isJsonCompatibleResponse(response)) {
                    log.warn("Content type of access token response is not compatible with {}",
                            MediaType.APPLICATION_JSON);
                    return Mono.just(response);
                }
                return readParameters(response)
                        .map(this::addTokenTypeIfNecessary)
                        .map(params -> createNewResponse(response, params));
            });
}

private boolean isJsonCompatibleResponse(ClientResponse response) {
    return response.headers().contentType()
            .map(mediaType -> mediaType.isCompatibleWith(MediaType.APPLICATION_JSON))
            .orElse(false);
}

private Mono<? extends Map<String, Object>> readParameters(ClientResponse response) {
    ParameterizedTypeReference<Map<String, Object>> type = new ParameterizedTypeReference<>() {
    };
    BodyExtractor<Mono<Map<String, Object>>, ReactiveHttpInputMessage> extractor = BodyExtractors.toMono(type);
    return response.body(extractor);
}

private Map<String, Object> addTokenTypeIfNecessary(Map<String, Object> params) {
    if (params.get("token_type") == null) {
        Map<String, Object> newParams = new HashMap<>(params);
        newParams.put("token_type", defaultTokenType);
        return newParams;
    }
    return params;
}

private ClientResponse createNewResponse(ClientResponse originalResponse, Map<String, Object> params) {
    Publisher<Map<String, Object>> input = Mono.just(params);
    ResolvableType bodyType = ResolvableType.forInstance(params);
    HttpMessageEncoder<Object> encoder = new Jackson2JsonEncoder();
    MimeType elementType = MimeTypeUtils.APPLICATION_JSON;
    Map<String, Object> hints = Collections.emptyMap();
    Flux<DataBuffer> newBody = encoder.encode(input, dataBufferFactory, bodyType, elementType, hints);
    return ClientResponse.create(originalResponse.statusCode(), originalResponse.strategies())
            .headers(headers -> headers.addAll(originalResponse.headers().asHttpHeaders()))
            .body(newBody)
            .build();
}

[avinashkris9]

@rchigvintsev thanks !