Report Map lookups as comparisons

If map.get(currentKey) returns null, the new hook finds a valid key
targetKey in map that is closest to currentKey and invokes a suitable
compare hook.

2 files changed, 88 insertions, 0 deletions

diff --git a/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceCmpHooks.java b/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceCmpHooks.java
index e8d5d89f..352da8ea 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceCmpHooks.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceCmpHooks.java
@@ -18,6 +18,8 @@ import com.code_intelligence.jazzer.api.HookType;
 import com.code_intelligence.jazzer.api.MethodHook;
 import java.lang.invoke.MethodHandle;
 import java.util.Arrays;
+import java.util.Map;
+import java.util.TreeMap;
 
 @SuppressWarnings("unused")
 final public class TraceCmpHooks {
@@ -253,4 +255,76 @@ final public class TraceCmpHooks {
       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;
+
+  @MethodHook(type = HookType.AFTER, targetClassName = "com.google.common.collect.ImmutableMap",
+      targetMethod = "get")
+  @MethodHook(
+      type = HookType.AFTER, targetClassName = "java.util.AbstractMap", targetMethod = "get")
+  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.EnumMap", targetMethod = "get")
+  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.HashMap", targetMethod = "get")
+  @MethodHook(
+      type = HookType.AFTER, targetClassName = "java.util.LinkedHashMap", targetMethod = "get")
+  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.Map", targetMethod = "get")
+  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.SortedMap", targetMethod = "get")
+  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.TreeMap", targetMethod = "get")
+  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.concurrent.ConcurrentMap",
+      targetMethod = "get")
+  public static void
+  mapGet(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Object returnValue) {
+    if (returnValue != null)
+      return;
+    if (thisObject == null)
+      return;
+    final Map map = (Map) thisObject;
+    if (map.size() == 0)
+      return;
+    final Object currentKey = arguments[0];
+    if (currentKey == null)
+      return;
+    // Find two valid map keys that bracket currentKey.
+    // This is a generalization of libFuzzer's __sanitizer_cov_trace_switch:
+    // https://github.com/llvm/llvm-project/blob/318942de229beb3b2587df09e776a50327b5cef0/compiler-rt/lib/fuzzer/FuzzerTracePC.cpp#L564
+    Object lowerBoundKey = null;
+    Object upperBoundKey = null;
+    if (map instanceof TreeMap) {
+      final TreeMap treeMap = (TreeMap) map;
+      lowerBoundKey = treeMap.floorKey(currentKey);
+      upperBoundKey = treeMap.ceilingKey(currentKey);
+    } else if (currentKey instanceof Comparable) {
+      final Comparable comparableKey = (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 == null)
+          continue;
+        // 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.
+        if (comparableKey.compareTo(validKey) > 0
+            && (lowerBoundKey == null || ((Comparable) validKey).compareTo(lowerBoundKey) > 0)) {
+          lowerBoundKey = validKey;
+        }
+        if (comparableKey.compareTo(validKey) < 0
+            && (upperBoundKey == null || ((Comparable) validKey).compareTo(upperBoundKey) < 0)) {
+          upperBoundKey = validKey;
+        }
+        if (enumeratedKeys++ > MAX_NUM_KEYS_TO_ENUMERATE)
+          break;
+      }
+    }
+    // Modify the hook ID so that compares against distinct valid keys are traced separately.
+    if (lowerBoundKey != null) {
+      TraceDataFlowNativeCallbacks.traceGenericCmp(
+          currentKey, lowerBoundKey, hookId + lowerBoundKey.hashCode());
+    }
+    if (upperBoundKey != null) {
+      TraceDataFlowNativeCallbacks.traceGenericCmp(
+          currentKey, upperBoundKey, hookId + upperBoundKey.hashCode());
+    }
+  }
 }
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceDataFlowNativeCallbacks.java b/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceDataFlowNativeCallbacks.java
index 92b21478..456d0cb9 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceDataFlowNativeCallbacks.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/runtime/TraceDataFlowNativeCallbacks.java
@@ -73,5 +73,19 @@ final public class TraceDataFlowNativeCallbacks {
     return Long.compare(arg1, arg2);
   }
 
+  // The caller has to ensure that arg1 and arg2 have the same class.
+  public static void traceGenericCmp(Object arg1, Object arg2, int pc) {
+    if (arg1 instanceof String) {
+      traceStrcmp((String) arg1, (String) arg2, 1, pc);
+    } else if (arg1 instanceof Integer || arg1 instanceof Short || arg1 instanceof Byte
+        || arg1 instanceof Character) {
+      traceCmpInt((int) arg1, (int) arg2, pc);
+    } else if (arg1 instanceof Long) {
+      traceCmpLong((long) arg1, (long) arg2, pc);
+    } else if (arg1 instanceof byte[]) {
+      traceMemcmp((byte[]) arg1, (byte[]) arg2, 1, pc);
+    }
+  }
+
   public static native void handleLibraryLoad();
 }