From 9da6ddb253cd8f66f8951adbf86039526e4a3c8b Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <awilkinson@pivotal.io>
Date: Mon, 29 Jan 2018 15:24:24 +0000
Subject: [PATCH] Trace IDs of sessions created downstream of trace filter

Closes gh-11717
---
 .../actuate/trace/WebRequestTraceFilter.java  | 13 +++++++---
 .../trace/WebRequestTraceFilterTests.java     | 26 +++++++++----------
 2 files changed, 22 insertions(+), 17 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 a5f35701ba..522a7d7a97 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2017 the original author or authors.
+ * Copyright 2012-2018 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.
@@ -112,14 +112,21 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 		}
 		finally {
 			addTimeTaken(trace, startTime);
+			addSessionIdIfNecessary(request, trace);
 			enhanceTrace(trace, status == response.getStatus() ? response
 					: new CustomStatusResponseWrapper(response, status));
 			this.repository.add(trace);
 		}
 	}
 
-	protected Map<String, Object> getTrace(HttpServletRequest request) {
+	private void addSessionIdIfNecessary(HttpServletRequest request,
+			Map<String, Object> trace) {
 		HttpSession session = request.getSession(false);
+		add(trace, Include.SESSION_ID, "sessionId",
+				(session == null ? null : session.getId()));
+	}
+
+	protected Map<String, Object> getTrace(HttpServletRequest request) {
 		Throwable exception = (Throwable) request
 				.getAttribute("javax.servlet.error.exception");
 		Principal userPrincipal = request.getUserPrincipal();
@@ -143,8 +150,6 @@ public class WebRequestTraceFilter extends OncePerRequestFilter implements Order
 		add(trace, Include.QUERY_STRING, "query", request.getQueryString());
 		add(trace, Include.AUTH_TYPE, "authType", request.getAuthType());
 		add(trace, Include.REMOTE_ADDRESS, "remoteAddress", request.getRemoteAddr());
-		add(trace, Include.SESSION_ID, "sessionId",
-				(session == null ? null : session.getId()));
 		add(trace, Include.REMOTE_USER, "remoteUser", request.getRemoteUser());
 		if (isIncluded(Include.ERRORS) && exception != null
 				&& this.errorAttributes != null) {
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 56f6a8db97..997fb461e5 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2017 the original author or authors.
+ * Copyright 2012-2018 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.
@@ -16,7 +16,6 @@
 
 package org.springframework.boot.actuate.trace;
 
-import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
 import java.security.Principal;
@@ -28,6 +27,9 @@ import javax.servlet.FilterChain;
 import javax.servlet.ServletException;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 
 import org.junit.Test;
 
@@ -105,19 +107,16 @@ public class WebRequestTraceFilterTests {
 		MockHttpServletResponse response = new MockHttpServletResponse();
 		response.addHeader("Content-Type", "application/json");
 		response.addHeader("Set-Cookie", "a=b");
-		this.filter.doFilterInternal(request, response, new FilterChain() {
+		this.filter.doFilterInternal(request, response,
+				new MockFilterChain(new HttpServlet() {
 
-			@Override
-			public void doFilter(ServletRequest request, ServletResponse response)
-					throws IOException, ServletException {
-				BufferedReader bufferedReader = request.getReader();
-				while (bufferedReader.readLine() != null) {
-					// read the contents as normal (forces cache to fill up)
-				}
-				response.getWriter().println("Goodbye, World!");
-			}
+					@Override
+					protected void doGet(HttpServletRequest req, HttpServletResponse resp)
+							throws ServletException, IOException {
+						req.getSession(true);
+					}
 
-		});
+				}));
 		assertThat(this.repository.findAll()).hasSize(1);
 		Map<String, Object> trace = this.repository.findAll().iterator().next().getInfo();
 		Map<String, Object> map = (Map<String, Object>) trace.get("headers");
@@ -136,6 +135,7 @@ public class WebRequestTraceFilterTests {
 		assertThat(trace.get("authType")).isEqualTo("authType");
 		assertThat(map.get("request").toString())
 				.isEqualTo("{Accept=application/json, Cookie=testCookie=testValue;}");
+		assertThat(trace).containsKey("sessionId");
 	}
 
 	@Test