Add support for stringbuffer in setlength and add remove method to ta…

…intedobjects
DataDog · Dec 23, 2024 · e2345c2 · e2345c2
1 parent 46b5986
commit e2345c2
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 7 deletions.
diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/StringModuleImpl.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/StringModuleImpl.java
@@ -852,6 +852,8 @@ public void onStringBuilderSetLength(@Nonnull CharSequence self, int length) {
     Range[] newRanges = Ranges.forSubstring(0, length, rangesSelf);
     if (newRanges != null && newRanges.length > 0) {
       selfTainted.setRanges(newRanges);
+    } else {
+      taintedObjects.untaint(self);
     }
   }
 

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedMap.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedMap.java
@@ -85,6 +85,8 @@ static TaintedMap buildWithPurge(final int capacity, int maxAge, TimeUnit maxAge
 
   void put(final @Nonnull TaintedObject entry);
 
+  boolean remove(final @Nonnull Object obj);
+
   int count();
 
   void clear();
@@ -210,6 +212,18 @@ public void put(final @Nonnull TaintedObject entry) {
       }
     }
 
+    @Override
+    public boolean remove(final @Nonnull Object obj) {
+      TaintedObject entry = get(obj);
+      if (entry == null) {
+        return false;
+      }
+
+      final int index = index(entry.positiveHashCode);
+      table[index] = null;
+      return true;
+    }
+
     @Override
     public void clear() {
       Arrays.fill(table, null);
@@ -352,6 +366,11 @@ public void put(@Nonnull final TaintedObject entry) {
       }
     }
 
+    @Override
+    public boolean remove(final @Nonnull Object obj) {
+      return delegate.remove(obj);
+    }
+
     @Nullable
     @Override
     public TaintedObject get(@Nonnull final Object key) {
@@ -455,6 +474,11 @@ public TaintedObject get(@Nonnull Object key) {
     @Override
     public void put(@Nonnull TaintedObject entry) {}
 
+    @Override
+    public boolean remove(@Nonnull Object obj) {
+      return false;
+    }
+
     @Override
     public int count() {
       return 0;

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedObjects.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedObjects.java
@@ -28,6 +28,8 @@ static TaintedObjects build(@Nonnull final TaintedMap map) {
   @Nullable
   TaintedObject taint(@Nonnull Object obj, @Nonnull Range[] ranges);
 
+  void untaint(@Nonnull Object obj);
+
   @Nullable
   TaintedObject get(@Nonnull Object obj);
 
@@ -56,6 +58,11 @@ public TaintedObject taint(final @Nonnull Object obj, final @Nonnull Range[] ran
       }
     }
 
+    @Override
+    public void untaint(final @Nonnull Object obj) {
+      map.remove(obj);
+    }
+
     @Nullable
     @Override
     public TaintedObject get(final @Nonnull Object obj) {
@@ -98,6 +105,12 @@ public TaintedObject taint(final @Nonnull Object obj, final @Nonnull Range[] ran
       return tainted;
     }
 
+    @Override
+    public void untaint(final @Nonnull Object obj) {
+      delegated.untaint(obj);
+      logUntainted(obj);
+    }
+
     @Nullable
     @Override
     public TaintedObject get(final @Nonnull Object obj) {
@@ -145,6 +158,20 @@ private void logTainted(@Nullable final TaintedObject tainted) {
       }
     }
 
+    private void logUntainted(@Nullable final Object obj) {
+      if (LOGGER.isDebugEnabled()) {
+        try {
+          if (obj == null) {
+            LOGGER.debug("untaint {}: not removed", id);
+          } else {
+            LOGGER.debug("untaint {}: untainted={}", id, obj.hashCode());
+          }
+        } catch (final Throwable e) {
+          LOGGER.error("Failed to debug untainted object", e);
+        }
+      }
+    }
+
     @Override
     public TaintedObjectsImpl unwrap() {
       return delegated;
@@ -161,6 +188,9 @@ public TaintedObject taint(@Nonnull final Object obj, @Nonnull final Range[] ran
       return null;
     }
 
+    @Override
+    public void untaint(final @Nonnull Object obj) {}
+
     @Nullable
     @Override
     public TaintedObject get(@Nonnull final Object obj) {

diff --git a/...gent-iast/src/main/java/com/datadog/iast/telemetry/taint/TaintedObjectsWithTelemetry.java b/...gent-iast/src/main/java/com/datadog/iast/telemetry/taint/TaintedObjectsWithTelemetry.java
@@ -45,6 +45,11 @@ public TaintedObject taint(@Nonnull Object obj, @Nonnull Range[] ranges) {
     return result;
   }
 
+  @Override
+  public void untaint(final @Nonnull Object obj) {
+    delegate.untaint(obj);
+  }
+
   @Nullable
   @Override
   public TaintedObject get(@Nonnull Object obj) {

diff --git a/...ava-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy b/...ava-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/StringModuleTest.groovy
@@ -1448,9 +1448,11 @@ class StringModuleTest extends IastModuleImplTestBase {
     0 * _
 
     where:
-    self      | length | mockCalls
-    sb("123") | 2      | 0
-    sb()      | 0      | 1
+    self       | length | mockCalls
+    sb("123")  | 2      | 0
+    sb()       | 0      | 1
+    sbf("123") | 2      | 0
+    sbf()      | 0      | 1
   }
 
   void 'onStringBuilderSetLength (#input, #length)'() {
@@ -1472,10 +1474,38 @@ class StringModuleTest extends IastModuleImplTestBase {
     taintFormat(result, taintedObject.getRanges()) == expected
 
     where:
-    input                         | length | expected
-    sb("==>0123<==")              | 3      | "==>012<=="
-    sb("0123==>456<==78")         | 5      | "0123==>4<=="
-    sb("01==>234<==5==>678<==90") | 8      | "01==>234<==5==>67<=="
+    input                          | length | expected
+    sb("==>0123<==")               | 3      | "==>012<=="
+    sb("0123==>456<==78")          | 5      | "0123==>4<=="
+    sb("01==>234<==5==>678<==90")  | 8      | "01==>234<==5==>67<=="
+    sbf("==>0123<==")              | 3      | "==>012<=="
+    sbf("0123==>456<==78")         | 5      | "0123==>4<=="
+    sbf("01==>234<==5==>678<==90") | 8      | "01==>234<==5==>67<=="
+  }
+
+  void 'onStringBuilderSetLength untainting after setLength (#input, #length)'() {
+    final taintedObjects = ctx.getTaintedObjects()
+    def self = addFromTaintFormat(taintedObjects, input)
+    if (self instanceof StringBuilder) {
+      ((StringBuilder) self).setLength(length)
+    } else if (self instanceof StringBuffer) {
+      ((StringBuffer) self).setLength(length)
+    }
+
+    when:
+    module.onStringBuilderSetLength(self, length)
+    def taintedObject = taintedObjects.get(self)
+
+    then:
+    1 * tracer.activeSpan() >> span
+    taintedObject == null
+
+    where:
+    input                  | length
+    sb("==>0123<==")       | 0
+    sb("0123==>456<==78")  | 3
+    sbf("==>0123<==")      | 0
+    sbf("0123==>456<==78") | 3
   }
 
   private static Date date(final String pattern, final String value) {

diff --git a/...ava-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java b/...ava-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringBuilderCallSite.java
@@ -183,6 +183,7 @@ public static CharSequence afterSubSequence(
   }
 
   @CallSite.After("void java.lang.StringBuilder.setLength(int)")
+  @CallSite.After("void java.lang.StringBuffer.setLength(int)")
   public static void afterSetLength(
       @CallSite.This final CharSequence self, @CallSite.Argument final int length) {
     final StringModule module = InstrumentationBridge.STRING;

diff --git a/.../src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy b/.../src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy
@@ -266,6 +266,7 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
     where:
     type      | suite                        | param         | length | expected
     "builder" | new TestStringBuilderSuite() | sb('012345')  | 5      | '01234'
+    "buffer"  | new TestStringBufferSuite()  | sbf('012345') | 5      | '01234'
   }
 
   private static class BrokenToString {