From a53182f17640de81470c1b2e8522972c8d5c1469 Mon Sep 17 00:00:00 2001
From: SylvainJuge <763082+SylvainJuge@users.noreply.github.com>
Date: Wed, 14 Feb 2024 15:07:46 +0100
Subject: [PATCH] cleanup servlet request attribute (#3527)

---
 CHANGELOG.asciidoc                            |  4 ++++
 .../apm/agent/servlet/ServletApiAdvice.java   | 22 +++++++++++++++----
 .../adapter/JakartaServletApiAdapter.java     |  5 ++---
 .../adapter/JavaxServletApiAdapter.java       |  5 ++---
 .../adapter/ServletRequestAdapter.java        |  3 +--
 ...trumentationWithExceptionResolverTest.java | 13 +++++++++++
 6 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc
index 4d5a7a312e..589b635329 100644
--- a/CHANGELOG.asciidoc
+++ b/CHANGELOG.asciidoc
@@ -31,6 +31,10 @@ Use subheadings with the "=====" level for adding notes for unreleased changes:
 
 === Unreleased
 
+[float]
+===== Bug fixes
+* Cleanup extra servlet request attribute used for Spring exception handler - {pull}3527[#3527]
+
 [[release-notes-1.x]]
 === Java Agent version 1.x
 
diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletApiAdvice.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletApiAdvice.java
index fa41a89bb0..c52377be12 100644
--- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletApiAdvice.java
+++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/ServletApiAdvice.java
@@ -55,7 +55,15 @@ public abstract class ServletApiAdvice {
     private static final DetachedThreadLocal<Object> servletPathTL = GlobalVariables.get(ServletApiAdvice.class, "servletPath", WeakConcurrent.buildThreadLocal());
     private static final DetachedThreadLocal<Object> pathInfoTL = GlobalVariables.get(ServletApiAdvice.class, "pathInfo", WeakConcurrent.buildThreadLocal());
 
-    private static final List<String> requestExceptionAttributes = Arrays.asList("javax.servlet.error.exception", "jakarta.servlet.error.exception", "exception", "org.springframework.web.servlet.DispatcherServlet.EXCEPTION", "co.elastic.apm.exception");
+    private static final String ELASTIC_EXCEPTION = "co.elastic.apm.exception";
+    private static final String JAVAX_ERROR_EXCEPTION = "javax.servlet.error.exception";
+    private static final String JAKARTA_ERROR_EXCEPTION = "jakarta.servlet.error.exception";
+    private static final List<String> requestExceptionAttributes = Arrays.asList(
+        JAVAX_ERROR_EXCEPTION,
+        JAKARTA_ERROR_EXCEPTION,
+        "exception",
+        "org.springframework.web.servlet.DispatcherServlet.EXCEPTION",
+        ELASTIC_EXCEPTION);
 
     @Nullable
     public static <HttpServletRequest, HttpServletResponse, ServletContext, ServletContextEvent, FilterConfig, ServletConfig> Object onServletEnter(
@@ -202,7 +210,7 @@ public static <HttpServletRequest, HttpServletResponse, ServletContext, ServletC
             httpServletRequest != null &&
             httpServletResponse != null) {
 
-            if (adapter.getHttpAttribute(httpServletRequest, ServletTransactionHelper.ASYNC_ATTRIBUTE) != null) {
+            if (adapter.getAttribute(httpServletRequest, ServletTransactionHelper.ASYNC_ATTRIBUTE) != null) {
                 // HttpServletRequest.startAsync was invoked on this httpServletRequest.
                 // The transaction should be handled from now on by the other thread committing the response
                 transaction.deactivate();
@@ -230,10 +238,16 @@ public static <HttpServletRequest, HttpServletResponse, ServletContext, ServletC
                     final int size = requestExceptionAttributes.size();
                     for (int i = 0; i < size; i++) {
                         String attributeName = requestExceptionAttributes.get(i);
-                        Object throwable = adapter.getHttpAttribute(httpServletRequest, attributeName);
+                        Object throwable = adapter.getAttribute(httpServletRequest, attributeName);
                         if (throwable instanceof Throwable) {
                             t2 = (Throwable) throwable;
-                            if (!attributeName.equals("javax.servlet.error.exception") && !attributeName.equals("jakarta.servlet.error.exception")) {
+
+                            // elastic exception can be removed as it's not needed after transaction end
+                            if (attributeName.equals(ELASTIC_EXCEPTION)) {
+                                adapter.removeAttribute(httpServletRequest, attributeName);
+                            }
+
+                            if (!attributeName.equals(JAVAX_ERROR_EXCEPTION) && !attributeName.equals(JAKARTA_ERROR_EXCEPTION)) {
                                 overrideStatusCodeOnThrowable = false;
                             }
                             break;
diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JakartaServletApiAdapter.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JakartaServletApiAdapter.java
index 1775a803d9..a941725e7f 100644
--- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JakartaServletApiAdapter.java
+++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JakartaServletApiAdapter.java
@@ -242,10 +242,9 @@ public void setAttribute(HttpServletRequest servletRequest, String attributeName
         servletRequest.setAttribute(attributeName, value);
     }
 
-    @Nullable
     @Override
-    public Object getHttpAttribute(HttpServletRequest httpServletRequest, String attributeName) {
-        return httpServletRequest.getAttribute(attributeName);
+    public void removeAttribute(HttpServletRequest httpServletRequest, String attributeName) {
+        httpServletRequest.removeAttribute(attributeName);
     }
 
     @Override
diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JavaxServletApiAdapter.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JavaxServletApiAdapter.java
index 47a1fe9253..585745afec 100644
--- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JavaxServletApiAdapter.java
+++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/JavaxServletApiAdapter.java
@@ -256,10 +256,9 @@ public void setAttribute(HttpServletRequest servletRequest, String attributeName
         servletRequest.setAttribute(attributeName, value);
     }
 
-    @Nullable
     @Override
-    public Object getHttpAttribute(HttpServletRequest httpServletRequest, String attributeName) {
-        return httpServletRequest.getAttribute(attributeName);
+    public void removeAttribute(HttpServletRequest servletRequest, String attributeName) {
+        servletRequest.removeAttribute(attributeName);
     }
 
     @Override
diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/ServletRequestAdapter.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/ServletRequestAdapter.java
index 27ec9e6be0..c13625c587 100644
--- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/ServletRequestAdapter.java
+++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/adapter/ServletRequestAdapter.java
@@ -90,8 +90,7 @@ public interface ServletRequestAdapter<HttpServletRequest, ServletContext> {
 
     void setAttribute(HttpServletRequest request, String attributeName, Object value);
 
-    @Nullable
-    Object getHttpAttribute(HttpServletRequest request, String attributeName);
+    void removeAttribute(HttpServletRequest request, String attributeName);
 
     Map<String, String[]> getParameterMap(HttpServletRequest httpServletRequest);
 
diff --git a/apm-agent-plugins/apm-spring-webmvc/apm-spring-webmvc-spring5/src/test/java/co/elastic/apm/agent/springwebmvc/exception/AbstractExceptionHandlerInstrumentationWithExceptionResolverTest.java b/apm-agent-plugins/apm-spring-webmvc/apm-spring-webmvc-spring5/src/test/java/co/elastic/apm/agent/springwebmvc/exception/AbstractExceptionHandlerInstrumentationWithExceptionResolverTest.java
index 1f62d500d2..0800056204 100644
--- a/apm-agent-plugins/apm-spring-webmvc/apm-spring-webmvc-spring5/src/test/java/co/elastic/apm/agent/springwebmvc/exception/AbstractExceptionHandlerInstrumentationWithExceptionResolverTest.java
+++ b/apm-agent-plugins/apm-spring-webmvc/apm-spring-webmvc-spring5/src/test/java/co/elastic/apm/agent/springwebmvc/exception/AbstractExceptionHandlerInstrumentationWithExceptionResolverTest.java
@@ -26,6 +26,9 @@
 import org.springframework.test.web.servlet.MvcResult;
 import org.springframework.test.web.servlet.ResultActions;
 
+import java.util.Enumeration;
+
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
 
@@ -45,6 +48,16 @@ public void testCallApiWithExceptionThrown() throws Exception {
         MvcResult result = resultActions.andReturn();
         MockHttpServletResponse response = result.getResponse();
 
+
+        Enumeration<String> attributeNames = result.getRequest().getAttributeNames();
+        while (attributeNames.hasMoreElements()) {
+            String attributeName = attributeNames.nextElement();
+            assertThat(attributeName)
+                .describedAs("elastic attributes should be removed after usage")
+                .doesNotStartWith("co.elastic.");
+        }
+
+
         assertExceptionCapture(ExceptionResolverRuntimeException.class, response, 200, "", "runtime exception occurred", "View#render error-page");
         assertEquals("error-page", response.getForwardedUrl());
         assertEquals("runtime exception occurred", result.getModelAndView().getModel().get("message"));