Allow WebRequestTraceFilter header post processing

Update WebRequestTraceFilter so that additional post processing can be
applied to traced request headers. The postProcessRequestHeaders method
can be used to remove or change map entries before they are returned.

Fixes gh-6309

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/TraceWebFilterAutoConfiguration.java

@@ -27,6 +27,7 @@ import org.springframework.boot.actuate.trace.WebRequestTraceFilter;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.web.ErrorAttributes;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
@@ -60,10 +61,10 @@ public class TraceWebFilterAutoConfiguration {
 	}
 
 	@Bean
+	@ConditionalOnMissingBean
 	public WebRequestTraceFilter webRequestLoggingFilter(BeanFactory beanFactory) {
 		WebRequestTraceFilter filter = new WebRequestTraceFilter(this.traceRepository,
 				this.traceProperties);
-
 		if (this.errorAttributes != null) {
 			filter.setErrorAttributes(this.errorAttributes);
 		}

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java

@@ -166,9 +166,18 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 		if (!isIncluded(Include.COOKIES)) {
 			headers.remove("Cookie");
 		}
+		postProcessRequestHeaders(headers);
 		return headers;
 	}
 
+	/**
+	 * Post process request headers before they are added to the trace.
+	 * @param headers a mutable map containing the request headers to trace
+	 * @since 1.4.0
+	 */
+	protected void postProcessRequestHeaders(Map<String, Object> headers) {
+	}
+
 	@SuppressWarnings("unchecked")
 	protected void enhanceTrace(Map<String, Object> trace, HttpServletResponse response) {
 		if (isIncluded(Include.RESPONSE_HEADERS)) {

spring-boot-actuator/src/test/java/org/springframework/boot/actuate/autoconfigure/TraceWebFilterAutoConfigurationTests.java

@@ -16,11 +16,17 @@
 
 package org.springframework.boot.actuate.autoconfigure;
 
+import java.util.Map;
+
 import org.junit.Test;
 
+import org.springframework.boot.actuate.trace.TraceProperties;
+import org.springframework.boot.actuate.trace.TraceRepository;
 import org.springframework.boot.actuate.trace.WebRequestTraceFilter;
 import org.springframework.boot.autoconfigure.PropertyPlaceholderAutoConfiguration;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -41,4 +47,40 @@ public class TraceWebFilterAutoConfigurationTests {
 		context.close();
 	}
 
+	@Test
+	public void overrideTraceFilter() throws Exception {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(
+				CustomTraceFilterConfig.class, PropertyPlaceholderAutoConfiguration.class,
+				TraceRepositoryAutoConfiguration.class,
+				TraceWebFilterAutoConfiguration.class);
+		WebRequestTraceFilter filter = context.getBean(WebRequestTraceFilter.class);
+		assertThat(filter).isInstanceOf(TestWebRequestTraceFilter.class);
+		context.close();
+	}
+
+	@Configuration
+	static class CustomTraceFilterConfig {
+
+		@Bean
+		public TestWebRequestTraceFilter testWebRequestTraceFilter(
+				TraceRepository repository, TraceProperties properties) {
+			return new TestWebRequestTraceFilter(repository, properties);
+		}
+
+	}
+
+	static class TestWebRequestTraceFilter extends WebRequestTraceFilter {
+
+		TestWebRequestTraceFilter(TraceRepository repository,
+				TraceProperties properties) {
+			super(repository, properties);
+		}
+
+		@Override
+		protected void postProcessRequestHeaders(Map<String, Object> headers) {
+			headers.clear();
+		}
+
+	}
+
 }

spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java

@@ -243,4 +243,23 @@ public class WebRequestTraceFilterTests {
 		}
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void postProcessRequestHeaders() throws Exception {
+		this.filter = new WebRequestTraceFilter(this.repository, this.properties) {
+
+			@Override
+			protected void postProcessRequestHeaders(Map<String, Object> headers) {
+				headers.remove("Test");
+			}
+
+		};
+		MockHttpServletRequest request = spy(new MockHttpServletRequest("GET", "/foo"));
+		request.addHeader("Accept", "application/json");
+		request.addHeader("Test", "spring");
+		Map<String, Object> map = (Map<String, Object>) this.filter.getTrace(request)
+				.get("headers");
+		assertThat(map.get("request").toString()).isEqualTo("{Accept=application/json}");
+	}
+
 }