Merge methods together

Fraunhofer-AISEC · Mar 19, 2024 · 47f0737 · 47f0737
1 parent 4b3152a
commit 47f0737
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 62 deletions.
diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilder.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilder.kt
@@ -563,14 +563,18 @@ fun <T> Literal<T>.duplicate(implicit: Boolean): Literal<T> {
     duplicate.file = this.file
     duplicate.name = this.name.clone()
     for (next in this.nextDFGEdges) {
-        if (next is ContextSensitiveDataflow)
-            duplicate.addNextDFGWithContext(next.end, next.callingContext, next.granularity)
-        else duplicate.addNextDFG(next.end, next.granularity)
+        duplicate.addNextDFG(
+            next.end,
+            next.granularity,
+            (next as? ContextSensitiveDataflow)?.callingContext
+        )
     }
     for (next in this.prevDFGEdges) {
-        if (next is ContextSensitiveDataflow)
-            duplicate.addPrevDFGWithContext(next.start, next.callingContext, next.granularity)
-        else duplicate.addNextDFG(next.start, next.granularity)
+        duplicate.addPrevDFG(
+            next.start,
+            next.granularity,
+            (next as? ContextSensitiveDataflow)?.callingContext
+        )
     }
     // TODO: This loses the properties of the edges.
     duplicate.nextEOG = this.nextEOG

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
@@ -251,22 +251,14 @@ open class Node : IVisitable<Node>, Persistable, LanguageProvider, ScopeProvider
     fun addNextDFG(
         next: Node,
         granularity: Granularity = default(),
+        callingContext: CallingContext? = null,
     ) {
-        val edge = Dataflow(this, next, granularity)
-        nextDFGEdges.add(edge)
-        next.prevDFGEdges.add(edge)
-    }
-
-    /**
-     * Adds a [Dataflow] edge from this node to [next], with the given [CallingContext] and
-     * [Granularity].
-     */
-    fun addNextDFGWithContext(
-        next: Node,
-        callingContext: CallingContext,
-        granularity: Granularity = default(),
-    ) {
-        val edge = ContextSensitiveDataflow(this, next, callingContext, granularity)
+        val edge =
+            if (callingContext != null) {
+                ContextSensitiveDataflow(this, next, callingContext, granularity)
+            } else {
+                Dataflow(this, next, granularity)
+            }
         nextDFGEdges.add(edge)
         next.prevDFGEdges.add(edge)
     }
@@ -283,26 +275,21 @@ open class Node : IVisitable<Node>, Persistable, LanguageProvider, ScopeProvider
         }
     }
 
-    /** Adds a [Dataflow] edge from [prev] node to this node, with the given [Granularity]. */
-    open fun addPrevDFG(
-        prev: Node,
-        granularity: Granularity = default(),
-    ) {
-        val edge = Dataflow(prev, this, granularity)
-        prevDFGEdges.add(edge)
-        prev.nextDFGEdges.add(edge)
-    }
-
     /**
-     * Adds a [Dataflow] edge from [prev] node to this node, with the given [CallingContext] and
-     * [Granularity].
+     * Adds a [Dataflow] edge from [prev] node to this node, with the given [Granularity] and
+     * [CallingContext].
      */
-    open fun addPrevDFGWithContext(
+    open fun addPrevDFG(
         prev: Node,
-        callingContext: CallingContext,
         granularity: Granularity = default(),
+        callingContext: CallingContext? = null,
     ) {
-        val edge = ContextSensitiveDataflow(prev, this, callingContext, granularity)
+        val edge =
+            if (callingContext != null) {
+                ContextSensitiveDataflow(prev, this, callingContext, granularity)
+            } else {
+                Dataflow(prev, this, granularity)
+            }
         prevDFGEdges.add(edge)
         prev.nextDFGEdges.add(edge)
     }
@@ -316,24 +303,16 @@ open class Node : IVisitable<Node>, Persistable, LanguageProvider, ScopeProvider
         prev.nextCDGEdges.add(edge)
     }
 
-    /** Adds a [Dataflow] edge from all [prev] nodes to this node, with the given [Granularity]. */
-    fun addAllPrevDFG(
-        prev: Collection<Node>,
-        granularity: Granularity = full(),
-    ) {
-        prev.forEach { addPrevDFG(it, granularity) }
-    }
-
     /**
-     * Adds a [Dataflow] edge from all [prev] nodes to this node, with the given [CallingContext]
-     * and [Granularity].
+     * Adds a [Dataflow] edge from all [prev] nodes to this node, with the given [Granularity] and
+     * [CallingContext] if applicable.
      */
-    fun addAllPrevDFGWithContext(
+    fun addAllPrevDFG(
         prev: Collection<Node>,
-        callingContext: CallingContext,
         granularity: Granularity = full(),
+        callingContext: CallingContext? = null,
     ) {
-        prev.forEach { addPrevDFGWithContext(it, callingContext, granularity) }
+        prev.forEach { addPrevDFG(it, granularity, callingContext) }
     }
 
     fun addAllPrevPDG(prev: Collection<Node>, dependenceType: DependenceType) {

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/statements/expressions/Reference.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/statements/expressions/Reference.kt
@@ -32,6 +32,7 @@ import de.fraunhofer.aisec.cpg.graph.Node
 import de.fraunhofer.aisec.cpg.graph.declarations.Declaration
 import de.fraunhofer.aisec.cpg.graph.declarations.ValueDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.VariableDeclaration
+import de.fraunhofer.aisec.cpg.graph.edge.CallingContext
 import de.fraunhofer.aisec.cpg.graph.edge.Granularity
 import de.fraunhofer.aisec.cpg.graph.scopes.Scope
 import de.fraunhofer.aisec.cpg.graph.types.HasType
@@ -147,8 +148,8 @@ open class Reference : Expression(), HasType.TypeObserver, HasAliases {
         return super.hashCode()
     }
 
-    override fun addPrevDFG(prev: Node, granularity: Granularity) {
-        super.addPrevDFG(prev, granularity)
+    override fun addPrevDFG(prev: Node, granularity: Granularity, callingContext: CallingContext?) {
+        super.addPrevDFG(prev, granularity, callingContext)
 
         // We want to propagate assigned types all through the previous DFG nodes. Therefore, we
         // override the DFG adding function here and add a type observer to the previous node (if it

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
@@ -357,7 +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?.addNextDFGWithContext(receiver, CallingContextIn(call))
+                call.base?.addNextDFG(receiver, callingContext = CallingContextIn(call))
             }
         }
 
@@ -373,12 +373,12 @@ object Util {
                 if (param.isVariadic) {
                     while (j < arguments.size) {
                         // map all the following arguments to this variadic param
-                        param.addPrevDFGWithContext(arguments[j], CallingContextIn(call))
+                        param.addPrevDFG(arguments[j], callingContext = CallingContextIn(call))
                         j++
                     }
                     break
                 } else {
-                    param.addPrevDFGWithContext(arguments[j], CallingContextIn(call))
+                    param.addPrevDFG(arguments[j], callingContext = CallingContextIn(call))
                 }
             }
             j++

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
@@ -127,9 +127,9 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
                         Pair(it, key) in edgePropertiesMap &&
                             edgePropertiesMap[Pair(it, key)] is CallingContext
                     ) {
-                        key.addPrevDFGWithContext(
+                        key.addPrevDFG(
                             it,
-                            (edgePropertiesMap[Pair(it, key)] as CallingContext)
+                            callingContext = (edgePropertiesMap[Pair(it, key)] as? CallingContext)
                         )
                     } else {
                         key.addPrevDFG(it)
@@ -471,7 +471,6 @@ open class ControlFlowSensitiveDFGPass(ctx: TranslationContext) : EOGStarterPass
                             }
                         doubleState.declarationsState[arg?.refersTo] =
                             PowersetLattice(identitySetOf(param))
-                        // TODO: Map probably needs param and a second node!
                         edgePropertiesMap[Pair(param, null)] = CallingContextOut(currentNode)
                     }
                 }

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
@@ -70,10 +70,10 @@ class DFGPass(ctx: TranslationContext) : ComponentPass(ctx) {
                     if (param == (invoked as? MethodDeclaration)?.receiver) {
                         (call as? MemberCallExpression)
                             ?.base
-                            ?.addPrevDFGWithContext(param, CallingContextOut(call))
+                            ?.addPrevDFG(param, callingContext = CallingContextOut(call))
                     } else if (param is ParameterDeclaration) {
                         val arg = call.arguments[param.argumentIndex]
-                        arg.addPrevDFGWithContext(param, CallingContextOut(call))
+                        arg.addPrevDFG(param, callingContext = CallingContextOut(call))
                         (arg as? Reference)?.let {
                             it.access = AccessValues.READWRITE
                             it.refersTo?.let { it1 -> it.addNextDFG(it1) }
@@ -463,7 +463,7 @@ class DFGPass(ctx: TranslationContext) : ComponentPass(ctx) {
         } else if (call.invokes.isNotEmpty()) {
             call.invokes.forEach {
                 Util.attachCallParameters(it, call)
-                call.addPrevDFGWithContext(it, CallingContextOut(call))
+                call.addPrevDFG(it, callingContext = CallingContextOut(call))
                 if (it.isInferred) {
                     callsInferredFunctions.add(call)
                 }

diff --git a/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilderTest.kt b/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ExpressionBuilderTest.kt
@@ -43,7 +43,7 @@ class ExpressionBuilderTest {
         val node2 = Reference()
         val granularity = PartialDataflowGranularity(FieldDeclaration())
         val callingContextIn = CallingContextIn(CallExpression())
-        node1.addPrevDFGWithContext(node2, callingContextIn, granularity)
+        node1.addPrevDFG(node2, granularity, callingContextIn)
 
         val clone = node1.duplicate(false)
         val clonedPrevDFG = clone.prevDFGEdges.single()
@@ -60,7 +60,7 @@ class ExpressionBuilderTest {
         val node2 = Reference()
         val granularity = PartialDataflowGranularity(FieldDeclaration())
         val callingContextIn = CallingContextIn(CallExpression())
-        node1.addNextDFGWithContext(node2, callingContextIn, granularity)
+        node1.addNextDFG(node2, granularity, callingContextIn)
 
         val clone = node1.duplicate(false)
         val clonedPrevDFG = clone.nextDFGEdges.single()