runtime: Instrument Map#getOrDefault and Map#containsKey

Author: fmeum

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

@@ -756,18 +756,46 @@ public static void arraysCompareRange(
   // key closest to the current lookup key in the mapGet hook.
   private static final int MAX_NUM_KEYS_TO_ENUMERATE = 100;
 
-  @SuppressWarnings({"rawtypes", "unchecked"})
-  @MethodHook(type = HookType.AFTER, targetClassName = "java.util.Map", targetMethod = "get")
+  @MethodHook(
+      type = HookType.AFTER,
+      targetClassName = "java.util.Map",
+      targetMethod = "containsKey",
+      targetMethodDescriptor = "(Ljava/lang/Object;)Z")
+  public static void containsKey(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Boolean isContained) {
+    if (!isContained) {
+      mapHookInternal((Map) thisObject, arguments[0], hookId);
+    }
+  }
+
+  @MethodHook(
+      type = HookType.AFTER,
+      targetClassName = "java.util.Map",
+      targetMethod = "get",
+      targetMethodDescriptor = "(Ljava/lang/Object;)Ljava/lang/Object;")
   public static void mapGet(
-      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Object returnValue) {
-    if (returnValue != null) return;
-    if (arguments.length != 1) {
-      return;
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Object value) {
+    if (value == null) {
+      mapHookInternal((Map) thisObject, arguments[0], hookId);
     }
-    if (thisObject == null) return;
-    final Map map = (Map) thisObject;
-    if (map.size() == 0) return;
-    final Object currentKey = arguments[0];
+  }
+
+  @MethodHook(
+      type = HookType.AFTER,
+      targetClassName = "java.util.Map",
+      targetMethod = "getOrDefault",
+      targetMethodDescriptor = "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;")
+  public static void mapGetOrDefault(
+      MethodHandle method, Object thisObject, Object[] arguments, int hookId, Object value) {
+    Object defaultValue = arguments[1];
+    if (value == defaultValue) {
+      mapHookInternal((Map) thisObject, arguments[0], hookId);
+    }
+  }
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static <K, V> void mapHookInternal(Map<K, V> map, K currentKey, int hookId) {
+    if (map == null || map.isEmpty()) return;
     if (currentKey == null) return;
     // Find two valid map keys that bracket currentKey.
     // This is a generalization of libFuzzer's __sanitizer_cov_trace_switch:
@@ -776,7 +804,7 @@ public static void mapGet(
     Object upperBoundKey = null;
     try {
       if (map instanceof TreeMap) {
-        final TreeMap treeMap = (TreeMap) map;
+        final TreeMap<K, V> treeMap = (TreeMap<K, V>) map;
         try {
           lowerBoundKey = treeMap.floorKey(currentKey);
           upperBoundKey = treeMap.ceilingKey(currentKey);

tests/BUILD.bazel

@@ -576,6 +576,21 @@ java_fuzz_target_test(
     ],
 )
 
+java_fuzz_target_test(
+    name = "MapFuzzer",
+    srcs = ["src/test/java/com/example/MapFuzzer.java"],
+    allowed_findings = ["com.code_intelligence.jazzer.api.FuzzerSecurityIssueMedium"],
+    fuzzer_args = [
+        "--experimental_mutator",
+        "-use_value_profile=1",
+    ],
+    target_class = "com.example.MapFuzzer",
+    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/MapFuzzer.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2023 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.Map;
+
+public class MapFuzzer {
+  public static void fuzzerTestOneInput(@NotNull Map<@NotNull String, @NotNull String> map) {
+    if (map.getOrDefault("some_key", "").startsWith("prefix")) {
+      if (map.containsKey("other_key")) {
+        throw new FuzzerSecurityIssueMedium();
+      }
+    }
+  }
+}