New property edges

Fraunhofer-AISEC · Mar 7, 2024 · d34e7ab · d34e7ab
1 parent eed6009
commit d34e7ab
Show file tree

Hide file tree

Showing 5 changed files with 101 additions and 30 deletions.
diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Node.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Node.kt
@@ -257,6 +257,20 @@ open class Node : IVisitable<Node>, Persistable, LanguageProvider, ScopeProvider
         next.prevDFGEdges.add(edge)
     }
 
+    /**
+     * Adds a [Dataflow] edge from this node to [next], with the given [CallingContext] and
+     * [Granularity].
+     */
+    fun addNextDFGContext(
+        next: Node,
+        callingContext: CallingContext,
+        granularity: Granularity = default(),
+    ) {
+        val edge = ContextsensitiveDataflow(this, next, callingContext, granularity)
+        nextDFGEdges.add(edge)
+        next.prevDFGEdges.add(edge)
+    }
+
     fun removeNextDFG(next: Node?) {
         if (next != null) {
             val thisRemove =
@@ -279,6 +293,20 @@ open class Node : IVisitable<Node>, Persistable, LanguageProvider, ScopeProvider
         prev.nextDFGEdges.add(edge)
     }
 
+    /**
+     * Adds a [Dataflow] edge from [prev] node to this node, with the given [CallingContext] and
+     * [Granularity].
+     */
+    open fun addPrevDFGContext(
+        prev: Node,
+        callingContext: CallingContext,
+        granularity: Granularity = default(),
+    ) {
+        val edge = ContextsensitiveDataflow(prev, this, callingContext, granularity)
+        prevDFGEdges.add(edge)
+        prev.nextDFGEdges.add(edge)
+    }
+
     fun addPrevCDG(
         prev: Node,
         properties: MutableMap<Properties, Any?> = EnumMap(Properties::class.java)
@@ -296,6 +324,18 @@ open class Node : IVisitable<Node>, Persistable, LanguageProvider, ScopeProvider
         prev.forEach { addPrevDFG(it, granularity) }
     }
 
+    /**
+     * Adds a [Dataflow] edge from all [prev] nodes to this node, with the given [CallingContext]
+     * and [Granularity].
+     */
+    fun addAllPrevDFGContext(
+        prev: Collection<Node>,
+        callingContext: CallingContext,
+        granularity: Granularity = full(),
+    ) {
+        prev.forEach { addPrevDFGContext(it, callingContext, granularity) }
+    }
+
     fun addAllPrevPDG(prev: Collection<Node>, dependenceType: DependenceType) {
         addAllPrevPDGEdges(prev.map { PropertyEdge(it, this) }, dependenceType)
     }

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edge/Dataflow.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edge/Dataflow.kt
@@ -30,6 +30,7 @@ import de.fraunhofer.aisec.cpg.graph.declarations.Declaration
 import de.fraunhofer.aisec.cpg.graph.declarations.FieldDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.TupleDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.VariableDeclaration
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.CallExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.MemberExpression
 import org.neo4j.ogm.annotation.RelationshipEntity
 
@@ -80,11 +81,43 @@ fun partial(target: Declaration?): PartialDataflowGranularity {
  * [granularity].
  */
 @RelationshipEntity
-class Dataflow(
+open class Dataflow(
     start: Node,
     end: Node,
     /** The granularity of this dataflow. */
-    val granularity: Granularity = default(),
+    val granularity: Granularity = default()
 ) : PropertyEdge<Node>(start, end) {
     override val label: String = "DFG"
 }
+
+sealed interface CallingContext
+
+class CallingContextIn(
+    /** The call expression that affects this dataflow edge. */
+    val callExpression: CallExpression
+) : CallingContext
+
+class CallingContextOut(
+    /** The call expression that affects this dataflow edge. */
+    val callExpression: CallExpression
+) : CallingContext
+
+fun callIn(callExpression: CallExpression): CallingContext {
+    return CallingContextIn(callExpression)
+}
+
+/**
+ * This edge class defines a flow of data between [start] and [end]. The flow must have a
+ * [callingContext] and can have a certain [granularity].
+ */
+@RelationshipEntity
+class ContextsensitiveDataflow(
+    start: Node,
+    end: Node,
+    /** The calling context affecting this dataflow. */
+    val callingContext: CallingContext,
+    /** The granularity of this dataflow. */
+    granularity: Granularity,
+) : Dataflow(start, end, granularity) {
+    override val label: String = "DFG"
+}
diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/Util.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/Util.kt
@@ -29,6 +29,7 @@ import de.fraunhofer.aisec.cpg.frontends.LanguageFrontend
 import de.fraunhofer.aisec.cpg.graph.Node
 import de.fraunhofer.aisec.cpg.graph.declarations.FunctionDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.MethodDeclaration
+import de.fraunhofer.aisec.cpg.graph.edge.CallingContextIn
 import de.fraunhofer.aisec.cpg.graph.edge.Properties
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
 import de.fraunhofer.aisec.cpg.sarif.PhysicalLocation
@@ -356,10 +357,7 @@ object Util {
         // Add an incoming DFG edge from a member call's base to the method's receiver
         if (target is MethodDeclaration && call is MemberCallExpression && !call.isStatic) {
             target.receiver?.let { receiver ->
-                call.base?.addNextDFG(
-                    receiver,
-                    mutableMapOf(Pair(Properties.CALLING_CONTEXT_IN, call))
-                )
+                call.base?.addNextDFGContext(receiver, CallingContextIn(call))
             }
         }
 
@@ -375,18 +373,12 @@ object Util {
                 if (param.isVariadic) {
                     while (j < arguments.size) {
                         // map all the following arguments to this variadic param
-                        param.addPrevDFG(
-                            arguments[j],
-                            mutableMapOf(Pair(Properties.CALLING_CONTEXT_IN, call))
-                        )
+                        param.addPrevDFGContext(arguments[j], CallingContextIn(call))
                         j++
                     }
                     break
                 } else {
-                    param.addPrevDFG(
-                        arguments[j],
-                        mutableMapOf(Pair(Properties.CALLING_CONTEXT_IN, call))
-                    )
+                    param.addPrevDFGContext(arguments[j], CallingContextIn(call))
                 }
             }
             j++

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/DFGPass.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/passes/DFGPass.kt
@@ -28,6 +28,7 @@ package de.fraunhofer.aisec.cpg.passes
 import de.fraunhofer.aisec.cpg.TranslationContext
 import de.fraunhofer.aisec.cpg.graph.*
 import de.fraunhofer.aisec.cpg.graph.declarations.*
+import de.fraunhofer.aisec.cpg.graph.edge.CallingContextOut
 import de.fraunhofer.aisec.cpg.graph.edge.partial
 import de.fraunhofer.aisec.cpg.graph.statements.*
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.*
@@ -66,19 +67,12 @@ class DFGPass(ctx: TranslationContext) : ComponentPass(ctx) {
                     functionSummaries.functionToChangedParameters[invoked] ?: mapOf()
                 for ((param, _) in changedParams) {
                     if (param == (invoked as? MethodDeclaration)?.receiver) {
-                        (call as? MemberCallExpression)?.base?.let { base ->
-                            base.addPrevDFG(
-                                param,
-                                mutableMapOf(Pair(Properties.CALLING_CONTEXT_OUT, call))
-                            )
-                            // (base as? Reference)?.access = AccessValues.READWRITE
-                        }
+                        (call as? MemberCallExpression)
+                            ?.base
+                            ?.addPrevDFGContext(param, CallingContextOut(call))
                     } else if (param is ParameterDeclaration) {
                         val arg = call.arguments[param.argumentIndex]
-                        arg.addPrevDFG(
-                            param,
-                            mutableMapOf(Pair(Properties.CALLING_CONTEXT_OUT, call))
-                        )
+                        arg.addPrevDFGContext(param, CallingContextOut(call))
                         // (arg as? Reference)?.access = AccessValues.READWRITE
                     }
                 }
@@ -465,7 +459,7 @@ class DFGPass(ctx: TranslationContext) : ComponentPass(ctx) {
         } else if (call.invokes.isNotEmpty()) {
             call.invokes.forEach {
                 Util.attachCallParameters(it, call)
-                call.addPrevDFG(it)
+                call.addPrevDFGContext(it, CallingContextOut(call))
                 if (it.isInferred) {
                     callsInferredFunctions.add(call)
                 }

diff --git a/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/enhancements/DFGFunctionSummariesTest.kt b/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/enhancements/DFGFunctionSummariesTest.kt
@@ -32,7 +32,9 @@ import de.fraunhofer.aisec.cpg.TranslationResult
 import de.fraunhofer.aisec.cpg.frontends.TestLanguage
 import de.fraunhofer.aisec.cpg.graph.*
 import de.fraunhofer.aisec.cpg.graph.builder.*
-import de.fraunhofer.aisec.cpg.graph.edge.Properties
+import de.fraunhofer.aisec.cpg.graph.edge.CallingContextIn
+import de.fraunhofer.aisec.cpg.graph.edge.CallingContextOut
+import de.fraunhofer.aisec.cpg.graph.edge.ContextsensitiveDataflow
 import de.fraunhofer.aisec.cpg.graph.functions
 import de.fraunhofer.aisec.cpg.graph.pointer
 import de.fraunhofer.aisec.cpg.passes.inference.DFGFunctionSummaries
@@ -79,14 +81,23 @@ class DFGFunctionSummariesTest {
 
         assertEquals(1, argA.nextDFG.size)
         assertEquals(2, argA.prevDFG.size)
+        /*
+        The flows should be as follows:
+        VariableDeclaration["a"] -> Reference["a" (argument of call)] -CallingContextIn-> ParameterDeclaration -CallingContextOut-> Reference["a" (return)]
+         */
 
         val nextDfg = argA.nextDFGEdges.single()
-        assertEquals(call, nextDfg.getProperty(Properties.CALLING_CONTEXT_IN))
+        assertEquals(
+            call,
+            ((nextDfg as? ContextsensitiveDataflow)?.callingContext as? CallingContextIn)
+                ?.callExpression
+        )
         assertEquals(param0, nextDfg.end)
 
         val prevDfgThroughFunction =
             argA.prevDFGEdges.singleOrNull {
-                it.containsProperties(mapOf(Pair(Properties.CALLING_CONTEXT_OUT, call)))
+                ((it as? ContextsensitiveDataflow)?.callingContext as? CallingContextOut)
+                    ?.callExpression != call
             }
         assertNotNull(prevDfgThroughFunction)
         assertEquals(param0, prevDfgThroughFunction.start)
@@ -96,7 +107,8 @@ class DFGFunctionSummariesTest {
 
         val prevDfgNotThroughFunction =
             argA.prevDFGEdges.singleOrNull {
-                !it.containsProperties(mapOf(Pair(Properties.CALLING_CONTEXT_OUT, call)))
+                ((it as? ContextsensitiveDataflow)?.callingContext as? CallingContextOut)
+                    ?.callExpression != call
             }
         assertNotNull(prevDfgNotThroughFunction)
         assertEquals(variableDeclA, prevDfgNotThroughFunction.start)