From f14eea011cff60ae86e068da2e0eac958625332b Mon Sep 17 00:00:00 2001
From: Xavier Pinho <mail@xavierp.org>
Date: Sat, 13 Jan 2024 18:24:29 +0000
Subject: [PATCH] javasrc2cpg - workaround for failing to resolve a method's
 external return type when it contains type arguments (#4044)

---
 .../declarations/AstForMethodsCreator.scala   | 12 ++++
 .../javasrc2cpg/querying/MethodTests.scala    | 72 +++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/joern-cli/frontends/javasrc2cpg/src/main/scala/io/joern/javasrc2cpg/astcreation/declarations/AstForMethodsCreator.scala b/joern-cli/frontends/javasrc2cpg/src/main/scala/io/joern/javasrc2cpg/astcreation/declarations/AstForMethodsCreator.scala
index eadb1620b711..bcd85cd277b1 100644
--- a/joern-cli/frontends/javasrc2cpg/src/main/scala/io/joern/javasrc2cpg/astcreation/declarations/AstForMethodsCreator.scala
+++ b/joern-cli/frontends/javasrc2cpg/src/main/scala/io/joern/javasrc2cpg/astcreation/declarations/AstForMethodsCreator.scala
@@ -40,6 +40,7 @@ import io.shiftleft.codepropertygraph.generated.Operators
 import io.shiftleft.codepropertygraph.generated.DispatchTypes
 import io.shiftleft.codepropertygraph.generated.EdgeTypes
 import com.github.javaparser.ast.Node
+import com.github.javaparser.symbolsolver.javaparsermodel.declarations.JavaParserParameterDeclaration
 import io.joern.javasrc2cpg.astcreation.declarations.AstForMethodsCreator.PartialConstructorDeclaration
 
 private[declarations] trait AstForMethodsCreator { this: AstCreator =>
@@ -54,6 +55,9 @@ private[declarations] trait AstForMethodsCreator { this: AstCreator =>
     val returnTypeFullName = expectedReturnType
       .flatMap(typeInfoCalc.fullName)
       .orElse(scope.lookupType(simpleMethodReturnType))
+      .orElse(
+        Try(methodDeclaration.getType.asClassOrInterfaceType).toOption.flatMap(t => scope.lookupType(t.getNameAsString))
+      )
       .orElse(typeParameters.find(_.name == simpleMethodReturnType).map(_.typeFullName))
 
     scope.pushMethodScope(
@@ -260,6 +264,14 @@ private[declarations] trait AstForMethodsCreator { this: AstCreator =>
         .map { param =>
           Try(param.getType).toOption
             .flatMap(paramType => typeInfoCalc.fullName(paramType, typeParamValues))
+            // In a scenario where we have an import of an external type e.g. `import foo.bar.Baz` and
+            // this parameter's type is e.g. `Baz<String>`, the lookup will fail. However, if we lookup
+            // for `Baz` instead (i.e. without type arguments), then the lookup will succeed.
+            .orElse(
+              Try(
+                param.asInstanceOf[JavaParserParameterDeclaration].getWrappedNode.getType.asClassOrInterfaceType
+              ).toOption.flatMap(t => scope.lookupType(t.getNameAsString))
+            )
         }
 
     toOptionList(parameterTypes)
diff --git a/joern-cli/frontends/javasrc2cpg/src/test/scala/io/joern/javasrc2cpg/querying/MethodTests.scala b/joern-cli/frontends/javasrc2cpg/src/test/scala/io/joern/javasrc2cpg/querying/MethodTests.scala
index 3380d238b602..a93782a0b32a 100644
--- a/joern-cli/frontends/javasrc2cpg/src/test/scala/io/joern/javasrc2cpg/querying/MethodTests.scala
+++ b/joern-cli/frontends/javasrc2cpg/src/test/scala/io/joern/javasrc2cpg/querying/MethodTests.scala
@@ -103,3 +103,75 @@ class MethodTests3 extends JavaSrcCode2CpgFixture {
     cpg.method("virtualMethod").isVirtual.fullName.head shouldBe "Foo.virtualMethod:void(java.lang.Integer)"
   }
 }
+
+class MethodTests4 extends JavaSrcCode2CpgFixture {
+
+  "List<String> in the method return type" should {
+    val cpg = code("""
+        |import java.util.*;
+        |class Foo {
+        | List<String> run() {
+        |   return null;
+        | }
+        |}
+        |""".stripMargin)
+
+    "have correct signature and full name" in {
+      val List(method) = cpg.method("run").l
+      method.signature shouldBe "java.util.List()"
+      method.fullName shouldBe "Foo.run:java.util.List()"
+    }
+  }
+
+  "Baz<String> in the method return type" should {
+    val cpg = code("""
+        |import foo.bar.Baz;
+        |class Foo {
+        | Baz<String> run() {
+        |   return null;
+        | }
+        |}
+        |""".stripMargin)
+
+    "have correct signature and full name" in {
+      val List(method) = cpg.method("run").l
+      method.signature shouldBe "foo.bar.Baz()"
+      method.fullName shouldBe "Foo.run:foo.bar.Baz()"
+    }
+  }
+
+  "Identity method for Baz<String>" should {
+    val cpg = code("""
+        |import foo.bar.Baz;
+        |class Foo {
+        | Baz<String> run(Baz<String> x) {
+        |   return x;
+        | }
+        |}
+        |""".stripMargin)
+
+    "have correct signature and full name" in {
+      val List(method) = cpg.method("run").l
+      method.signature shouldBe "foo.bar.Baz(foo.bar.Baz)"
+      method.fullName shouldBe "Foo.run:foo.bar.Baz(foo.bar.Baz)"
+    }
+  }
+
+  "Generic identity method for Baz<T>" should {
+    val cpg = code("""
+        |import foo.bar.Baz;
+        |class Foo {
+        | <T> Baz<T> run(Baz<T> x) {
+        |   return x;
+        | }
+        |}
+        |""".stripMargin)
+
+    "have correct signature and full name" in {
+      val List(method) = cpg.method("run").l
+      method.signature shouldBe "foo.bar.Baz(foo.bar.Baz)"
+      method.fullName shouldBe "Foo.run:foo.bar.Baz(foo.bar.Baz)"
+    }
+  }
+
+}