feat: add hooks for Set.contains & Set.remove

Author: onionpsy

src/main/java/com/code_intelligence/jazzer/runtime/TraceCmpHooks.java

@@ -802,9 +802,32 @@ public static void arraysCompareRange(
     TraceDataFlowNativeCallbacks.traceMemcmp(first, second, returnValue, hookId);
   }
 
-  // The maximal number of elements of a non-TreeMap Map that will be sorted and searched for the
-  // key closest to the current lookup key in the mapGet hook.
-  private static final int MAX_NUM_KEYS_TO_ENUMERATE = 100;
+  // The maximal number of elements of collections we enumerate to find values close to a lookup.
+  private static final int MAX_NUM_ELEMENTS_TO_ENUMERATE = 100;
+
+  @MethodHook(
+      type = HookType.AFTER,
+      targetClassName = "java.util.Set",
+      targetMethod = "contains",
+      targetMethodDescriptor = "(Ljava/lang/Object;)Z")
+  public static void setContains(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Boolean isContained) {
+    if (!isContained) {
+      setHookInternal((Set) thisObject, arguments[0], hookId);
+    }
+  }
+
+  @MethodHook(
+      type = HookType.AFTER,
+      targetClassName = "java.util.Set",
+      targetMethod = "remove",
+      targetMethodDescriptor = "(Ljava/lang/Object;)Z")
+  public static void setRemove(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Boolean wasRemoved) {
+    if (!wasRemoved) {
+      setHookInternal((Set) thisObject, arguments[0], hookId);
+    }
+  }
 
   @MethodHook(
       type = HookType.AFTER,
@@ -843,6 +866,54 @@ public static void mapGetOrDefault(
     }
   }
 
+  private static final class Bounds {
+    private final Object lower;
+    private final Object upper;
+
+    private Bounds(Object lower, Object upper) {
+      this.lower = lower;
+      this.upper = upper;
+    }
+
+    private Object getLower() {
+      return lower;
+    }
+
+    private Object getUpper() {
+      return upper;
+    }
+  }
+
+  private static <E> Bounds getLowerUpperBounds(Set<E> elements, E currentElement) {
+    int enumeratedElements = 0;
+    Comparable comparableElement = (Comparable) currentElement;
+
+    Object lowerBound = null;
+    Object upperBound = null;
+    for (Object validElement : elements) {
+      if (!(validElement instanceof Comparable)) continue;
+      final Comparable comparableValidElement = (Comparable) validElement;
+      // If the element sorts lower than the non-existing elements, but higher than the current
+      // lower bound, update the lower bound and vice versa for the upper bound.
+      try {
+        if (comparableValidElement.compareTo(comparableElement) < 0
+            && (lowerBound == null || comparableValidElement.compareTo(lowerBound) > 0)) {
+          lowerBound = validElement;
+        }
+        if (comparableValidElement.compareTo(comparableElement) > 0
+            && (upperBound == null || comparableValidElement.compareTo(upperBound) < 0)) {
+          upperBound = validElement;
+        }
+      } catch (ClassCastException ignored) {
+        // Can be thrown by Comparable.compareTo if comparableElement is of a type that can't be
+        // compared to the elements set.
+      }
+      if (enumeratedElements++ > MAX_NUM_ELEMENTS_TO_ENUMERATE) break;
+    }
+
+    return new Bounds(lowerBound, upperBound);
+  }
+
   @SuppressWarnings({"rawtypes", "unchecked"})
   private static <K, V> void mapHookInternal(Map<K, V> map, K currentKey, int hookId) {
     if (map == null || map.isEmpty()) return;
@@ -853,8 +924,8 @@ private static <K, V> void mapHookInternal(Map<K, V> map, K currentKey, int hook
     Object lowerBoundKey = null;
     Object upperBoundKey = null;
     try {
-      if (map instanceof TreeMap) {
-        final TreeMap<K, V> treeMap = (TreeMap<K, V>) map;
+      if (map instanceof NavigableMap) {
+        final NavigableMap<K, V> treeMap = (NavigableMap<K, V>) map;
         try {
           lowerBoundKey = treeMap.floorKey(currentKey);
           upperBoundKey = treeMap.ceilingKey(currentKey);
@@ -863,30 +934,9 @@ private static <K, V> void mapHookInternal(Map<K, V> map, K currentKey, int hook
           // compared to the maps keys.
         }
       } else if (currentKey instanceof Comparable) {
-        final Comparable comparableCurrentKey = (Comparable) currentKey;
-        // Find two keys that bracket currentKey.
-        // Note: This is not deterministic if map.size() > MAX_NUM_KEYS_TO_ENUMERATE.
-        int enumeratedKeys = 0;
-        for (Object validKey : map.keySet()) {
-          if (!(validKey instanceof Comparable)) continue;
-          final Comparable comparableValidKey = (Comparable) validKey;
-          // If the key sorts lower than the non-existing key, but higher than the current lower
-          // bound, update the lower bound and vice versa for the upper bound.
-          try {
-            if (comparableValidKey.compareTo(comparableCurrentKey) < 0
-                && (lowerBoundKey == null || comparableValidKey.compareTo(lowerBoundKey) > 0)) {
-              lowerBoundKey = validKey;
-            }
-            if (comparableValidKey.compareTo(comparableCurrentKey) > 0
-                && (upperBoundKey == null || comparableValidKey.compareTo(upperBoundKey) < 0)) {
-              upperBoundKey = validKey;
-            }
-          } catch (ClassCastException ignored) {
-            // Can be thrown by floorKey and ceilingKey if currentKey is of a type that can't be
-            // compared to the maps keys.
-          }
-          if (enumeratedKeys++ > MAX_NUM_KEYS_TO_ENUMERATE) break;
-        }
+        Bounds bounds = getLowerUpperBounds(map.keySet(), currentKey);
+        lowerBoundKey = bounds.getLower();
+        upperBoundKey = bounds.getUpper();
       }
     } catch (ConcurrentModificationException ignored) {
       // map was modified by another thread, skip this invocation
@@ -901,6 +951,44 @@ private static <K, V> void mapHookInternal(Map<K, V> map, K currentKey, int hook
     }
   }
 
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static <E> void setHookInternal(Set<E> set, E currentElement, int hookId) {
+    if (set == null || set.isEmpty()) return;
+    if (currentElement == null) return;
+
+    Object lowerBoundElement = null;
+    Object upperBoundElement = null;
+
+    try {
+      if (set instanceof NavigableSet) {
+        final NavigableSet<E> navigableSet = (NavigableSet<E>) set;
+        try {
+          lowerBoundElement = navigableSet.floor(currentElement);
+          upperBoundElement = navigableSet.ceiling(currentElement);
+        } catch (ClassCastException ignored) {
+          // Can be thrown by NavigableSet.floor and NavigableSet.ceiling if the element cannot be
+          // compared to elements in the set.
+        }
+
+      } else if (currentElement instanceof Comparable) {
+        Bounds bounds = getLowerUpperBounds(set, currentElement);
+        lowerBoundElement = bounds.getLower();
+        upperBoundElement = bounds.getUpper();
+      }
+    } catch (ConcurrentModificationException ignored) {
+      // set was modified by another thread, skip this invocation
+      return;
+    }
+
+    if (lowerBoundElement != null) {
+      TraceDataFlowNativeCallbacks.traceGenericCmp(currentElement, lowerBoundElement, hookId);
+    }
+    if (upperBoundElement != null) {
+      TraceDataFlowNativeCallbacks.traceGenericCmp(
+          currentElement, upperBoundElement, 31 * hookId + 11);
+    }
+  }
+
   @MethodHook(
       type = HookType.AFTER,
       targetClassName = "org.junit.jupiter.api.Assertions",

tests/BUILD.bazel

@@ -696,6 +696,21 @@ java_fuzz_target_test(
     ],
 )
 
+java_fuzz_target_test(
+    name = "SetFuzzer",
+    srcs = ["src/test/java/com/example/SetFuzzer.java"],
+    allowed_findings = ["com.code_intelligence.jazzer.api.FuzzerSecurityIssueMedium"],
+    fuzzer_args = [
+        "-use_value_profile=1",
+        "-runs=1000000",
+    ],
+    target_class = "com.example.SetFuzzer",
+    verify_crash_reproducer = False,
+    deps = [
+        "//src/main/java/com/code_intelligence/jazzer/mutation/annotation",
+    ],
+)
+
 sh_test(
     name = "jazzer_from_path_test",
     srcs = ["src/test/shell/jazzer_from_path_test.sh"],

tests/src/test/java/com/example/SetFuzzer.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2024 Code Intelligence GmbH
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example;
+
+import com.code_intelligence.jazzer.api.FuzzerSecurityIssueMedium;
+import com.code_intelligence.jazzer.mutation.annotation.NotNull;
+import java.util.Set;
+
+public class SetFuzzer {
+  public static void fuzzerTestOneInput(@NotNull Set<@NotNull String> set) {
+    if (set.contains("9ab8439jf983") && set.remove("87sdfj89d")) {
+      throw new FuzzerSecurityIssueMedium();
+    }
+  }
+}