✅ Type Handling Tests (#54)

* TYPE and TYPE_DECL handling

* ✨ Added TYPE_DECL to getProgramStructure

* ⚡ Improved TG queries with v2 syntax

* ✨ TypeDecl is in properly

* ✨ TypeDecl tests are green

Co-authored-by: David Baker Effendi <[email protected]>

cpgconv/src/main/scala/io/github/plume/oss/Traversals.scala

@@ -1,10 +1,9 @@
 package io.github.plume.oss
 
 import java.util
-
 import io.shiftleft.codepropertygraph.Cpg
 import io.shiftleft.codepropertygraph.generated.EdgeTypes
-import io.shiftleft.codepropertygraph.generated.nodes.{AstNode, HasOrder, StoredNode}
+import io.shiftleft.codepropertygraph.generated.nodes.{AstNode, HasOrder, StoredNode, TypeDecl}
 import io.shiftleft.semanticcpg.language._
 import overflowdb.{Edge, Graph}
 import io.shiftleft.codepropertygraph.generated.nodes
@@ -55,6 +54,8 @@ object Traversals {
     (edgesFromFile ++ edgesFromNamespaceBlock).asJava
   }
 
+  def getTypeDecls(graph: Graph):util.List[TypeDecl] = Cpg(graph).typeDecl.l.asJava
+
   def getNeighbours(graph: Graph, nodeId: Long): util.List[Edge] = {
     Cpg(graph)
       .id[StoredNode](nodeId)

gradle.properties

@@ -6,7 +6,7 @@ jgVersion=0.5.3
 neo4jDriverVersion=4.2.0
 khttpVersion=1.0.0
 jacksonVersion=2.12.0
-shiftleftVersion=1.3.61
+shiftleftVersion=1.3.62
 sootVersion=4.2.1
 lz4Version=1.7.1
 

plume/src/main/kotlin/io/github/plume/oss/Extractor.kt

@@ -32,12 +32,16 @@ import io.github.plume.oss.util.ResourceCompilationUtil.compileJavaFiles
 import io.github.plume.oss.util.ResourceCompilationUtil.deleteClassFiles
 import io.github.plume.oss.util.ResourceCompilationUtil.moveClassFiles
 import io.github.plume.oss.util.SootToPlumeUtil
+import io.shiftleft.codepropertygraph.generated.EdgeTypes
+import io.shiftleft.codepropertygraph.generated.NodeKeyNames.FULL_NAME
+import io.shiftleft.codepropertygraph.generated.NodeTypes.TYPE_DECL
 import io.shiftleft.codepropertygraph.generated.nodes.*
 import org.apache.logging.log4j.LogManager
 import org.apache.logging.log4j.Logger
 import overflowdb.Graph
 import overflowdb.Node
 import soot.*
+import soot.jimple.*
 import soot.jimple.spark.SparkTransformer
 import soot.jimple.toolkits.callgraph.CHATransformer
 import soot.jimple.toolkits.callgraph.Edge
@@ -248,6 +252,27 @@ class Extractor(val driver: IDriver) {
             }
             .distinct().toList().let { if (it.size >= 100000) it.parallelStream() else it.stream() }
             .filter { !it.isPhantom }.map { BriefUnitGraph(it.retrieveActiveBody()) }.toList()
+        // Build types from fields
+        classStream.asSequence().map { it.fields }.flatten().map { it.type }.distinct()
+            .filter { t ->
+                !classStream.any { it.name == t.toQuotedString() }
+                        &&
+                        !programStructure.nodes { n -> n == TYPE_DECL }.asSequence()
+                            .any { n -> n.property(FULL_NAME) == t.toQuotedString() }
+            }
+            .map(SootToPlumeUtil::buildTypeDeclaration)
+            .forEach(driver::addVertex)
+        // Build types from locals
+        graphs.asSequence().map { it.body.locals + it.body.parameterLocals }.flatten().map { it.type }
+            .distinct()
+            .filter { t ->
+                !classStream.any { it.name == t.toQuotedString() }
+                        &&
+                        !programStructure.nodes { n -> n == TYPE_DECL }.asSequence()
+                            .any { n -> n.property(FULL_NAME) == t.toQuotedString() }
+            }
+            .map(SootToPlumeUtil::buildTypeDeclaration)
+            .forEach(driver::addVertex)
         // Construct the CPGs for methods
         graphs.map(this::constructCPG)
             .toList().asSequence()
@@ -281,7 +306,14 @@ class Extractor(val driver: IDriver) {
         if (programStructure.nodes { it == ODBFile.Label() }.asSequence().none { it.property("NAME") == cls.name }) {
             logger.debug("Building file, namespace, and type declaration for ${cls.name}")
             SootToPlumeUtil.buildClassStructure(cls, driver)
-            SootToPlumeUtil.buildTypeDeclaration(cls, driver)
+            val typeDecl = SootToPlumeUtil.buildTypeDeclaration(cls.type, false)
+            addSootToPlumeAssociation(cls, typeDecl)
+            cls.fields.forEachIndexed { i, field ->
+                SootToPlumeUtil.projectMember(field, i + 1).let { memberVertex ->
+                    driver.addEdge(typeDecl, memberVertex, EdgeTypes.AST)
+                    addSootToPlumeAssociation(field, memberVertex)
+                }
+            }
         }
     }
 
@@ -313,7 +345,7 @@ class Extractor(val driver: IDriver) {
         val currentFileHash = getFileHashPair(cls)
         val files = programStructure.nodes { it == ODBFile.Label() }.asSequence()
         logger.debug("Looking for existing file vertex for ${cls.name} from given file hash $currentFileHash")
-        files.firstOrNull { it.property("NAME") == SootToPlumeUtil.sootClassToFileName(cls)}?.let { fileV ->
+        files.firstOrNull { it.property("NAME") == SootToPlumeUtil.sootClassToFileName(cls) }?.let { fileV ->
             if (fileV.property("HASH") != currentFileHash) {
                 logger.debug("Existing class was found and file hashes do not match, marking for rebuild.")
                 // Rebuild

plume/src/main/kotlin/io/github/plume/oss/domain/mappers/VertexMapper.kt

@@ -101,6 +101,7 @@ object VertexMapper {
                 .name(map["NAME"] as String)
                 .fullname(map["FULL_NAME"] as String)
                 .order(map["ORDER"] as Int)
+                .isexternal(map["IS_EXTERNAL"] as Boolean)
                 .id(map["id"] as Long)
             TypeParameter.Label() -> NewTypeParameterBuilder()
                 .name(map["NAME"] as String)
@@ -332,7 +333,6 @@ object VertexMapper {
                 "POLICY_DIRECTORIES" -> Optional.empty()
                 "INHERITS_FROM_TYPE_FULL_NAME" -> Optional.empty()
                 "OVERLAYS" -> Optional.empty()
-                "IS_EXTERNAL" -> Optional.empty()
                 else -> Optional.of(it.key)
             }
             if (key.isPresent) attributes[key.get()] = it.value

plume/src/main/kotlin/io/github/plume/oss/drivers/GremlinDriver.kt

@@ -166,13 +166,24 @@ abstract class GremlinDriver : IDriver {
     }
 
     override fun getProgramStructure(): overflowdb.Graph {
-        val programStructureSubGraph = g.V().hasLabel(File.Label())
+        val sg = g.V().hasLabel(File.Label())
             .repeat(un.outE(AST).inV()).emit()
             .inE()
             .subgraph("sg")
             .cap<Graph>("sg")
             .next()
-        return gremlinToPlume(programStructureSubGraph.traversal())
+        // Transfer type decl vertices to the result, this needs to be done with the tokens step to get all properties
+        // from the remote server
+        g.V().hasLabel(TypeDecl.Label()).valueMap<String>()
+            .with(WithOptions.tokens)
+            .by(un.unfold<Any>())
+            .toStream()
+            .map { Pair(sg.addVertex(T.label, it[T.label], T.id, it[T.id]), it as Map<*, *>) }
+            .forEach { (v, map) ->
+                map.filter { it.key != T.id && it.key != T.label }
+                    .forEach { (key, value) -> v.property(key.toString(), value) }
+            }
+        return gremlinToPlume(sg.traversal())
     }
 
     override fun getNeighbours(v: NewNodeBuilder): overflowdb.Graph {

plume/src/main/kotlin/io/github/plume/oss/drivers/IDriver.kt

@@ -74,7 +74,7 @@ interface IDriver : AutoCloseable {
     fun getMethod(fullName: String, signature: String, includeBody: Boolean = false): Graph
 
     /**
-     * Obtains all program structure related vertices.
+     * Obtains all program structure related vertices. These are NAMESPACE_BLOCK, FILE, and TYPE_DECL vertices.
      *
      * @return The [Graph] containing the program structure related sub-graphs.
      */

plume/src/main/kotlin/io/github/plume/oss/drivers/Neo4jDriver.kt

@@ -5,8 +5,7 @@ import io.github.plume.oss.domain.mappers.VertexMapper
 import io.github.plume.oss.domain.mappers.VertexMapper.checkSchemaConstraints
 import io.github.plume.oss.domain.mappers.VertexMapper.mapToVertex
 import io.shiftleft.codepropertygraph.generated.EdgeTypes.AST
-import io.shiftleft.codepropertygraph.generated.NodeTypes.FILE
-import io.shiftleft.codepropertygraph.generated.NodeTypes.METHOD
+import io.shiftleft.codepropertygraph.generated.NodeTypes.*
 import io.shiftleft.codepropertygraph.generated.nodes.NewNodeBuilder
 import org.apache.logging.log4j.LogManager
 import org.neo4j.driver.AuthTokens
@@ -320,6 +319,17 @@ class Neo4jDriver : IDriver {
                 ).list().map { it["x"] }
             }
             neo4jToOverflowGraph(result, graph)
+            val typeDecl = session.writeTransaction { tx ->
+                tx.run(
+                    """
+                    MATCH (n:$TYPE_DECL)
+                    RETURN n
+                    """.trimIndent()
+                ).list()
+            }
+            typeDecl.map { it["n"].asNode() }
+                .map { mapToVertex(it.asMap() + mapOf("id" to it.id())) }
+                .forEach { addNodeToGraph(graph, it) }
         }
         return graph
     }

plume/src/main/kotlin/io/github/plume/oss/drivers/OverflowDbDriver.kt

@@ -134,7 +134,12 @@ class OverflowDbDriver : IDriver {
     }
 
     override fun getProgramStructure(): Graph {
-        return deepCopyGraph(Traversals.getProgramStructure(graph))
+        val g = deepCopyGraph(Traversals.getProgramStructure(graph))
+        Traversals.getTypeDecls(graph).forEach { t ->
+            val node = g.addNode(t.id(), t.label())
+            t.propertyMap().forEach { (key, value) -> node.setProperty(key, value) }
+        }
+        return g
     }
 
     override fun getNeighbours(v: NewNodeBuilder): Graph {

plume/src/main/kotlin/io/github/plume/oss/graphio/GraphMLWriter.kt

@@ -54,6 +54,7 @@ object GraphMLWriter {
                 when (u) {
                     is String -> keySet[t] = "string"
                     is Int -> keySet[t] = "int"
+                    is Boolean -> keySet[t] = "boolean"
                     is `$colon$colon`<*> -> keySet[t] = "string"
                     is `Nil$` -> keySet[t] = "string"
                 }

plume/src/main/kotlin/io/github/plume/oss/graphio/GraphSONWriter.kt

@@ -99,6 +99,7 @@ object GraphSONWriter {
                 is String -> sb.append(",\"value\":\"$v\"")
                 is Int -> sb.append(",\"value\":{\"@type\":\"g:Int32\",\"@value\":$v}")
                 is Long -> sb.append(",\"value\":{\"@type\":\"g:Int64\",\"@value\":$v}")
+                is Boolean -> sb.append(",\"value\":$v")
                 else -> println("Unsupported type $v ${v.javaClass}")
             }
             sb.append("}]")

plume/src/main/kotlin/io/github/plume/oss/util/SootToPlumeUtil.kt

@@ -21,6 +21,7 @@ import io.github.plume.oss.domain.mappers.VertexMapper.mapToVertex
 import io.github.plume.oss.drivers.IDriver
 import io.github.plume.oss.util.SootParserUtil.determineEvaluationStrategy
 import io.shiftleft.codepropertygraph.generated.EdgeTypes.*
+import io.shiftleft.codepropertygraph.generated.NodeTypes.UNKNOWN
 import io.shiftleft.codepropertygraph.generated.nodes.*
 import scala.Option
 import scala.jdk.CollectionConverters
@@ -44,7 +45,7 @@ object SootToPlumeUtil {
      *
      * @param field The [SootField] from which the class member information is constructed from.
      */
-    private fun projectMember(field: SootField, childIdx: Int): NewMemberBuilder =
+    fun projectMember(field: SootField, childIdx: Int): NewMemberBuilder =
         NewMemberBuilder()
             .name(field.name)
             .code(field.declaration)
@@ -153,8 +154,16 @@ object SootToPlumeUtil {
             (Extractor.getSootAssociation(mtd)?.size ?: 0) + 1
         )
         // Create program structure
-        buildClassStructure(mtd.declaringClass, driver)
-        buildTypeDeclaration(mtd.declaringClass, driver)
+        val cls = mtd.declaringClass
+        buildClassStructure(cls, driver)
+        val typeDecl = buildTypeDeclaration(cls.type)
+        addSootToPlumeAssociation(cls, typeDecl)
+        cls.fields.forEachIndexed { i, field ->
+            projectMember(field, i + 1).let { memberVertex ->
+                driver.addEdge(typeDecl, memberVertex, AST)
+                addSootToPlumeAssociation(field, memberVertex)
+            }
+        }
         connectMethodToTypeDecls(mtd, driver)
         return mtdVertex
     }
@@ -245,10 +254,10 @@ object SootToPlumeUtil {
      * @param cls the soot class
      * @return the filename in string form
      * */
-    fun sootClassToFileName(cls : SootClass) : String {
+    fun sootClassToFileName(cls: SootClass): String {
         val packageName = cls.packageName
         return if (packageName != null) {
-            "/" + cls.name.replace(".", "/")  + ".class"
+            "/" + cls.name.replace(".", "/") + ".class"
         } else {
             io.shiftleft.semanticcpg.language.types.structure.File.UNKNOWN()
         }
@@ -260,7 +269,11 @@ object SootToPlumeUtil {
      * @param namespaceList A list of package names.
      * @return The final [NewNamespaceBlockBuilder] in the chain (the one associated with the file).
      */
-    private fun populateNamespaceChain(namespaceList: Array<String>, filename : String, driver: IDriver): NewNamespaceBlockBuilder {
+    private fun populateNamespaceChain(
+        namespaceList: Array<String>,
+        filename: String,
+        driver: IDriver
+    ): NewNamespaceBlockBuilder {
         var prevNamespaceBlock: NewNamespaceBlockBuilder
         var currNamespaceBlock: NewNamespaceBlockBuilder? = null
         driver.getProgramStructure().use { programStructure ->
@@ -304,30 +317,36 @@ object SootToPlumeUtil {
     }
 
     /**
-     * Given a class will construct a type declaration with members.
+     * Given a type will construct a type declaration with members.
      *
-     * @param cls The [SootClass] to create the declaration from.
-     * @param driver The driver to construct the type declaration to.
+     * @param type The [soot.Type] to create the declaration from.
+     * @param isExternal Whether the type is part of the application or it is external.
      * @return The [NewTypeDecl] representing this newly created vertex.
      */
-    fun buildTypeDeclaration(cls: SootClass, driver: IDriver): NewTypeDeclBuilder =
-        NewTypeDeclBuilder()
-            .name(cls.shortName)
-            .fullname(cls.name)
-            .filename(sootClassToFileName(cls))
-            .astparentfullname(cls.getPackageName())
+    fun buildTypeDeclaration(type: soot.Type, isExternal: Boolean = true): NewTypeDeclBuilder {
+        val filename = if (isExternal) {
+            "<${UNKNOWN.toLowerCase()}>"
+        } else {
+            if (type.toQuotedString().contains('.')) "/${
+                type.toQuotedString().replace(".", "/").removeSuffix("[]")
+            }.class"
+            else type.toQuotedString()
+        }
+        val parentType = if (type.toQuotedString().contains('.')) type.toQuotedString().substringBeforeLast(".")
+        else type.toQuotedString()
+        val shortName = if (type.toQuotedString().contains('.')) type.toQuotedString().substringAfterLast('.')
+        else type.toQuotedString()
+
+        return NewTypeDeclBuilder()
+            .name(shortName)
+            .fullname(type.toQuotedString())
+            .filename(filename)
+            .astparentfullname(parentType)
             .astparenttype("NAMESPACE_BLOCK")
-            .order(0)
-            .apply {
-                // Attach fields to the TypeDecl
-                cls.fields.forEachIndexed { i, field ->
-                    projectMember(field, i + 1).let { memberVertex ->
-                        driver.addEdge(this, memberVertex, AST)
-                        addSootToPlumeAssociation(field, memberVertex)
-                    }
-                }
-                addSootToPlumeAssociation(cls, this)
-            }
+            .order(if (isExternal) -1 else 1)
+            .isexternal(isExternal)
+            .apply { addSootToPlumeAssociation(type, this) }
+    }
 
     /**
      * Connects the given method's [BriefUnitGraph] to its type declaration and source file (if present).

plume/src/test/kotlin/io/github/plume/oss/TestDomainResources.kt

@@ -24,6 +24,7 @@ class TestDomainResources {
         val EVAL_2 = BY_SHARING
         val MOD_1 = ABSTRACT
         val MOD_2 = CONSTRUCTOR
+        val BOOL_1 = false
 
         val vertices = listOf<NewNodeBuilder>(
             NewArrayInitializerBuilder().order(INT_1),
@@ -63,7 +64,7 @@ class TestDomainResources {
             NewReturnBuilder().order(INT_1).argumentindex(INT_1).code(STRING_1).linenumber(Option.apply(INT_1))
                 .columnnumber(Option.apply(INT_1)),
             NewTypeArgumentBuilder().order(INT_1),
-            NewTypeDeclBuilder().name(STRING_1).fullname(STRING_1).order(INT_1),
+            NewTypeDeclBuilder().name(STRING_1).fullname(STRING_1).order(INT_1).isexternal(BOOL_1),
             NewTypeParameterBuilder().name(STRING_1).order(INT_1),
             NewTypeRefBuilder().typefullname(STRING_1)
                 .dynamictypehintfullname(createScalaList(STRING_1))

plume/src/test/kotlin/io/github/plume/oss/drivers/JanusGraphDriverIntTest.kt

@@ -31,6 +31,7 @@ import io.github.plume.oss.TestDomainResources.Companion.typeParameterVertex
 import io.github.plume.oss.TestDomainResources.Companion.typeRefVertex
 import io.github.plume.oss.TestDomainResources.Companion.unknownVertex
 import io.github.plume.oss.domain.exceptions.PlumeSchemaViolationException
+import io.github.plume.oss.graphio.GraphMLWriter
 import io.github.plume.oss.util.SootToPlumeUtil
 import io.shiftleft.codepropertygraph.generated.EdgeTypes.*
 import io.shiftleft.codepropertygraph.generated.nodes.*
@@ -39,6 +40,7 @@ import org.junit.jupiter.api.*
 import org.junit.jupiter.api.Assertions.*
 import overflowdb.Graph
 import scala.Option
+import java.io.FileWriter
 import kotlin.properties.Delegates
 
 class JanusGraphDriverIntTest {
@@ -504,15 +506,16 @@ class JanusGraphDriverIntTest {
             g = driver.getProgramStructure()
             val ns = g.nodes().asSequence().toList()
             val es = g.edges().asSequence().toList()
-            assertEquals(3, ns.size)
+            assertEquals(4, ns.size)
             assertEquals(2, es.size)
 
             val file = g.V(fileVertex.id()).next()
             val ns1 = g.V(namespaceBlockVertex1.id()).next()
-            // Assert no program structure vertices part of the method body
+            // Assert program structure vertices are present
             assertTrue(ns.any { it.id() == namespaceBlockVertex2.id() })
             assertTrue(ns.any { it.id() == namespaceBlockVertex1.id() })
             assertTrue(ns.any { it.id() == fileVertex.id() })
+            assertTrue(ns.any { it.id() == typeDeclVertex.id() })
             // Check that vertices are connected by AST edges
             assertTrue(file.out(AST).asSequence().any { it.id() == namespaceBlockVertex1.id() })
             assertTrue(ns1.out(AST).asSequence().any { it.id() == namespaceBlockVertex2.id() })

plume/src/test/kotlin/io/github/plume/oss/drivers/Neo4jDriverIntTest.kt

@@ -508,15 +508,16 @@ class Neo4jDriverIntTest {
             g = driver.getProgramStructure()
             val ns = g.nodes().asSequence().toList()
             val es = g.edges().asSequence().toList()
-            assertEquals(3, ns.size)
+            assertEquals(4, ns.size)
             assertEquals(2, es.size)
 
             val file = g.V(fileVertex.id()).next()
             val ns1 = g.V(namespaceBlockVertex1.id()).next()
-            // Assert no program structure vertices part of the method body
+            // Assert program structure vertices are present
             assertTrue(ns.any { it.id() == namespaceBlockVertex2.id() })
             assertTrue(ns.any { it.id() == namespaceBlockVertex1.id() })
             assertTrue(ns.any { it.id() == fileVertex.id() })
+            assertTrue(ns.any { it.id() == typeDeclVertex.id() })
             // Check that vertices are connected by AST edges
             assertTrue(file.out(AST).asSequence().any { it.id() == namespaceBlockVertex1.id() })
             assertTrue(ns1.out(AST).asSequence().any { it.id() == namespaceBlockVertex2.id() })