From cff0ff6be93bc255c935b31c7053b8b819f96da3 Mon Sep 17 00:00:00 2001
From: Alexander Kuechler <alexander.kuechler@aisec.fraunhofer.de>
Date: Wed, 11 Oct 2023 10:31:38 +0200
Subject: [PATCH] More fixes

---
 .../cpg/passes/ControlFlowSensitiveDFGPass.kt | 14 ++++++++---
 .../frontends/python/PythonFrontendTest.kt    | 24 +++++++++----------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/ControlFlowSensitiveDFGPass.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/ControlFlowSensitiveDFGPass.kt
index 7bd1c16dcc0..2cbc3af7752 100644
--- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/ControlFlowSensitiveDFGPass.kt
+++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/ControlFlowSensitiveDFGPass.kt
@@ -195,7 +195,11 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
                 // the other steps
                 state.push(currentNode, it)
             }
-        } else if ((currentNode as? Reference)?.access == AccessValues.READWRITE) {
+        } else if (
+            (currentNode as? Reference)?.access == AccessValues.READWRITE &&
+                currentNode.nextEOG.none { isCompoundAssignment(it) || isIncOrDec(it) }
+        ) {
+            /* This branch collects all READWRITE accesses which are not handled separately as compoundAssignment or inc/dec unary operation. This could for example be a pointer passed to an unknown function which is modified in this function but other things are also possible. */
             // We can only find a change if there's a state for the variable
             doubleState.declarationsState[currentNode.refersTo]?.let {
                 // We only read the variable => Get previous write which have been collected in
@@ -206,7 +210,11 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
             // there was probably some other kind of DFG edge into the reference
             doubleState.declarationsState[currentNode.refersTo] =
                 PowersetLattice(setOf(currentNode))
-        } else if ((currentNode as? Reference)?.access == AccessValues.WRITE) {
+        } else if (
+            (currentNode as? Reference)?.access == AccessValues.WRITE &&
+                currentNode.nextEOG.none { it is ForEachStatement }
+        ) {
+            /* Also here, we want/have to filter out variables in ForEachStatements because this must be handled separately.  */
             // We write to the variable => Update the declarationState accordingly because
             // there was probably some other kind of DFG edge into the reference
             doubleState.declarationsState[currentNode.refersTo] =
@@ -229,7 +237,7 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : TranslationUni
                             null
                         }
                     }
-                    else -> currentNode.variable
+                    else -> variable
                 }
 
             // We wrote something to this variable declaration
diff --git a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt
index 2ca7c4fb665..b86448db62a 100644
--- a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt
+++ b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt
@@ -1043,35 +1043,35 @@ class PythonFrontendTest : BaseTest() {
         assertNotNull(barCall)
 
         // no dataflow from var declaration to loop variable because it's a write access
-        assert((firstLoop.variable?.prevDFG?.contains(varDefinedBeforeLoop) == false))
+        assertFalse(
+            varDefinedBeforeLoop in (firstLoop.variable?.prevDFG ?: setOf(varDefinedBeforeLoop))
+        )
 
         // dataflow from range call to loop variable
         val firstLoopIterable = firstLoop.iterable as? CallExpression
         assertNotNull(firstLoopIterable)
-        assert((firstLoop.variable?.prevDFG?.contains((firstLoopIterable)) == true))
+        assertContains((firstLoop.variable?.prevDFG ?: setOf()), firstLoopIterable)
 
         // dataflow from var declaration to loop iterable call
-        assert(
-            firstLoopIterable.arguments.firstOrNull()?.prevDFG?.contains(varDefinedBeforeLoop) ==
-                true
+        assertContains(
+            firstLoopIterable.arguments.firstOrNull()?.prevDFG ?: setOf(),
+            varDefinedBeforeLoop
         )
 
         // dataflow from first loop to foo call
         val loopVar = firstLoop.variable as? Reference
         assertNotNull(loopVar)
-        assert(fooCall.arguments.first().prevDFG.contains(loopVar))
+        assertContains(fooCall.arguments.first().prevDFG, loopVar)
 
         // dataflow from var declaration to foo call (in case for loop is not executed)
-        assert(fooCall.arguments.first().prevDFG.contains(varDefinedBeforeLoop))
+        assertContains(fooCall.arguments.first().prevDFG, varDefinedBeforeLoop)
 
         // dataflow from range call to loop variable
         val secondLoopIterable = secondLoop.iterable as? CallExpression
         assertNotNull(secondLoopIterable)
-        assert(
-            ((secondLoop.variable as DeclarationStatement)
-                .singleDeclaration
-                ?.prevDFG
-                ?.contains((secondLoopIterable)) == true)
+        assertContains(
+            (secondLoop.variable as DeclarationStatement).singleDeclaration?.prevDFG ?: setOf(),
+            secondLoopIterable
         )
 
         // dataflow from second loop var to bar call