From 6d25738ec41aab2d2ebe3c152be7271a3dc57649 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Wed, 23 Aug 2017 18:02:59 +0200
Subject: [PATCH] Adding the default Trace representation of
 UserInfoRestTemplateCustomizer

---
 .../TraceWebClientAutoConfiguration.java      | 47 ++++++++++++++-----
 1 file changed, 36 insertions(+), 11 deletions(-)

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 076fa5015..e6f22c701 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
@@ -86,15 +86,22 @@ public class TraceWebClientAutoConfiguration {
 	}
 
 	@Configuration
-	@ConditionalOnClass(UserInfoRestTemplateCustomizer.class)
+	@ConditionalOnClass({ UserInfoRestTemplateCustomizer.class, OAuth2RestTemplate.class })
 	protected static class TraceOAuthConfiguration {
 
 		@Autowired BeanFactory beanFactory;
 
-		@Bean UserInfoRestTemplateCustomizerBPP userInfoRestTemplateCustomizerBeanPostProcessor() {
+		@Bean
+		UserInfoRestTemplateCustomizerBPP userInfoRestTemplateCustomizerBeanPostProcessor() {
 			return new UserInfoRestTemplateCustomizerBPP(this.beanFactory);
 		}
 
+		@Bean
+		@ConditionalOnMissingBean
+		UserInfoRestTemplateCustomizer traceUserInfoRestTemplateCustomizer() {
+			return new TraceUserInfoRestTemplateCustomizer(this.beanFactory);
+		}
+
 		class UserInfoRestTemplateCustomizerBPP implements BeanPostProcessor {
 
 			private final BeanFactory beanFactory;
@@ -113,15 +120,9 @@ public class TraceWebClientAutoConfiguration {
 			public Object postProcessAfterInitialization(final Object bean,
 					String beanName) throws BeansException {
 				final BeanFactory beanFactory = this.beanFactory;
-				if (bean instanceof UserInfoRestTemplateCustomizer) {
-					return new UserInfoRestTemplateCustomizer() {
-						@Override public void customize(OAuth2RestTemplate template) {
-							final TraceRestTemplateInterceptor interceptor =
-									beanFactory.getBean(TraceRestTemplateInterceptor.class);
-							new RestTemplateInterceptorInjector(interceptor).inject(template);
-							((UserInfoRestTemplateCustomizer) bean).customize(template);
-						}
-					};
+				if (bean instanceof UserInfoRestTemplateCustomizer &&
+						!(bean instanceof TraceUserInfoRestTemplateCustomizer)) {
+					return new TraceUserInfoRestTemplateCustomizer(beanFactory, bean);
 				}
 				return bean;
 			}
@@ -142,4 +143,28 @@ class RestTemplateInterceptorInjector {
 		interceptors.add(this.interceptor);
 		restTemplate.setInterceptors(interceptors);
 	}
+}
+
+class TraceUserInfoRestTemplateCustomizer implements UserInfoRestTemplateCustomizer {
+
+	private final BeanFactory beanFactory;
+	private final Object delegate;
+
+	TraceUserInfoRestTemplateCustomizer(BeanFactory beanFactory) {
+		this(beanFactory, null);
+	}
+
+	TraceUserInfoRestTemplateCustomizer(BeanFactory beanFactory, Object bean) {
+		this.beanFactory = beanFactory;
+		this.delegate = bean;
+	}
+
+	@Override public void customize(OAuth2RestTemplate template) {
+		final TraceRestTemplateInterceptor interceptor =
+				this.beanFactory.getBean(TraceRestTemplateInterceptor.class);
+		new RestTemplateInterceptorInjector(interceptor).inject(template);
+		if (this.delegate != null) {
+			((UserInfoRestTemplateCustomizer) this.delegate).customize(template);
+		}
+	}
 }
\ No newline at end of file