From 6c313fa9c4e0c1d5678f1cd8cc1725e15e11a42d Mon Sep 17 00:00:00 2001
From: Jonas Kunz <jonas.kunz@elastic.co>
Date: Tue, 10 Dec 2024 15:08:08 +0100
Subject: [PATCH] Support unwrapping webflux requests

---
 .../springwebflux/WebfluxServletHelper.java   | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/apm-agent-plugins/apm-spring-webflux/apm-spring-webflux-spring5/src/main/java/co/elastic/apm/agent/springwebflux/WebfluxServletHelper.java b/apm-agent-plugins/apm-spring-webflux/apm-spring-webflux-spring5/src/main/java/co/elastic/apm/agent/springwebflux/WebfluxServletHelper.java
index cc525a7f3a..3ee10c86d7 100644
--- a/apm-agent-plugins/apm-spring-webflux/apm-spring-webflux-spring5/src/main/java/co/elastic/apm/agent/springwebflux/WebfluxServletHelper.java
+++ b/apm-agent-plugins/apm-spring-webflux/apm-spring-webflux-spring5/src/main/java/co/elastic/apm/agent/springwebflux/WebfluxServletHelper.java
@@ -18,12 +18,13 @@
  */
 package co.elastic.apm.agent.springwebflux;
 
-import co.elastic.apm.agent.sdk.weakconcurrent.WeakKeySoftValueLoadingCache;
-import co.elastic.apm.agent.tracer.Transaction;
 import co.elastic.apm.agent.sdk.logging.Logger;
 import co.elastic.apm.agent.sdk.logging.LoggerFactory;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakKeySoftValueLoadingCache;
+import co.elastic.apm.agent.tracer.Transaction;
 import org.springframework.http.server.reactive.AbstractServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
 import org.springframework.web.server.ServerWebExchange;
 
 import javax.annotation.Nullable;
@@ -60,19 +61,23 @@ public static Transaction<?> getServletTransaction(ServerWebExchange exchange) {
 
         // While the active transaction is the one created by Servlet, it would rely on the fact that we are on the
         // same thread as the one that created the transaction, which is an implementation detail.
-
-        Transaction<?> transaction = null;
-
         ServerHttpRequest exchangeRequest = exchange.getRequest();
+        return getTransactionFromRequest(exchangeRequest);
+    }
+
+    @Nullable
+    private static Transaction<?> getTransactionFromRequest(ServerHttpRequest exchangeRequest) {
         if (exchangeRequest instanceof AbstractServerHttpRequest) {
             Object nativeRequest = ((AbstractServerHttpRequest) exchangeRequest).getNativeRequest();
 
             // note: attribute name is defined in Servlet plugin and should be kept in sync
-            transaction = (Transaction<?>) getServletAttribute(nativeRequest, "co.elastic.apm.agent.servlet.ServletApiAdvice.transaction");
+            return (Transaction<?>) getServletAttribute(nativeRequest, "co.elastic.apm.agent.servlet.ServletApiAdvice.transaction");
 
+        } else if (exchangeRequest instanceof ServerHttpRequestDecorator) {
+            ServerHttpRequestDecorator decorator = (ServerHttpRequestDecorator) exchangeRequest;
+            return getTransactionFromRequest(decorator.getDelegate());
         }
-
-        return transaction;
+        return null;
     }
 
     /**