Polish "Add ability to filter cookies in trace data"

Closes gh-6018

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

@@ -46,7 +46,8 @@ public class TraceProperties {
 	}
 
 	/**
-	 * Items to be included in the trace. Defaults to request/response headers and errors.
+	 * Items to be included in the trace. Defaults to request/response headers (including cookies)
+	 * and errors.
 	 */
 	private Set<Include> include = new HashSet<Include>(DEFAULT_INCLUDES);
 
@@ -74,7 +75,7 @@ public class TraceProperties {
 		RESPONSE_HEADERS,
 
 		/**
-		 * Include Cookie in request and Set-Cookie in response headers.
+		 * Include "Cookie" in request and "Set-Cookie" in response headers.
 		 */
 		COOKIES,
 

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

@@ -123,11 +123,7 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 		trace.put("path", request.getRequestURI());
 		trace.put("headers", headers);
 		if (isIncluded(Include.REQUEST_HEADERS)) {
-			Map<String, Object> requestHeaders = getRequestHeaders(request);
-			if (!isIncluded(Include.COOKIES)) {
-				requestHeaders.remove("Cookie");
-			}
-			headers.put("request", requestHeaders);
+			headers.put("request", getRequestHeaders(request));
 		}
 		add(trace, Include.PATH_INFO, "pathInfo", request.getPathInfo());
 		add(trace, Include.PATH_TRANSLATED, "pathTranslated",
@@ -167,6 +163,9 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 			}
 			headers.put(name, value);
 		}
+		if (!isIncluded(Include.COOKIES)) {
+			headers.remove("Cookie");
+		}
 		return headers;
 	}
 
@@ -174,11 +173,7 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 	protected void enhanceTrace(Map<String, Object> trace, HttpServletResponse response) {
 		if (isIncluded(Include.RESPONSE_HEADERS)) {
 			Map<String, Object> headers = (Map<String, Object>) trace.get("headers");
-			Map<String, String> responseHeaders = getResponseHeaders(response);
-			if (!isIncluded(Include.COOKIES)) {
-				responseHeaders.remove("Set-Cookie");
-			}
-			headers.put("response", responseHeaders);
+			headers.put("response", getResponseHeaders(response));
 		}
 	}
 
@@ -188,6 +183,9 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 			String value = response.getHeader(header);
 			headers.put(header, value);
 		}
+		if (!isIncluded(Include.COOKIES)) {
+			headers.remove("Set-Cookie");
+		}
 		headers.put("status", "" + response.getStatus());
 		return headers;
 	}

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

@@ -28,7 +28,6 @@ import javax.servlet.FilterChain;
 import javax.servlet.ServletException;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
-import javax.servlet.http.Cookie;
 
 import org.junit.Test;
 
@@ -51,6 +50,7 @@ import static org.mockito.Mockito.verify;
  * @author Phillip Webb
  * @author Andy Wilkinson
  * @author Venil Noronha
+ * @author Stephane Nicoll
  */
 public class WebRequestTraceFilterTests {
 
@@ -80,6 +80,7 @@ public class WebRequestTraceFilterTests {
 		this.properties.setInclude(EnumSet.allOf(Include.class));
 		MockHttpServletRequest request = new MockHttpServletRequest("GET", "/foo");
 		request.addHeader("Accept", "application/json");
+		request.addHeader("Cookie", "testCookie=testValue;");
 		request.setContextPath("some.context.path");
 		request.setContent("Hello, World!".getBytes());
 		request.setRemoteAddr("some.remote.addr");
@@ -89,8 +90,6 @@ public class WebRequestTraceFilterTests {
 		String url = tmp.toURI().toURL().toString();
 		request.setPathInfo(url);
 		tmp.deleteOnExit();
-		Cookie cookie = new Cookie("testCookie", "testValue");
-		request.setCookies(cookie);
 		request.setAuthType("authType");
 		Principal principal = new Principal() {
 
@@ -103,6 +102,7 @@ public class WebRequestTraceFilterTests {
 		request.setUserPrincipal(principal);
 		MockHttpServletResponse response = new MockHttpServletResponse();
 		response.addHeader("Content-Type", "application/json");
+		response.addHeader("Set-Cookie", "a=b");
 		this.filter.doFilterInternal(request, response, new FilterChain() {
 
 			@Override
@@ -121,7 +121,7 @@ public class WebRequestTraceFilterTests {
 		Map<String, Object> map = (Map<String, Object>) trace.get("headers");
 
 		assertThat(map.get("response").toString())
-				.isEqualTo("{Content-Type=application/json, status=200}");
+				.isEqualTo("{Content-Type=application/json, Set-Cookie=a=b, status=200}");
 		assertThat(trace.get("method")).isEqualTo("GET");
 		assertThat(trace.get("path")).isEqualTo("/foo");
 		assertThat(((String[]) ((Map) trace.get("parameters")).get("param"))[0])
@@ -132,7 +132,8 @@ public class WebRequestTraceFilterTests {
 		assertThat(trace.get("contextPath")).isEqualTo("some.context.path");
 		assertThat(trace.get("pathInfo")).isEqualTo(url);
 		assertThat(trace.get("authType")).isEqualTo("authType");
-		assertThat(map.get("request").toString()).isEqualTo("{Accept=application/json}");
+		assertThat(map.get("request").toString())
+				.isEqualTo("{Accept=application/json, Cookie=testCookie=testValue;}");
 	}
 
 	@Test

spring-boot-docs/src/main/asciidoc/appendix-application-properties.adoc

@@ -1072,7 +1072,7 @@ content into your application; rather pick only the properties that you need.
 	management.shell.telnet.port=5000 # Telnet port.
 
 	# TRACING ({sc-spring-boot-actuator}/trace/TraceProperties.{sc-ext}[TraceProperties])
-	management.trace.include=request-headers,response-headers,errors # Items to be included in the trace.
+	management.trace.include=request-headers,response-headers,cookies,errors # Items to be included in the trace.
 
 	# METRICS EXPORT ({sc-spring-boot-actuator}/metrics/export/MetricExportProperties.{sc-ext}[MetricExportProperties])
 	spring.metrics.export.aggregate.key-pattern= # Pattern that tells the aggregator what to do with the keys from the source repository.