Add ability to filter cookies in trace data

See gh-6018

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2015 the original author or authors.
+ * Copyright 2012-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.
@@ -28,6 +28,7 @@ import org.springframework.boot.context.properties.ConfigurationProperties;
  *
  * @author Wallace Wadge
  * @author Phillip Webb
+ * @author Venil Noronha
  * @since 1.3.0
  */
 @ConfigurationProperties(prefix = "management.trace")
@@ -39,6 +40,7 @@ public class TraceProperties {
 		Set<Include> defaultIncludes = new LinkedHashSet<Include>();
 		defaultIncludes.add(Include.REQUEST_HEADERS);
 		defaultIncludes.add(Include.RESPONSE_HEADERS);
+		defaultIncludes.add(Include.COOKIES);
 		defaultIncludes.add(Include.ERRORS);
 		DEFAULT_INCLUDES = Collections.unmodifiableSet(defaultIncludes);
 	}
@@ -71,6 +73,11 @@ public class TraceProperties {
 		 */
 		RESPONSE_HEADERS,
 
+		/**
+		 * Include Cookie in request and Set-Cookie in response headers.
+		 */
+		COOKIES,
+
 		/**
 		 * Include errors (if any).
 		 */

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

@@ -48,6 +48,7 @@ import org.springframework.web.filter.OncePerRequestFilter;
  * @author Dave Syer
  * @author Wallace Wadge
  * @author Andy Wilkinson
+ * @author Venil Noronha
  */
 public class WebRequestTraceFilter extends OncePerRequestFilter implements Ordered {
 
@@ -122,7 +123,11 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 		trace.put("path", request.getRequestURI());
 		trace.put("headers", headers);
 		if (isIncluded(Include.REQUEST_HEADERS)) {
-			headers.put("request", getRequestHeaders(request));
+			Map<String, Object> requestHeaders = getRequestHeaders(request);
+			if (!isIncluded(Include.COOKIES)) {
+				requestHeaders.remove("Cookie");
+			}
+			headers.put("request", requestHeaders);
 		}
 		add(trace, Include.PATH_INFO, "pathInfo", request.getPathInfo());
 		add(trace, Include.PATH_TRANSLATED, "pathTranslated",
@@ -169,7 +174,11 @@ 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");
-			headers.put("response", getResponseHeaders(response));
+			Map<String, String> responseHeaders = getResponseHeaders(response);
+			if (!isIncluded(Include.COOKIES)) {
+				responseHeaders.remove("Set-Cookie");
+			}
+			headers.put("response", responseHeaders);
 		}
 	}
 

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

@@ -50,6 +50,7 @@ import static org.mockito.Mockito.verify;
  * @author Wallace Wadge
  * @author Phillip Webb
  * @author Andy Wilkinson
+ * @author Venil Noronha
  */
 public class WebRequestTraceFilterTests {
 
@@ -153,6 +154,35 @@ public class WebRequestTraceFilterTests {
 		assertThat(headers.get("response") == null).isTrue();
 	}
 
+	@Test
+	@SuppressWarnings({ "unchecked" })
+	public void filterDoesNotAddRequestCookiesWithCookiesExclude()
+			throws ServletException, IOException {
+		this.properties.setInclude(Collections.singleton(Include.REQUEST_HEADERS));
+		MockHttpServletRequest request = spy(new MockHttpServletRequest("GET", "/foo"));
+		request.addHeader("Accept", "application/json");
+		request.addHeader("Cookie", "testCookie=testValue;");
+		Map<String, Object> map = (Map<String, Object>) this.filter.getTrace(request)
+				.get("headers");
+		assertThat(map.get("request").toString()).isEqualTo("{Accept=application/json}");
+	}
+
+	@Test
+	@SuppressWarnings({ "unchecked" })
+	public void filterDoesNotAddResponseCookiesWithCookiesExclude()
+			throws ServletException, IOException {
+		this.properties.setInclude(Collections.singleton(Include.RESPONSE_HEADERS));
+		MockHttpServletRequest request = new MockHttpServletRequest("GET", "/foo");
+		MockHttpServletResponse response = new MockHttpServletResponse();
+		response.addHeader("Content-Type", "application/json");
+		response.addHeader("Set-Cookie", "testCookie=testValue;");
+		Map<String, Object> trace = this.filter.getTrace(request);
+		this.filter.enhanceTrace(trace, response);
+		Map<String, Object> map = (Map<String, Object>) trace.get("headers");
+		assertThat(map.get("response").toString())
+				.isEqualTo("{Content-Type=application/json, status=200}");
+	}
+
 	@Test
 	public void filterHasResponseStatus() {
 		MockHttpServletRequest request = new MockHttpServletRequest("GET", "/foo");