Add support for var-args method calls with @PolymorphicSignature causing unmatched method signatures

sootup.callgraph/pom.xml

@@ -56,6 +56,10 @@
 			<artifactId>sootup.jimple.frontend</artifactId>
 			<scope>test</scope>
 		</dependency>
-	</dependencies>
+        <dependency>
+            <groupId>org.soot-oss</groupId>
+            <artifactId>sootup.java.core</artifactId>
+        </dependency>
+    </dependencies>
 
 </project>

sootup.callgraph/src/main/java/sootup/callgraph/AbstractCallGraphAlgorithm.java

@@ -49,6 +49,8 @@
 import sootup.core.types.ClassType;
 import sootup.core.types.VoidType;
 import sootup.core.views.View;
+import sootup.java.core.AnnotationUsage;
+import sootup.java.core.JavaSootClass;
 
 /**
  * The AbstractCallGraphAlgorithm class is the super class of all call graph algorithm. It provides
@@ -61,10 +63,42 @@ public abstract class AbstractCallGraphAlgorithm implements CallGraphAlgorithm {
 
   @NonNull protected final View view;
   @NonNull protected final TypeHierarchy typeHierarchy;
+  @NonNull protected final List<MethodSignature> preanalysis;
 
   protected AbstractCallGraphAlgorithm(@NonNull View view) {
     this.view = view;
     this.typeHierarchy = view.getTypeHierarchy();
+    this.preanalysis = getMethodsWithPolymorphicAnnotation();
+  }
+
+  /**
+   * Collects all method signatures from the current view that are annotated with {@code
+   * java.lang.invoke.MethodHandle$PolymorphicSignature}.
+   *
+   * @return a list of method signatures annotated with {@code @PolymorphicSignature}
+   */
+  protected List<MethodSignature> getMethodsWithPolymorphicAnnotation() {
+    List<MethodSignature> polymorphicMethodSigs = new ArrayList<>();
+    ClassType polymorphicAnnotationType =
+        view.getIdentifierFactory()
+            .getClassType("java.lang.invoke.MethodHandle$PolymorphicSignature");
+    view.getClasses()
+        .filter(sootClass -> sootClass instanceof JavaSootClass)
+        .map(sootClass -> (JavaSootClass) sootClass)
+        .flatMap(javaSootClass -> javaSootClass.getMethods().stream())
+        .forEach(
+            javaSootMethod -> {
+              Iterable<AnnotationUsage> annotationUsages = javaSootMethod.getAnnotations();
+              if (annotationUsages.equals(Collections.emptyList())) {
+                return;
+              }
+              for (AnnotationUsage annotationUsage : annotationUsages) {
+                if (annotationUsage.getAnnotation().equals(polymorphicAnnotationType)) {
+                  polymorphicMethodSigs.add(javaSootMethod.getSignature());
+                }
+              }
+            });
+    return polymorphicMethodSigs;
   }
 
   /**
@@ -592,18 +626,7 @@ protected static Optional<? extends SootMethod> findConcreteMethod(
     }
 
     // search method in interfaces
-    Optional<? extends SootMethod> defaultMethod = findDefaultMethod(view, startClass, methodSig);
-    if (defaultMethod.isPresent()) {
-      return defaultMethod;
-    }
-
-    logger.warn(
-        "Could not find \""
-            + sig.getSubSignature()
-            + "\" in "
-            + sig.getDeclClassType().getClassName()
-            + " and in its superclasses and interfaces");
-    return Optional.empty();
+    return findDefaultMethod(view, startClass, methodSig);
   }
 
   protected static Optional<SootMethod> findMethodInHierarchy(
@@ -667,4 +690,30 @@ protected static Optional<? extends SootMethod> findDefaultMethod(
   protected boolean isInterface(ClassType classType) {
     return typeHierarchy.isInterface(classType);
   }
+
+  /**
+   * Attempts to resolve a previously unresolved method signature by comparing it with all known
+   * method signatures annotated with {@code @PolymorphicSignature}. If a match is found (ignoring
+   * the parameters), the corresponding SootMethod is returned.
+   *
+   * @param targetMethodSignature the signature of the searched method
+   * @return the found method object, or null if the method was not found.
+   */
+  protected Optional<? extends SootMethod> findMatchingPolymorphicMethod(
+      @NonNull MethodSignature targetMethodSignature) {
+    for (MethodSignature polymorphicMethodSigs : preanalysis) {
+      if (targetMethodSignature.getDeclClassType().equals(polymorphicMethodSigs.getDeclClassType())
+          && targetMethodSignature.getName().equals(polymorphicMethodSigs.getName())) {
+        // return the actualTargetMethodOpt
+        return view.getMethod(polymorphicMethodSigs).map(sootMethod -> (SootMethod) sootMethod);
+      }
+    }
+    logger.warn(
+        "Could not find \""
+            + targetMethodSignature.getSubSignature()
+            + "\" in "
+            + targetMethodSignature.getDeclClassType().getClassName()
+            + " and in its superclasses and interfaces");
+    return Optional.empty();
+  }
 }

sootup.callgraph/src/main/java/sootup/callgraph/ClassHierarchyAnalysisAlgorithm.java

@@ -93,9 +93,14 @@ protected Stream<MethodSignature> resolveCall(SootMethod method, InvokableStmt i
     if (actualTargetMethod == null) {
       // method not implemented, search for implementation in super classes or ínterfaces
       actualTargetMethod = findConcreteMethod(view, targetMethodSignature).orElse(null);
-      // method implementation isn't contained in the view. return the called method as target
       if (actualTargetMethod == null) {
-        return Stream.of(targetMethodSignature);
+        // method implementation isn't contained in the view.
+        // check if method got the PolymorphicSignature annotation
+        actualTargetMethod = findMatchingPolymorphicMethod(targetMethodSignature).orElse(null);
+        if (actualTargetMethod == null) {
+          // method couldn't be resolved, return the called method as target
+          return Stream.of(targetMethodSignature);
+        }
       }
     }
 

sootup.callgraph/src/main/java/sootup/callgraph/RapidTypeAnalysisAlgorithm.java

@@ -34,6 +34,7 @@
 import sootup.core.jimple.common.expr.JSpecialInvokeExpr;
 import sootup.core.jimple.common.stmt.InvokableStmt;
 import sootup.core.jimple.common.stmt.JAssignStmt;
+import sootup.core.jimple.common.stmt.Stmt;
 import sootup.core.model.MethodModifier;
 import sootup.core.model.SootClass;
 import sootup.core.model.SootMethod;
@@ -96,13 +97,21 @@ protected List<ClassType> collectInstantiatedClassesInMethod(SootMethod method)
     if (method == null || method.isAbstract() || method.isNative()) {
       return Collections.emptyList();
     }
-
     Set<ClassType> instantiated =
         method.getBody().getStmts().stream()
-            .filter(stmt -> stmt instanceof JAssignStmt)
+            .filter(Stmt::isJAssignStmt)
             .map(stmt -> ((JAssignStmt) stmt).getRightOp())
-            .filter(value -> value instanceof JNewExpr)
-            .map(value -> ((JNewExpr) value).getType())
+            .flatMap(
+                value -> {
+                  // ClassType as returnType
+                  if (value instanceof AbstractInvokeExpr && value.getType() instanceof ClassType) {
+                    return Stream.of((ClassType) value.getType());
+                    // new-expression
+                  } else if (value instanceof JNewExpr) {
+                    return Stream.of(((JNewExpr) value).getType());
+                  }
+                  return Stream.empty();
+                })
             .collect(Collectors.toSet());
     List<ClassType> newInstantiatedClassTypes =
         instantiated.stream()
@@ -141,9 +150,14 @@ protected Stream<MethodSignature> resolveCall(
     if (actualTargetMethod == null) {
       // method not implemented, search for implementation in super classes or ínterfaces
       actualTargetMethod = findConcreteMethod(view, targetMethodSignature).orElse(null);
-      // method implementation isn't contained in the view. return the called method as target
       if (actualTargetMethod == null) {
-        return Stream.of(targetMethodSignature);
+        // method implementation isn't contained in the view.
+        // check if method got the PolymorphicSignature annotation
+        actualTargetMethod = findMatchingPolymorphicMethod(targetMethodSignature).orElse(null);
+        if (actualTargetMethod == null) {
+          // method couldn't be resolved, return the called method as target
+          return Stream.of(targetMethodSignature);
+        }
       }
     }
     // special invokes and static invokes can only have one specific target
@@ -159,7 +173,6 @@ protected Stream<MethodSignature> resolveCall(
             .subtypesOf(targetMethodSignature.getDeclClassType())
             .flatMap(classType -> view.getClass(classType).stream())
             .toList();
-
     // get all targets of these subtypes
     Stream<MethodSignature> targets =
         resolveAllCallTargets(
@@ -317,8 +330,9 @@ protected void preProcessingMethod(
     if (method == null) {
       return;
     }
-
+    System.out.println("PreProcessed Method: " + method);
     List<ClassType> newInstantiatedClasses = collectInstantiatedClassesInMethod(method);
+    System.out.println("New Instantiated Classes: " + newInstantiatedClasses);
     newInstantiatedClasses.forEach(
         classType -> includeIgnoredCallsToClass(classType, cg, workList));
   }

sootup.callgraph/src/test/java/sootup/callgraph/CallGraphTestBase.java

@@ -2,9 +2,10 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
-import java.util.*;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import org.junit.jupiter.api.Test;
 import sootup.core.inputlocation.AnalysisInputLocation;
 import sootup.core.jimple.common.Value;
@@ -41,7 +42,7 @@ public abstract class CallGraphTestBase<T extends AbstractCallGraphAlgorithm> {
 
   private JavaView createViewForClassPath(String classPath) {
     List<AnalysisInputLocation> inputLocations = new ArrayList<>();
-    inputLocations.add(new DefaultRuntimeAnalysisInputLocation());
+    inputLocations.add(new DefaultRuntimeAnalysisInputLocation()); // SourceType.Application
     inputLocations.add(new JavaClassPathAnalysisInputLocation(classPath));
 
     return new JavaView(inputLocations);
@@ -1265,4 +1266,43 @@ public void testImplicitExample6() {
             identifierFactory.getClassType("t6.NoThread"), "run", "void", Collections.emptyList());
     assertFalse(cg.containsMethod(runMethodSig));
   }
+
+  @Test
+  public void testMethodHandleInvokeExample1() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e1.MethodHandleInvokeExample1");
+    MethodSignature invokeMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.MethodHandle"),
+            "invoke",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(invokeMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
+
+  @Test
+  public void testMethodHandleInvokeExample2() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e2.MethodHandleInvokeExample2");
+    MethodSignature invokeMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.MethodHandle"),
+            "invoke",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(invokeMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
+
+  @Test
+  public void testVarHandleGetExample3() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e3.VarHandleGetExample3");
+    MethodSignature getMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.VarHandle"),
+            "get",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(getMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
 }

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/MethodHandleInvokeExample1.java

@@ -0,0 +1,14 @@
+package e1;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+
+public class MethodHandleInvokeExample1 {
+    public static void main(String[] args) throws Throwable {
+        MethodHandle handle = MethodHandles.lookup()
+                .findVirtual(String.class, "length", MethodType.methodType(int.class));
+
+        handle.invoke("Hello");
+    }
+}

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/MethodHandleInvokeExample2.java

@@ -0,0 +1,14 @@
+package e2;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+
+public class MethodHandleInvokeExample2 {
+    public static void main(String[] args) throws Throwable {
+        MethodHandle handle = MethodHandles.lookup()
+                .findVirtual(String.class, "substring", MethodType.methodType(String.class, int.class, int.class));
+
+        handle.invoke("HelloWorld", 0, 5);
+    }
+}

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/VarHandleGetExample3.java

@@ -0,0 +1,15 @@
+package e3;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+
+public class VarHandleGetExample3 {
+    static int value = 42;
+
+    public static void main(String[] args) throws Throwable {
+        VarHandle handle = MethodHandles.lookup()
+                .findStaticVarHandle(VarHandleGetExample3.class, "value", int.class);
+
+        handle.get();
+    }
+}