From 71c8a114f121186bcbb643e56639f3f18defb606 Mon Sep 17 00:00:00 2001
From: Johnny Lim <izeye@naver.com>
Date: Tue, 20 Oct 2015 11:57:04 +0900
Subject: [PATCH] Make response headers optional in /trace endpoint

Update `WebRequestTraceFilter` so that the response respects the
TraceProperties.Include.RESPONSE enum.

See gh-3948
Closes gh-4242
---
 .../actuate/trace/WebRequestTraceFilter.java  |  6 +++--
 .../trace/WebRequestTraceFilterTests.java     | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java
index bcd4b80de0..02312c342b 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/trace/WebRequestTraceFilter.java
@@ -170,8 +170,10 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 
 	@SuppressWarnings("unchecked")
 	protected void enhanceTrace(Map<String, Object> trace, HttpServletResponse response) {
-		Map<String, Object> headers = (Map<String, Object>) trace.get("headers");
-		headers.put("response", getResponseHeaders(response));
+		if (isIncluded(Include.RESPONSE_HEADERS)) {
+			Map<String, Object> headers = (Map<String, Object>) trace.get("headers");
+			headers.put("response", getResponseHeaders(response));
+		}
 	}
 
 	private Map<String, String> getResponseHeaders(HttpServletResponse response) {
diff --git a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java
index a900d0e5dd..c122f50c25 100644
--- a/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java
+++ b/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/trace/WebRequestTraceFilterTests.java
@@ -20,6 +20,7 @@ import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
 import java.security.Principal;
+import java.util.Collections;
 import java.util.EnumSet;
 import java.util.Map;
 
@@ -37,6 +38,7 @@ import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockHttpServletResponse;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Tests for {@link WebRequestTraceFilter}.
@@ -126,6 +128,25 @@ public class WebRequestTraceFilterTests {
 		assertEquals("{Accept=application/json}", map.get("request").toString());
 	}
 
+	@Test
+	@SuppressWarnings({ "unchecked" })
+	public void filterDoesNotAddResponseHeadersWithoutResponseHeadersInclude()
+			throws ServletException, IOException {
+		this.properties.setInclude(Collections.singleton(Include.REQUEST_HEADERS));
+		MockHttpServletRequest request = new MockHttpServletRequest("GET", "/foo");
+		MockHttpServletResponse response = new MockHttpServletResponse();
+		response.addHeader("Content-Type", "application/json");
+		this.filter.doFilterInternal(request, response, new FilterChain() {
+			@Override
+			public void doFilter(ServletRequest request, ServletResponse response)
+					throws IOException, ServletException {
+			}
+		});
+		Map<String, Object> info = this.repository.findAll().iterator().next().getInfo();
+		Map<String, Object> headers = (Map<String, Object>) info.get("headers");
+		assertTrue(headers.get("response") == null);
+	}
+
 	@Test
 	public void filterHasResponseStatus() {
 		MockHttpServletRequest request = new MockHttpServletRequest("GET", "/foo");
@@ -156,4 +177,5 @@ public class WebRequestTraceFilterTests {
 		System.err.println(map);
 		assertEquals("Foo", map.get("message").toString());
 	}
+
 }