Add User-Agent header on Oauth2 Client - Spring Boot

README.md

Spring Security - add User-Agent header on oauth2-client

• create custom token response client class
• modify RestTemplate (add User-Agent header).

Spring Security Issue - 4958

DefaultPasswordTokenResponseClient

.
..
...
public CustomPasswordTokenResponseClient() {

  RestTemplate restTemplate = new RestTemplateBuilder(rt -> rt.getInterceptors().add((request, body, execution) -> {
            request.getHeaders().add(HttpHeaders.USER_AGENT, UUID.randomUUID().toString());
            return execution.execute(request, body);
  })).messageConverters(Arrays.asList(new FormHttpMessageConverter(), new OAuth2AccessTokenResponseHttpMessageConverter())).build();

  restTemplate.setErrorHandler(new OAuth2ErrorResponseErrorHandler());
  this.restOperations = restTemplate;
}
...
..
.

application.properties

*
*
*

spring.security.oauth2.client.registration.CUSTOM_PROVIDER.client-id=CLIENT_ID
spring.security.oauth2.client.registration.CUSTOM_PROVIDER.client-secret=CLIENT_SECRET
spring.security.oauth2.client.registration.CUSTOM_PROVIDER.client-authentication-method=basic
spring.security.oauth2.client.registration.CUSTOM_PROVIDER.authorization-grant-type=password
spring.security.oauth2.client.registration.CUSTOM_PROVIDER.provider=CUSTOM_PROVIDER
spring.security.oauth2.client.registration.CUSTOM_PROVIDER.scope=SCOPES

spring.security.oauth2.client.provider.CUSTOM_PROVIDER.token-uri=OAUTH_TOKEN_URL
  
*
*
*

build.gradle

plugins {
	id 'org.springframework.boot' version '2.3.1.RELEASE'
	id 'io.spring.dependency-management' version '1.0.9.RELEASE'
}

*
*
*

implementation 'org.springframework.boot:spring-boot-starter-webflux'
implementation 'org.springframework.boot:spring-boot-starter-security'
implementation 'org.springframework.boot:spring-boot-starter-oauth2-client'
  
*
*
*

CustomPasswordTokenResponseClient.java

package *;

import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.core.convert.converter.Converter;
import org.springframework.http.HttpHeaders;
import org.springframework.http.RequestEntity;
import org.springframework.http.ResponseEntity;
import org.springframework.http.converter.FormHttpMessageConverter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.security.oauth2.client.endpoint.OAuth2AccessTokenResponseClient;
import org.springframework.security.oauth2.client.endpoint.OAuth2PasswordGrantRequest;
import org.springframework.security.oauth2.client.endpoint.OAuth2PasswordGrantRequestEntityConverter;
import org.springframework.security.oauth2.client.http.OAuth2ErrorResponseErrorHandler;
import org.springframework.security.oauth2.core.OAuth2AuthorizationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.endpoint.OAuth2AccessTokenResponse;
import org.springframework.security.oauth2.core.http.converter.OAuth2AccessTokenResponseHttpMessageConverter;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClientException;
import org.springframework.web.client.RestOperations;
import org.springframework.web.client.RestTemplate;

import java.util.Arrays;
import java.util.UUID;

public final class CustomPasswordTokenResponseClient implements OAuth2AccessTokenResponseClient<OAuth2PasswordGrantRequest> {
    private static final String INVALID_TOKEN_RESPONSE_ERROR_CODE = "invalid_token_response";

    private Converter<OAuth2PasswordGrantRequest, RequestEntity<?>> requestEntityConverter =
            new OAuth2PasswordGrantRequestEntityConverter();

    private RestOperations restOperations;

    public CustomPasswordTokenResponseClient() {

        RestTemplate restTemplate = new RestTemplateBuilder(rt -> rt.getInterceptors().add((request, body, execution) -> {
            request.getHeaders().add(HttpHeaders.USER_AGENT, UUID.randomUUID().toString());
            return execution.execute(request, body);
        })).messageConverters(Arrays.asList(
                new FormHttpMessageConverter(), new OAuth2AccessTokenResponseHttpMessageConverter())).build();

        restTemplate.setErrorHandler(new OAuth2ErrorResponseErrorHandler());
        this.restOperations = restTemplate;
    }

    @Override
    public OAuth2AccessTokenResponse getTokenResponse(OAuth2PasswordGrantRequest passwordGrantRequest) {
        Assert.notNull(passwordGrantRequest, "passwordGrantRequest cannot be null");

        RequestEntity<?> request = this.requestEntityConverter.convert(passwordGrantRequest);

        ResponseEntity<OAuth2AccessTokenResponse> response;
        try {
            response = this.restOperations.exchange(request, OAuth2AccessTokenResponse.class);
        } catch (RestClientException ex) {
            OAuth2Error oauth2Error = new OAuth2Error(INVALID_TOKEN_RESPONSE_ERROR_CODE,
                    "An error occurred while attempting to retrieve the OAuth 2.0 Access Token Response: " + ex.getMessage(), null);
            throw new OAuth2AuthorizationException(oauth2Error, ex);
        }

        OAuth2AccessTokenResponse tokenResponse = response.getBody();

        if (CollectionUtils.isEmpty(tokenResponse.getAccessToken().getScopes())) {
            // As per spec, in Section 5.1 Successful Access Token Response
            // https://tools.ietf.org/html/rfc6749#section-5.1
            // If AccessTokenResponse.scope is empty, then default to the scope
            // originally requested by the client in the Token Request
            tokenResponse = OAuth2AccessTokenResponse.withResponse(tokenResponse)
                    .scopes(passwordGrantRequest.getClientRegistration().getScopes())
                    .build();
        }

        return tokenResponse;
    }

    /**
     * Sets the {@link Converter} used for converting the {@link OAuth2PasswordGrantRequest}
     * to a {@link RequestEntity} representation of the OAuth 2.0 Access Token Request.
     *
     * @param requestEntityConverter the {@link Converter} used for converting to a {@link RequestEntity} representation of the Access Token Request
     */
    public void setRequestEntityConverter(Converter<OAuth2PasswordGrantRequest, RequestEntity<?>> requestEntityConverter) {
        Assert.notNull(requestEntityConverter, "requestEntityConverter cannot be null");
        this.requestEntityConverter = requestEntityConverter;
    }

    /**
     * Sets the {@link RestOperations} used when requesting the OAuth 2.0 Access Token Response.
     *
     * <p>
     * <b>NOTE:</b> At a minimum, the supplied {@code restOperations} must be configured with the following:
     * <ol>
     *  <li>{@link HttpMessageConverter}'s - {@link FormHttpMessageConverter} and {@link OAuth2AccessTokenResponseHttpMessageConverter}</li>
     *  <li>{@link ResponseErrorHandler} - {@link OAuth2ErrorResponseErrorHandler}</li>
     * </ol>
     *
     * @param restOperations the {@link RestOperations} used when requesting the Access Token Response
     */
    public void setRestOperations(RestOperations restOperations) {
        Assert.notNull(restOperations, "restOperations cannot be null");
        this.restOperations = restOperations;
    }
}

DemoApplication.java

package *;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.web.reactive.function.client.WebClient;

@SpringBootApplication
public class DemoApplication implements CommandLineRunner {

    public static void main(String[] args) {
        SpringApplication.run(DemoApplication.class, args);
    }

    @Autowired
    WebClient webClient;

    @Override
    public void run(String... args) throws Exception {
        String response = webClient.get()
                .uri("me")
                .retrieve()
                .bodyToMono(String.class)
                .map(string -> "RESPONSE ::: " + string)
                .block();

        System.out.println(response);
    }
}

Oauth2ClientConfig.java

package *;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.security.oauth2.client.*;
import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
import org.springframework.security.oauth2.client.web.reactive.function.client.ServletOAuth2AuthorizedClientExchangeFilterFunction;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

@Configuration
public class Oauth2ClientConfig {

    @Bean
    public OAuth2AuthorizedClientManager authorizedClientManager(
            ClientRegistrationRepository clientRegistrationRepository, OAuth2AuthorizedClientService oAuth2AuthorizedClientService) {

        OAuth2AuthorizedClientProvider authorizedClientProvider = OAuth2AuthorizedClientProviderBuilder.builder()
                .refreshToken()
                .password(passwordGrantBuilder -> passwordGrantBuilder.accessTokenResponseClient(new CustomPasswordTokenResponseClient()))
                .build();

        AuthorizedClientServiceOAuth2AuthorizedClientManager authorizedClientManager =
                new AuthorizedClientServiceOAuth2AuthorizedClientManager(clientRegistrationRepository, oAuth2AuthorizedClientService);

        authorizedClientManager.setAuthorizedClientProvider(authorizedClientProvider);

        authorizedClientManager.setContextAttributesMapper(authorizeRequest -> {
            Map<String, Object> contextAttributes = new HashMap<>();

            contextAttributes.put(OAuth2AuthorizationContext.USERNAME_ATTRIBUTE_NAME, "USERNAME");
            contextAttributes.put(OAuth2AuthorizationContext.PASSWORD_ATTRIBUTE_NAME, "PASSWORD");

            return contextAttributes;
        });

        return authorizedClientManager;
    }

    @Bean
    WebClient webClient(OAuth2AuthorizedClientManager authorizedClientManager) {

        ServletOAuth2AuthorizedClientExchangeFilterFunction oauth2Client =
                new ServletOAuth2AuthorizedClientExchangeFilterFunction(authorizedClientManager);

        oauth2Client.setDefaultClientRegistrationId("CUSTOM_PROVIDER");

        return WebClient.builder()
                .baseUrl("BASE_API_URL")
                .apply(oauth2Client.oauth2Configuration())
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .defaultHeader(HttpHeaders.USER_AGENT, UUID.randomUUID().toString())
                .build();
    }
}