From 7c959129b28695806e16a0cea223047766e6f7b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alejandro=20Gonz=C3=A1lez=20Garc=C3=ADa?=
 <alejandro.gonzalez@datadoghq.com>
Date: Thu, 22 Jun 2023 10:55:54 +0200
Subject: [PATCH] Fix unvalidated redirect detection in Jetty

What Does This Do
Make method matchers in IAST instrumenter for servlet more precise. Avoids matching the wrong methods in Jetty.

Motivation
Some unvalidated redirects were not detected in Jetty.
---
 .../servlet/HttpServletResponseInstrumentation.java   | 10 ++++++++--
 .../JakartaHttpServletResponseInstrumentation.java    | 11 +++++++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java
index 7b7c1128e36..0336d67ed73 100644
--- a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java
+++ b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java
@@ -5,6 +5,7 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
@@ -49,8 +50,13 @@ public void adviceTransformations(AdviceTransformation transformation) {
         namedOneOf("setHeader", "addHeader").and(takesArguments(String.class, String.class)),
         getClass().getName() + "$AddHeaderAdvice");
     transformation.applyAdvice(
-        namedOneOf("encodeRedirectURL", "encodeURL"), getClass().getName() + "$EncodeURLAdvice");
-    transformation.applyAdvice(named("sendRedirect"), getClass().getName() + "$SendRedirectAdvice");
+        namedOneOf("encodeRedirectURL", "encodeURL")
+            .and(takesArgument(0, String.class))
+            .and(returns(String.class)),
+        getClass().getName() + "$EncodeURLAdvice");
+    transformation.applyAdvice(
+        named("sendRedirect").and(takesArgument(0, String.class)),
+        getClass().getName() + "$SendRedirectAdvice");
   }
 
   public static class AddCookieAdvice {
diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java
index 0046ff7070a..06878cb2b2a 100644
--- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java
+++ b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java
@@ -5,6 +5,8 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -39,8 +41,13 @@ public void adviceTransformations(AdviceTransformation transformation) {
     transformation.applyAdvice(
         namedOneOf("setHeader", "addHeader"), getClass().getName() + "$AddHeaderAdvice");
     transformation.applyAdvice(
-        namedOneOf("encodeRedirectURL", "encodeURL"), getClass().getName() + "$EncodeURLAdvice");
-    transformation.applyAdvice(named("sendRedirect"), getClass().getName() + "$SendRedirectAdvice");
+        namedOneOf("encodeRedirectURL", "encodeURL")
+            .and(takesArgument(0, String.class))
+            .and(returns(String.class)),
+        getClass().getName() + "$EncodeURLAdvice");
+    transformation.applyAdvice(
+        named("sendRedirect").and(takesArgument(0, String.class)),
+        getClass().getName() + "$SendRedirectAdvice");
   }
 
   public static class AddCookieAdvice {