Not relying on AST enymore

Fraunhofer-AISEC · Aug 14, 2024 · 9208655 · 9208655
1 parent a990cf0
commit 9208655
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 73 deletions.
diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/ast/AstEdge.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/ast/AstEdge.kt
@@ -75,7 +75,7 @@ fun <NodeType : Node> Node.astEdgeOf(
 }
 
 /** This property edge list describes elements that are AST children of a node. */
-open class AstEdges<NodeType : Node, PropertyEdgeType : Edge<NodeType>>(
+open class AstEdges<NodeType : Node, PropertyEdgeType : AstEdge<NodeType>>(
     thisRef: Node,
     onAdd: ((PropertyEdgeType) -> Unit)? = null,
     onRemove: ((PropertyEdgeType) -> Unit)? = null,
@@ -89,4 +89,4 @@ open class AstEdges<NodeType : Node, PropertyEdgeType : Edge<NodeType>>(
         init = init,
         onAdd = onAdd,
         onRemove = onRemove,
-    )
+    ) {}
diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeCollection.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeCollection.kt
@@ -27,6 +27,7 @@ package de.fraunhofer.aisec.cpg.graph.edges.collections
 
 import de.fraunhofer.aisec.cpg.graph.Node
 import de.fraunhofer.aisec.cpg.graph.edges.Edge
+import de.fraunhofer.aisec.cpg.passes.EdgeType
 
 /**
  * This interfaces is an extension of [MutableCollection] that holds specific functions for the
@@ -139,7 +140,7 @@ interface EdgeCollection<
      * Note, that is an immutable list and only a snapshot. If you want a magic container that is in
      * sync with this [EdgeCollection], please use [unwrap].
      */
-    fun toNodeCollection(): Collection<NodeType>
+    fun toNodeCollection(predicate: ((EdgeType) -> Boolean)? = null): Collection<NodeType>
 
     /**
      * Returns an [UnwrappedEdgeCollection] magic container which holds a structure that provides
@@ -170,14 +171,23 @@ interface EdgeCollection<
 @Suppress("UNCHECKED_CAST")
 internal fun <
     NodeType : Node,
+    EdgeType : Edge<NodeType>,
     CollectionType : MutableCollection<NodeType>
 > internalToNodeCollection(
-    edges: EdgeCollection<NodeType, out Edge<NodeType>>,
+    edges: EdgeCollection<NodeType, EdgeType>,
     outgoing: Boolean = true,
+    predicate: ((EdgeType) -> Boolean)? = null,
     createCollection: () -> CollectionType
 ): CollectionType {
     var unwrapped = createCollection()
-    edges.mapTo(unwrapped) { if (outgoing) it.end else it.start as NodeType }
+    for (edge in edges) {
+        if (predicate != null && !predicate(edge)) {
+            continue
+        }
+
+        unwrapped += if (outgoing) edge.end else edge.start as NodeType
+    }
+
     return unwrapped
 }
 

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeList.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeList.kt
@@ -106,8 +106,8 @@ abstract class EdgeList<NodeType : Node, EdgeType : Edge<NodeType>>(
         }
     }
 
-    override fun toNodeCollection(): List<NodeType> {
-        return internalToNodeCollection(this, outgoing, ::ArrayList)
+    override fun toNodeCollection(predicate: ((EdgeType) -> Boolean)?): List<NodeType> {
+        return internalToNodeCollection(this, outgoing, predicate, ::ArrayList)
     }
 
     /**

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeSet.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeSet.kt
@@ -80,8 +80,8 @@ abstract class EdgeSet<NodeType : Node, EdgeType : Edge<NodeType>>(
         edges.forEach { handleOnRemove(it) }
     }
 
-    override fun toNodeCollection(): MutableSet<NodeType> {
-        return internalToNodeCollection(this, outgoing, ::HashSet)
+    override fun toNodeCollection(predicate: ((EdgeType) -> Boolean)?): MutableSet<NodeType> {
+        return internalToNodeCollection(this, outgoing, predicate, ::HashSet)
     }
 
     /**

diff --git a/...core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeSingletonList.kt b/...core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/edges/collections/EdgeSingletonList.kt
@@ -107,8 +107,8 @@ class EdgeSingletonList<NodeType : Node, NullableNodeType : NodeType?, EdgeType
         set(value) {}
 
     @Suppress("UNCHECKED_CAST")
-    override fun toNodeCollection(): Collection<NodeType> {
-        var elements = this.toList()
+    override fun toNodeCollection(predicate: ((EdgeType) -> Boolean)?): Collection<NodeType> {
+        var elements = predicate?.let { toList().filter(it) } ?: toList()
         return elements.map {
             if (outgoing) {
                 it.end

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/SubgraphWalker.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/SubgraphWalker.kt
@@ -30,15 +30,14 @@ import de.fraunhofer.aisec.cpg.frontends.LanguageFrontend
 import de.fraunhofer.aisec.cpg.graph.AST
 import de.fraunhofer.aisec.cpg.graph.Node
 import de.fraunhofer.aisec.cpg.graph.declarations.RecordDeclaration
-import de.fraunhofer.aisec.cpg.graph.edges.Edge
+import de.fraunhofer.aisec.cpg.graph.edges.ast.AstEdge
 import de.fraunhofer.aisec.cpg.graph.edges.collections.EdgeCollection
 import de.fraunhofer.aisec.cpg.processing.strategy.Strategy
 import java.lang.annotation.AnnotationFormatError
 import java.lang.reflect.Field
 import java.util.*
 import java.util.function.BiConsumer
 import java.util.function.Consumer
-import org.neo4j.ogm.annotation.Relationship
 import org.slf4j.LoggerFactory
 
 /** A type for a node visitor callback for the [SubgraphWalker]. */
@@ -67,7 +66,7 @@ object SubgraphWalker {
             }
             val fields = ArrayList<Field>()
             fields.addAll(getAllFields(classType.superclass))
-            fields.addAll(listOf(*classType.declaredFields))
+            fields.addAll(listOf(*classType.declaredFields).filter { it.name.contains("Edge") })
 
             // update the cache
             fieldCache[cacheKey] = fields
@@ -93,54 +92,9 @@ object SubgraphWalker {
         if (node == null) return children
         val classType: Class<*> = node.javaClass
 
-        /*for (member in node::class.members) {
-            val subGraph = member.findAnnotation<SubGraph>()
-            if (subGraph != null && listOf(*subGraph.value).contains("AST")) {
-                val old = member.isAccessible
-
-                member.isAccessible = true
-
-                val obj = member.call(node)
-
-                // skip, if null
-                if (obj == null) {
-                    continue
-                }
-
-                member.isAccessible = old
-
-                var outgoing = true // default
-                var relationship = member.findAnnotation<Relationship>()
-                if (relationship != null) {
-                    outgoing =
-                        relationship.direction ==
-                                Relationship.Direction.OUTGOING)
-                }
-                if (checkForPropertyEdge(field, obj)) {
-                    obj = unwrap(obj as List<PropertyEdge<Node>>, outgoing)
-                }
-                when (obj) {
-                    is Node -> {
-                        children.add(obj)
-                    }
-                    is Collection<*> -> {
-                        children.addAll(obj as Collection<Node>)
-                    }
-                    else -> {
-                        throw AnnotationFormatError(
-                            "Found @field:SubGraph(\"AST\") on field of type " +
-                                    obj.javaClass +
-                                    " but can only used with node graph classes or collections of graph nodes"
-                        )
-                    }
-                }
-            }
-        }*/
-
         // We currently need to stick to pure Java reflection, since Kotlin reflection
         // is EXTREMELY slow. See https://youtrack.jetbrains.com/issue/KT-32198
         for (field in getAllFields(classType)) {
-            field.getAnnotation(AST::class.java) ?: continue
             try {
                 // We need to synchronize access to the field, because otherwise different
                 // threads might restore the isAccessible property while this thread is still
@@ -156,22 +110,9 @@ object SubgraphWalker {
                         obj
                     } ?: continue
 
-                // skip, if null
-                var outgoing = true // default
-                if (field.getAnnotation(Relationship::class.java) != null) {
-                    outgoing =
-                        (field.getAnnotation(Relationship::class.java).direction ==
-                            Relationship.Direction.OUTGOING)
-                }
                 when (obj) {
-                    is Edge<*> -> {
-                        // This will most likely also not happen because of a bug in Neo4J we cannot
-                        // use our edge class directly
-                        children.add(obj.end)
-                    }
                     is EdgeCollection<*, *> -> {
-                        // This is the default for edge lists
-                        children.addAll(obj.toNodeCollection())
+                        children.addAll(obj.toNodeCollection({ it is AstEdge<*> }))
                     }
                     else -> {
                         throw AnnotationFormatError(