[dataflowengineoss] globalFromLiteral taking into account 3AC assignm…

…ents (#4649)

dataflowengineoss/src/main/scala/io/joern/dataflowengineoss/package.scala

@@ -2,6 +2,8 @@ package io.joern
 
 import io.shiftleft.codepropertygraph.generated.nodes.{Declaration, Expression, Identifier, Literal}
 import io.shiftleft.semanticcpg.language.*
+import io.shiftleft.semanticcpg.language.operatorextension.OpNodes.Assignment
+import io.shiftleft.semanticcpg.utils.MemberAccess.isFieldAccess
 
 package object dataflowengineoss {
 
@@ -14,10 +16,23 @@ package object dataflowengineoss {
     * @return
     *   the LHS of the assignment
     */
-  def globalFromLiteral(lit: Literal, recursive: Boolean = true): Iterator[Expression] = lit.start
-    .where(_.method.isModule)
-    .flatMap(t => if (recursive) t.inAssignment else t.inCall.isAssignment)
-    .target
+  def globalFromLiteral(lit: Literal, recursive: Boolean = true): Iterator[Expression] = {
+
+    /** Frontends often create three-address code representations of compound literals, e.g. in pysrc2cpg the dictionary
+      * literal `{"x": y}` is lowered as a block `{tmp0 = {}; tmp0["x"] = y; tmp0}`.
+      *
+      * In an assignment `foo = {"x": "y"}`, if [[lit]] is "y", we don't want to pick the intermediate assignment
+      * `tmp0["x"] = "y"`, since `tmp0` is never a global/module-level variable. Instead, we want to pick `foo`.
+      */
+    def skipLowLevelAssignments(assignments: Iterator[Assignment]): Iterator[Assignment] = {
+      assignments.repeat(_.parentBlock.inCall.isAssignment)(_.emit).lastOption.fold(assignments)(_.iterator)
+    }
+
+    lit.start
+      .where(_.method.isModule)
+      .flatMap(t => if (recursive) t.inAssignment else skipLowLevelAssignments(t.inCall.isAssignment))
+      .target
+  }
 
   def identifierToFirstUsages(node: Identifier): List[Identifier] = node.refsTo.flatMap(identifiersFromCapturedScopes).l
 

dataflowengineoss/src/main/scala/io/joern/dataflowengineoss/queryengine/SourcesToStartingPoints.scala

@@ -194,7 +194,7 @@ abstract class BaseSourceToStartingPoints extends Callable[Unit] {
         val usageInput = targetsToClassIdentifierPair(literalToInitializedMembers(lit), src)
         val uses       = usages(usageInput)
         val globals = globalFromLiteral(lit, recursive = false).flatMap {
-          case x: Identifier if x.isModuleVariable => x :: moduleVariableToFirstUsagesAcrossProgram(x)
+          case x: Identifier if x.isModuleVariable => moduleVariableToFirstUsagesAcrossProgram(x)
           case x                                   => x :: Nil
         }
         (lit :: (uses ++ globals), usageInput)

joern-cli/frontends/pysrc2cpg/src/test/scala/io/joern/pysrc2cpg/dataflow/DataFlowTests.scala

@@ -512,6 +512,53 @@ class DataFlowTests extends PySrc2CpgFixture(withOssDataflow = true) {
     flow100 shouldBe List(("x = 100", 3), ("tmp0['Property'] = x", 4), ("tmp0", 4))
   }
 
+  "flow from literal in an imported dictionary literal to `print`" in {
+    val cpg = code(
+      """
+        |from p1.bar import baz
+        |print(baz)
+        |""".stripMargin,
+      "foo.py"
+    ).moreCode(
+      """
+        |baz = {"Property1": "foo"}
+        |""".stripMargin,
+      "p1/bar.py"
+    )
+    val source     = cpg.literal("\"foo\"")
+    val sink       = cpg.call("print").argument(1)
+    val List(flow) = sink.reachableByFlows(source).map(flowToResultPairs).l
+
+    flow shouldBe List(("tmp0[\"Property1\"] = \"foo\"", 2), ("baz = import(p1.bar, baz)", 2), ("print(baz)", 3))
+  }
+
+  "flow from literal in an imported method-returned dictionary to `print`" in {
+    val cpg = code(
+      """
+        |import bar
+        |print(bar.baz())
+        |""".stripMargin,
+      "main.py"
+    ).moreCode(
+      """
+        |def baz():
+        |  return {"Property": "foo"}
+        |""".stripMargin,
+      "bar.py"
+    )
+    val source     = cpg.literal("\"foo\"")
+    val sink       = cpg.call("print").argument(1)
+    val List(flow) = sink.reachableByFlows(source).map(flowToResultPairs).l
+
+    flow shouldBe List(
+      ("tmp0[\"Property\"] = \"foo\"", 3),
+      ("tmp0", 3),
+      ("return {\"Property\": \"foo\"}", 3),
+      ("RET", 2),
+      ("bar.baz()", 3)
+    )
+  }
+
   "flow from global variable defined in imported file and used as argument to `print`" in {
     val cpg = code("""
         |from models import FOOBAR