[TypeRecovery] include <returnValue> dummy for index accesses of calls (

#4657)
joernio · Jun 12, 2024 · 706a17b · 706a17b
1 parent e7940fe
commit 706a17b
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 0 deletions.
diff --git a/.../frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/passes/TypeRecoveryPassTests.scala b/.../frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/passes/TypeRecoveryPassTests.scala
@@ -2,6 +2,7 @@ package io.joern.pysrc2cpg.passes
 
 import io.joern.pysrc2cpg.PySrc2CpgFixture
 import io.joern.x2cpg.passes.frontend.XTypeHintCallLinker
+import io.shiftleft.codepropertygraph.generated.Operators
 import io.shiftleft.codepropertygraph.generated.nodes.{Call, Identifier, Member}
 import io.shiftleft.semanticcpg.language.*
 import io.shiftleft.semanticcpg.language.importresolver.*
@@ -1457,4 +1458,47 @@ class TypeRecoveryPassTests extends PySrc2CpgFixture(withOssDataflow = false) {
       }
     }
   }
+
+  "assignment to imported method call" should {
+    val cpg = code("""
+        |from helpers import foo
+        |x = foo()
+        |""".stripMargin)
+
+    "provide meaningful typeFullName for the target of the assignment" in {
+      cpg.assignment.target.isIdentifier.name("x").l match {
+        case List(x) => x.typeFullName shouldBe "helpers.py:<module>.foo.<returnValue>"
+        case result  => fail(s"Expected single assignment to x, but got $result")
+      }
+    }
+  }
+
+  "assignment to non-chained index access of an imported method call" should {
+    val cpg = code("""
+        |from helpers import foo
+        |x = foo()
+        |y = x[0]
+        |""".stripMargin)
+
+    "provide meaningful typeFullName for the target of the assignment" in {
+      cpg.assignment.target.isIdentifier.name("y").l match {
+        case List(y) => y.typeFullName shouldBe "helpers.py:<module>.foo.<returnValue>.<indexAccess>"
+        case result  => fail(s"Expected single assignment to y, but got $result")
+      }
+    }
+  }
+
+  "assignment to chained index access of an imported method call" should {
+    val cpg = code("""
+        |from helpers import foo
+        |x = foo()[0]
+        |""".stripMargin)
+
+    "provide meaningful typeFullName for the target of the assignment" in {
+      cpg.assignment.target.isIdentifier.name("x").l match {
+        case List(x) => x.typeFullName shouldBe "helpers.py:<module>.foo.<returnValue>.<indexAccess>"
+        case result  => fail(s"Expected single assignment to x, but got $result")
+      }
+    }
+  }
 }
diff --git a/joern-cli/frontends/x2cpg/src/main/scala/io/joern/x2cpg/passes/frontend/XTypeRecovery.scala b/joern-cli/frontends/x2cpg/src/main/scala/io/joern/x2cpg/passes/frontend/XTypeRecovery.scala
@@ -716,6 +716,10 @@ abstract class RecoverForXCompilationUnit[CompilationUnitType <: AstNode](
   protected def getIndexAccessTypes(ia: Call): Set[String] = indexAccessToCollectionVar(ia) match {
     case Some(cVar) if symbolTable.contains(cVar) =>
       symbolTable.get(cVar)
+    case Some(cVar) if ia.argument(1).isCall && symbolTable.contains(CallAlias(cVar.identifier)) =>
+      symbolTable
+        .get(CallAlias(cVar.identifier))
+        .map(x => s"$x$pathSep${XTypeRecovery.DummyReturnType}$pathSep${XTypeRecovery.DummyIndexAccess}")
     case Some(cVar) if symbolTable.contains(LocalVar(cVar.identifier)) =>
       symbolTable.get(LocalVar(cVar.identifier)).map(x => s"$x$pathSep${XTypeRecovery.DummyIndexAccess}")
     case _ => Set.empty