spring-cloud · marcingrzejszczak · Oct 27, 2016 · Oct 21, 2016
diff --git a/docs/src/main/asciidoc/spring-cloud-sleuth.adoc b/docs/src/main/asciidoc/spring-cloud-sleuth.adoc
@@ -228,20 +228,14 @@ to your bean definition.
 
 === HTTP
 
-For HTTP these are the beans responsible for creation of a Span from a `HttpServletRequest`
-  and filling in the `HttpServletResponse` with tracing information.
+For HTTP these are the beans responsible for creation of a Span from a `HttpServletRequest`.
 
 [source,java]
 ----
 @Bean
 public SpanExtractor<HttpServletRequest> httpServletRequestSpanExtractor() {
     ...
 }
-
-@Bean
-public SpanInjector<HttpServletResponse> httpServletResponseSpanInjector() {
-    ...
-}
 ----
 
 You can override them by providing your own implementation and by adding a `@Primary` annotation
@@ -262,18 +256,25 @@ This is a an example of a `SpanExtractor`
 include::../../../..//spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterCustomExtractorTests.java[tags=extractor,indent=0]
 ----
 
-The following `SpanInjector` could be created
+And you could register it like this:
 
 [source,java]
 ----
-include::../../../..//spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterCustomExtractorTests.java[tags=injector,indent=0]
+include::../../../..//spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterCustomExtractorTests.java[tags=configuration,indent=0]
 ----
 
-And you could register them like this:
+Spring Cloud Sleuth does not add trace/span related headers to the Http Response for security reasons. If you need the headers then a custom `SpanInjector`
+that injects the headers into the Http Response and a Servlet filter which makes use of this can be added the following way:
 
 [source,java]
 ----
-include::../../../..//spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterCustomExtractorTests.java[tags=configuration,indent=0]
+include::../../../..//spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceCustomFilterResponseInjectorTests.java[tags=injector,indent=0]
+----
+
+And you could register them like this:
+[source,java]
+----
+include::../../../..//spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceCustomFilterResponseInjectorTests.java[tags=configuration,indent=0]
 ----
 
 === Custom SA tag in Zipkin

diff --git a/...g/springframework/cloud/sleuth/instrument/web/TraceCustomFilterResponseInjectorTests.java b/...g/springframework/cloud/sleuth/instrument/web/TraceCustomFilterResponseInjectorTests.java
@@ -0,0 +1,158 @@
+/*
+ * Copyright 2013-2016 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.web;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.context.embedded.EmbeddedServletContainerInitializedEvent;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanInjector;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.RequestEntity;
+import org.springframework.http.ResponseEntity;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.web.bind.annotation.RequestHeader;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.client.RestTemplate;
+import org.springframework.web.filter.GenericFilterBean;
+
+import javax.servlet.FilterChain;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.BDDAssertions.then;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringBootTest(classes = TraceCustomFilterResponseInjectorTests.Config.class,
+		webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
+@DirtiesContext
+public class TraceCustomFilterResponseInjectorTests {
+	@Autowired RestTemplate restTemplate;
+	@Autowired Config config;
+	@Autowired CustomRestController customRestController;
+
+
+	@Test
+	@SuppressWarnings("unchecked")
+	public void should_inject_trace_and_span_ids_in_response_headers() {
+		RequestEntity<?> requestEntity = RequestEntity
+				.get(URI.create("http://localhost:" + this.config.port + "/headers"))
+				.build();
+
+		@SuppressWarnings("rawtypes")
+		ResponseEntity<Map> responseEntity = this.restTemplate.exchange(requestEntity, Map.class);
+
+		then(responseEntity.getHeaders())
+				.containsKeys(Span.TRACE_ID_NAME, Span.SPAN_ID_NAME)
+				.as("Trace headers must be present in response headers");
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	static class Config
+			implements ApplicationListener<EmbeddedServletContainerInitializedEvent> {
+		int port;
+
+		// tag::configuration[]
+		@Bean
+		SpanInjector<HttpServletResponse> customHttpServletResponseSpanInjector() {
+			return new CustomHttpServletResponseSpanInjector();
+		}
+
+		@Bean
+		HttpResponseInjectingTraceFilter responseInjectingTraceFilter(Tracer tracer) {
+			return new HttpResponseInjectingTraceFilter(tracer, customHttpServletResponseSpanInjector());
+		}
+		// end::configuration[]
+
+		@Override
+		public void onApplicationEvent(EmbeddedServletContainerInitializedEvent event) {
+			this.port = event.getEmbeddedServletContainer().getPort();
+		}
+
+		@Bean
+		public RestTemplate restTemplate() {
+			return new RestTemplate();
+		}
+
+		@Bean
+		CustomRestController customRestController() {
+			return new CustomRestController();
+		}
+
+
+	}
+
+	// tag::injector[]
+	static class CustomHttpServletResponseSpanInjector
+			implements SpanInjector<HttpServletResponse> {
+
+		@Override
+		public void inject(Span span, HttpServletResponse carrier) {
+			carrier.addHeader(Span.TRACE_ID_NAME, Span.idToHex(span.getTraceId()));
+			carrier.addHeader(Span.SPAN_ID_NAME, Span.idToHex(span.getSpanId()));
+		}
+	}
+
+	static class HttpResponseInjectingTraceFilter extends GenericFilterBean {
+
+		private final Tracer tracer;
+		private final SpanInjector<HttpServletResponse> spanInjector;
+
+		public HttpResponseInjectingTraceFilter(Tracer tracer, SpanInjector<HttpServletResponse> spanInjector) {
+			this.tracer = tracer;
+			this.spanInjector = spanInjector;
+		}
+
+		@Override
+		public void doFilter(ServletRequest request, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
+			HttpServletResponse response = (HttpServletResponse) servletResponse;
+			Span currentSpan = this.tracer.getCurrentSpan();
+			this.spanInjector.inject(currentSpan, response);
+			filterChain.doFilter(request, response);
+		}
+	}
+	// end::injector[]
+
+	@RestController
+	static class CustomRestController {
+
+		@RequestMapping("/headers")
+		public Map<String, String> headers(@RequestHeader HttpHeaders headers) {
+			Map<String, String> map = new HashMap<>();
+			for (String key : headers.keySet()) {
+				map.put(key, headers.getFirst(key));
+			}
+			return map;
+		}
+	}
+}
diff --git a/...java/org/springframework/cloud/sleuth/instrument/web/TraceFilterCustomExtractorTests.java b/...java/org/springframework/cloud/sleuth/instrument/web/TraceFilterCustomExtractorTests.java
@@ -16,13 +16,6 @@
 
 package org.springframework.cloud.sleuth.instrument.web;
 
-import java.net.URI;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Random;
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -33,7 +26,6 @@
 import org.springframework.cloud.sleuth.Sampler;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanExtractor;
-import org.springframework.cloud.sleuth.SpanInjector;
 import org.springframework.cloud.sleuth.SpanReporter;
 import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
 import org.springframework.cloud.sleuth.util.ArrayListSpanAccumulator;
@@ -51,6 +43,12 @@
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.client.RestTemplate;
 
+import javax.servlet.http.HttpServletRequest;
+import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+
 import static com.jayway.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.BDDAssertions.then;
@@ -83,13 +81,13 @@ public void should_create_a_valid_span_from_custom_headers() {
 				.header("mySpanId", Span.idToHex(spanId)).build();
 
 		@SuppressWarnings("rawtypes")
-		ResponseEntity<Map> responseHeaders = this.restTemplate.exchange(requestEntity,
+		ResponseEntity<Map> responseEntity = this.restTemplate.exchange(requestEntity,
 				Map.class);
 
 		await().atMost(5, SECONDS).until(() -> then(this.accumulator.getSpans().stream().filter(
 				span -> span.getSpanId() == spanId).findFirst().get())
 				.hasTraceIdEqualTo(traceId));
-		then(responseHeaders.getBody())
+		then(responseEntity.getBody())
 				.containsEntry("correlationid", Span.idToHex(traceId))
 				.containsEntry("myspanid", Span.idToHex(spanId))
 				.as("input request headers");
@@ -107,12 +105,6 @@ static class Config
 		SpanExtractor<HttpServletRequest> customHttpServletRequestSpanExtractor() {
 			return new CustomHttpServletRequestSpanExtractor();
 		}
-
-		@Bean
-		@Primary
-		SpanInjector<HttpServletResponse> customHttpServletResponseSpanInjector() {
-			return new CustomHttpServletResponseSpanInjector();
-		}
 		// end::configuration[]
 
 		@Override
@@ -157,19 +149,6 @@ public Span joinTrace(HttpServletRequest carrier) {
 	}
 	// end::extractor[]
 
-	// tag::injector[]
-	static class CustomHttpServletResponseSpanInjector
-			implements SpanInjector<HttpServletResponse> {
-
-		@Override
-		public void inject(Span span, HttpServletResponse carrier) {
-			carrier.addHeader("correlationId", Span.idToHex(span.getTraceId()));
-			carrier.addHeader("mySpanId", Span.idToHex(span.getSpanId()));
-			// inject the rest of Span values to the header
-		}
-	}
-	// end::injector[]
-
 	@RestController
 	static class CustomRestController {