From cb279e238da177de89b57b016e8a70f4bccb81a5 Mon Sep 17 00:00:00 2001
From: David Baker Effendi <davidbakereffendi@hotmail.com>
Date: Thu, 11 Feb 2021 12:57:41 +0200
Subject: [PATCH] :bug: Fixed Invalid Call Edges (#65)

* :sparkles: Improved build dir initialization and clean up

* :bug: Fixed ampersands in xml

* :bug: Fixed non-call vertices being sources of calls

* Added fixes to 0.1.7
---
 CHANGELOG.md                                                | 2 ++
 plume/src/main/kotlin/io/github/plume/oss/Extractor.kt      | 6 +++++-
 .../kotlin/io/github/plume/oss/graph/CallGraphBuilder.kt    | 2 +-
 .../kotlin/io/github/plume/oss/graphio/GraphMLWriter.kt     | 5 ++++-
 .../io/github/plume/oss/extractor/BasicExtractorTest.kt     | 1 -
 5 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 87368fa6..158ce76c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Fixed
 - Module not found bug introduced by improper class cleanup in temp dir.
+- Fixed instances where CallGraphBuilder would connect non-NewCallBuilder source nodes to methods. 
+- Fixed GraphML not escaping ampersands
 
 ## [0.1.6] - 2021-02-10
 
diff --git a/plume/src/main/kotlin/io/github/plume/oss/Extractor.kt b/plume/src/main/kotlin/io/github/plume/oss/Extractor.kt
index 3a82b840..c70f667c 100644
--- a/plume/src/main/kotlin/io/github/plume/oss/Extractor.kt
+++ b/plume/src/main/kotlin/io/github/plume/oss/Extractor.kt
@@ -81,6 +81,10 @@ class Extractor(val driver: IDriver) {
     private lateinit var programStructure: Graph
 
     init {
+        File(COMP_DIR).let { f ->
+            if (f.exists()) f.deleteRecursively()
+            f.deleteOnExit()
+        }
         checkDriverConnection(driver)
         astBuilder = ASTBuilder(driver)
         cfgBuilder = CFGBuilder(driver)
@@ -188,7 +192,7 @@ class Extractor(val driver: IDriver) {
      */
     @Throws(PlumeCompileException::class, NullPointerException::class, IOException::class)
     fun load(f: File) {
-        if (!File(COMP_DIR).exists()) File(COMP_DIR).mkdirs()
+        File(COMP_DIR).let { c -> if (!c.exists()) c.mkdirs() }
         if (!f.exists()) {
             throw NullPointerException("File '${f.name}' does not exist!")
         } else if (f.isDirectory) {
diff --git a/plume/src/main/kotlin/io/github/plume/oss/graph/CallGraphBuilder.kt b/plume/src/main/kotlin/io/github/plume/oss/graph/CallGraphBuilder.kt
index f0cd62cc..4f66c372 100644
--- a/plume/src/main/kotlin/io/github/plume/oss/graph/CallGraphBuilder.kt
+++ b/plume/src/main/kotlin/io/github/plume/oss/graph/CallGraphBuilder.kt
@@ -59,7 +59,7 @@ class CallGraphBuilder(private val driver: IDriver) : IGraphBuilder {
             // If Soot points to the assignment as the call source then this is most likely from the rightOp. Let's
             // hope this is not the source of a bug
             val srcUnit = if (unit is AssignStmt) unit.rightOp else unit
-            getSootAssociation(srcUnit)?.firstOrNull()?.let { srcPlumeVertex ->
+            getSootAssociation(srcUnit)?.filterIsInstance<NewCallBuilder>()?.firstOrNull()?.let { srcPlumeVertex ->
                 val tgtPlumeVertex = getSootAssociation(e.tgt.method())?.firstOrNull()
                     ?: constructPhantom(e.tgt.method(), driver)
                 runCatching {
diff --git a/plume/src/main/kotlin/io/github/plume/oss/graphio/GraphMLWriter.kt b/plume/src/main/kotlin/io/github/plume/oss/graphio/GraphMLWriter.kt
index 5ad97f8f..132e0e9b 100644
--- a/plume/src/main/kotlin/io/github/plume/oss/graphio/GraphMLWriter.kt
+++ b/plume/src/main/kotlin/io/github/plume/oss/graphio/GraphMLWriter.kt
@@ -84,7 +84,10 @@ object GraphMLWriter {
 
     private fun escape(o: Any) =
         when (o) {
-            is String -> o.replace("<", "&lt;").replace(">", "&gt;")
+            is String -> o
+                .replace("<", "&lt;")
+                .replace(">", "&gt;")
+                .replace("&", "&amp;")
             is `$colon$colon`<*> -> o.head()
             is `Nil$` -> ""
             else -> o.toString()
diff --git a/plume/src/test/kotlin/io/github/plume/oss/extractor/BasicExtractorTest.kt b/plume/src/test/kotlin/io/github/plume/oss/extractor/BasicExtractorTest.kt
index 68edeac1..fb4cfe6d 100644
--- a/plume/src/test/kotlin/io/github/plume/oss/extractor/BasicExtractorTest.kt
+++ b/plume/src/test/kotlin/io/github/plume/oss/extractor/BasicExtractorTest.kt
@@ -110,7 +110,6 @@ class BasicExtractorTest {
         extractor.load(validJarFile)
         extractor.project()
         g = driver.getWholeGraph()
-        GraphMLWriter.write(g, FileWriter("/tmp/plume/x.xml"))
         g = driver.getProgramStructure()
         val ns = g.nodes().asSequence().toList()
         ns.filterIsInstance<ODBFile>().let { fileList ->