Refactor advice and helper in the freemarker instrumentation

DataDog · Sep 2, 2024 · 2c48ca0 · 2c48ca0
1 parent 0b09105
commit 2c48ca0
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 18 deletions.
diff --git a/...r/core/DollarVariable24DatadogAdvice.java → ...marker24/DollarVariableDatadogAdvice.java b/...r/core/DollarVariable24DatadogAdvice.java → ...marker24/DollarVariableDatadogAdvice.java
@@ -1,24 +1,25 @@
-package freemarker.core;
+package datadog.trace.instrumentation.freemarker24;
 
 import datadog.trace.api.iast.InstrumentationBridge;
 import datadog.trace.api.iast.Sink;
 import datadog.trace.api.iast.VulnerabilityTypes;
 import datadog.trace.api.iast.sink.XssModule;
-import freemarker.template.TemplateException;
+import freemarker.core.DollarVariable24Helper;
+import freemarker.core.Environment;
 import net.bytebuddy.asm.Advice;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public final class DollarVariable24DatadogAdvice {
+public final class DollarVariableDatadogAdvice {
 
-  protected static final Logger log = LoggerFactory.getLogger(DollarVariable24DatadogAdvice.class);
+  protected static final Logger log = LoggerFactory.getLogger(DollarVariableDatadogAdvice.class);
 
   public static class DollarVariableAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     @Sink(VulnerabilityTypes.XSS)
     public static void onEnter(
-        @Advice.Argument(0) final Environment environment, @Advice.This final DollarVariable self) {
+        @Advice.Argument(0) final Environment environment, @Advice.This final Object self) {
       if (environment == null || self == null) {
         return;
       }
@@ -29,15 +30,9 @@ public static void onEnter(
       if (DollarVariable24Helper.fetchAutoEscape(self)) {
         return;
       }
-      String charSec = null;
-      try {
-        charSec = (String) self.calculateInterpolatedStringOrMarkup(environment);
-      } catch (TemplateException e) {
-        log.debug("Failed to get DollarVariable templateModel", e);
-        return;
-      }
+      String charSec = DollarVariable24Helper.fetchCharSec(self, environment);
       final String templateName = environment.getMainTemplate().getName();
-      final int line = self.beginLine;
+      final int line = DollarVariable24Helper.fetchBeginLine(self);
       xssModule.onXss(charSec, templateName, line);
     }
   }

diff --git a/...c/main/java/datadog/trace/instrumentation/freemarker24/DollarVariableInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/freemarker24/DollarVariableInstrumentation.java
@@ -14,7 +14,6 @@
 public class DollarVariableInstrumentation extends InstrumenterModule.Iast
     implements Instrumenter.ForSingleType {
   static final String FREEMARKER_CORE = "freemarker.core";
-  static final String ADVICE_BASE = FREEMARKER_CORE + ".DollarVariable24DatadogAdvice$";
 
   public DollarVariableInstrumentation() {
     super("freemarker");
@@ -42,7 +41,6 @@ public String instrumentedType() {
   public String[] helperClassNames() {
     return new String[] {
       FREEMARKER_CORE + ".DollarVariable24Helper",
-      FREEMARKER_CORE + ".DollarVariable24DatadogAdvice"
     };
   }
 
@@ -52,6 +50,6 @@ public void methodAdvice(MethodTransformer transformer) {
         named("accept")
             .and(isMethod())
             .and(takesArgument(0, named(FREEMARKER_CORE + ".Environment"))),
-        ADVICE_BASE + "DollarVariableAdvice");
+        packageName + ".DollarVariableDatadogAdvice$DollarVariableAdvice");
   }
 }
diff --git a/...strumentation/freemarker-2.3.24/src/main/java/freemarker/core/DollarVariable24Helper.java b/...strumentation/freemarker-2.3.24/src/main/java/freemarker/core/DollarVariable24Helper.java
@@ -1,5 +1,6 @@
 package freemarker.core;
 
+import freemarker.template.TemplateException;
 import java.lang.reflect.Field;
 import java.lang.reflect.UndeclaredThrowableException;
 import org.slf4j.Logger;
@@ -24,8 +25,8 @@ private static Field prepareAutoEscape() {
     return autoEscape;
   }
 
-  public static boolean fetchAutoEscape(DollarVariable dollarVariable) {
-    if (AUTO_ESCAPE == null) {
+  public static boolean fetchAutoEscape(Object dollarVariable) {
+    if (AUTO_ESCAPE == null || !(dollarVariable instanceof DollarVariable)) {
       return true;
     }
     try {
@@ -34,4 +35,22 @@ public static boolean fetchAutoEscape(DollarVariable dollarVariable) {
       throw new UndeclaredThrowableException(e);
     }
   }
+
+  public static String fetchCharSec(Object object, Environment environment) {
+    if (!(object instanceof DollarVariable)) {
+      return null;
+    }
+    try {
+      return (String) ((DollarVariable) object).calculateInterpolatedStringOrMarkup(environment);
+    } catch (TemplateException e) {
+      throw new UndeclaredThrowableException(e);
+    }
+  }
+
+  public static Integer fetchBeginLine(Object object) {
+    if (!(object instanceof DollarVariable)) {
+      return null;
+    }
+    return ((DollarVariable) object).beginLine;
+  }
 }