From 0d136893a3959cf5016e76dc86b882ba0cec0d7b Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Sat, 3 Mar 2018 19:51:05 +0100
Subject: [PATCH] Added UserInfoRestTemplateCustomizer support; fixes gh-879

---
 pom.xml                                       |  7 ++
 spring-cloud-sleuth-core/pom.xml              |  5 ++
 .../TraceWebClientAutoConfiguration.java      | 69 +++++++++++++++++++
 3 files changed, 81 insertions(+)

diff --git a/pom.xml b/pom.xml
index 895d9095c..2f20a1360 100644
--- a/pom.xml
+++ b/pom.xml
@@ -212,6 +212,12 @@
 				<version>${spock.version}</version>
 				<scope>test</scope>
 			</dependency>
+			<dependency>
+				<groupId>org.springframework.security.oauth.boot</groupId>
+				<artifactId>spring-security-oauth2-autoconfigure</artifactId>
+				<version>${spring-security-boot-autoconfigure.version}</version>
+				<optional>true</optional>
+			</dependency>
 			<dependency>
 				<groupId>cglib</groupId>
 				<artifactId>cglib-nodep</artifactId>
@@ -267,6 +273,7 @@
 		<spring-cloud-netflix.version>2.0.0.BUILD-SNAPSHOT</spring-cloud-netflix.version>
 		<spring-cloud-openfeign.version>2.0.0.BUILD-SNAPSHOT</spring-cloud-openfeign.version>
 		<brave.version>4.17.1</brave.version>
+		<spring-security-boot-autoconfigure.version>2.0.0.RELEASE</spring-security-boot-autoconfigure.version>
 	</properties>
 
 	<profiles>
diff --git a/spring-cloud-sleuth-core/pom.xml b/spring-cloud-sleuth-core/pom.xml
index 67e1cca3f..1f766b3f2 100644
--- a/spring-cloud-sleuth-core/pom.xml
+++ b/spring-cloud-sleuth-core/pom.xml
@@ -186,6 +186,11 @@
 			<artifactId>httpasyncclient</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework.security.oauth.boot</groupId>
+			<artifactId>spring-security-oauth2-autoconfigure</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<!-- BRAVE -->
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java
index 342761488..556d57389 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java
@@ -52,6 +52,7 @@ import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.security.oauth2.resource.UserInfoRestTemplateCustomizer;
 import org.springframework.boot.web.client.RestTemplateCustomizer;
 import org.springframework.cloud.commons.httpclient.HttpClientConfiguration;
 import org.springframework.cloud.sleuth.instrument.web.TraceWebServletAutoConfiguration;
@@ -62,6 +63,7 @@ import org.springframework.http.HttpRequest;
 import org.springframework.http.client.ClientHttpRequestExecution;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.http.client.ClientHttpResponse;
+import org.springframework.security.oauth2.client.OAuth2RestTemplate;
 import org.springframework.web.client.RestTemplate;
 import org.springframework.web.reactive.function.client.WebClient;
 import reactor.core.publisher.Mono;
@@ -146,6 +148,48 @@ public class TraceWebClientAutoConfiguration {
 			return new NettyAspect(httpTracing);
 		}
 	}
+
+	@Configuration
+	@ConditionalOnClass({ UserInfoRestTemplateCustomizer.class, OAuth2RestTemplate.class })
+	protected static class TraceOAuthConfiguration {
+
+		@Bean
+		UserInfoRestTemplateCustomizerBPP userInfoRestTemplateCustomizerBeanPostProcessor(BeanFactory beanFactory) {
+			return new UserInfoRestTemplateCustomizerBPP(beanFactory);
+		}
+
+		@Bean
+		@ConditionalOnMissingBean
+		UserInfoRestTemplateCustomizer traceUserInfoRestTemplateCustomizer(BeanFactory beanFactory) {
+			return new TraceUserInfoRestTemplateCustomizer(beanFactory);
+		}
+
+		private static class UserInfoRestTemplateCustomizerBPP implements BeanPostProcessor {
+
+			private final BeanFactory beanFactory;
+
+			UserInfoRestTemplateCustomizerBPP(BeanFactory beanFactory) {
+				this.beanFactory = beanFactory;
+			}
+
+			@Override
+			public Object postProcessBeforeInitialization(Object bean,
+					String beanName) throws BeansException {
+				return bean;
+			}
+
+			@Override
+			public Object postProcessAfterInitialization(final Object bean,
+					String beanName) throws BeansException {
+				final BeanFactory beanFactory = this.beanFactory;
+				if (bean instanceof UserInfoRestTemplateCustomizer &&
+						!(bean instanceof TraceUserInfoRestTemplateCustomizer)) {
+					return new TraceUserInfoRestTemplateCustomizer(beanFactory, bean);
+				}
+				return bean;
+			}
+		}
+	}
 }
 
 class RestTemplateInterceptorInjector {
@@ -374,3 +418,28 @@ class TracingHttpClientInstrumentation {
 		}
 	}
 }
+
+class TraceUserInfoRestTemplateCustomizer implements UserInfoRestTemplateCustomizer {
+
+	private final BeanFactory beanFactory;
+	private final Object delegate;
+
+	TraceUserInfoRestTemplateCustomizer(BeanFactory beanFactory) {
+		this.beanFactory = beanFactory;
+		this.delegate = null;
+	}
+
+	TraceUserInfoRestTemplateCustomizer(BeanFactory beanFactory, Object bean) {
+		this.beanFactory = beanFactory;
+		this.delegate = bean;
+	}
+
+	@Override public void customize(OAuth2RestTemplate template) {
+		final TracingClientHttpRequestInterceptor interceptor =
+				this.beanFactory.getBean(TracingClientHttpRequestInterceptor.class);
+		new RestTemplateInterceptorInjector(interceptor).inject(template);
+		if (this.delegate != null) {
+			((UserInfoRestTemplateCustomizer) this.delegate).customize(template);
+		}
+	}
+}