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,43 @@ public static void mapGetOrDefault(
     }
   }
 
+  static class LowerUpperBounds {
+    public Object lowerBound;
+    public Object upperBound;
+  }
+
+  private static <E> LowerUpperBounds getLowerUpperBounds(Set<E> elements, E currentElement) {
+    int enumeratedKeys = 0;
+    Comparable comparableElement = (Comparable) currentElement;
+
+    LowerUpperBounds bounds = new LowerUpperBounds();
+    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
+            && (bounds.lowerBound == null
+                || comparableValidElement.compareTo(bounds.lowerBound) > 0)) {
+          bounds.lowerBound = validElement;
+        }
+        if (comparableValidElement.compareTo(comparableElement) > 0
+            && (bounds.upperBound == null
+                || comparableValidElement.compareTo(bounds.upperBound) < 0)) {
+          bounds.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 (enumeratedKeys++ > MAX_NUM_ELEMENTS_TO_ENUMERATE) break;
+    }
+
+    return bounds;
+  }
+
   @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 +913,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 +923,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;
-        }
+        LowerUpperBounds bounds = getLowerUpperBounds(map.keySet(), currentKey);
+        lowerBoundKey = bounds.lowerBound;
+        upperBoundKey = bounds.upperBound;
       }
     } catch (ConcurrentModificationException ignored) {
       // map was modified by another thread, skip this invocation
@@ -901,6 +940,45 @@ 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 is of a type that can't be
+          // compared to the maps keys.
+        }
+
+      } else if (currentElement instanceof Comparable) {
+        LowerUpperBounds bounds = getLowerUpperBounds(set, currentElement);
+        lowerBoundElement = bounds.lowerBound;
+        upperBoundElement = bounds.upperBound;
+      }
+    } 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",