From 31094588f3cd3429ecf2552955c780304968cf9d Mon Sep 17 00:00:00 2001
From: Brendan Edmonds <bed88012@gmail.com>
Date: Thu, 5 Oct 2023 14:26:46 +1000
Subject: [PATCH] BASIL-19: Fixed issue with external functios

---
 src/main/scala/analysis/NonReturningFunctions.scala | 10 +++++++---
 src/main/scala/ir/Program.scala                     |  3 ++-
 src/main/scala/translating/BAPToIR.scala            |  5 +++--
 src/main/scala/util/RunUtils.scala                  |  6 +++---
 4 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/main/scala/analysis/NonReturningFunctions.scala b/src/main/scala/analysis/NonReturningFunctions.scala
index 1cada46da..b06670a1d 100644
--- a/src/main/scala/analysis/NonReturningFunctions.scala
+++ b/src/main/scala/analysis/NonReturningFunctions.scala
@@ -2,9 +2,11 @@ package analysis
 
 import bap.{BAPJump, BAPSubroutine}
 import ir.{Block, DirectCall, GoTo, IndirectCall, Jump, Procedure, Statement}
+import specification.ExternalFunction
 
 import scala.collection.mutable.{Map, Queue}
 import collection.parallel.CollectionConverters.seqIsParallelizable
+import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.parallel.CollectionConverters.*
 
@@ -15,6 +17,8 @@ class NonReturningFunctions {
     val mapJumpsToBlocks: Map[String, ArrayBuffer[(Jump, Block)]] = Map()
     val mapBlocksToProcedure: Map[String, (Procedure, Integer)] = Map()
 
+
+
     def isEndlessLoop(proc: Procedure, goTo: GoTo, index: Integer): Boolean = {
 
       if (goTo.condition.isEmpty && mapBlocksToProcedure.contains(goTo.target.label) && mapBlocksToProcedure(goTo.target.label)._2 < index) {
@@ -45,6 +49,7 @@ class NonReturningFunctions {
     for (proc <- procedures) {
       var numberOfReturns = 0
       for ((block, index) <- proc.blocks.zipWithIndex) {
+
         mapBlocksToProcedure.addOne(block.label, (proc, index))
         for (jump <- block.jumps) {
 
@@ -59,7 +64,7 @@ class NonReturningFunctions {
             case goTo: GoTo =>
               mapJumpsToBlocks.put(goTo.target.label, mapJumpsToBlocks.getOrElse(goTo.target.label, ArrayBuffer()).addOne((goTo, block)))
               if (proc.blocks.length > index && isEndlessLoop(proc, goTo, index)) {
-                blocksToRemove.enqueue(proc.blocks(index+1).label)
+                blocksToRemove.enqueue(proc.blocks(index + 1).label)
               }
             case _ =>
           }
@@ -72,8 +77,7 @@ class NonReturningFunctions {
     while (blocksDeleted) {
       blocksDeleted = false
       for (proc <- procedures) {
-
-        if (proc.calculateReturnCount() == 0) {
+        if (!proc.externalFunction && proc.calculateReturnCount() == 0) {
           mapJumpsToBlocks.get(proc.name) match {
             case Some(v) => for (block <- v) {
               val (_, containingBlock) = block
diff --git a/src/main/scala/ir/Program.scala b/src/main/scala/ir/Program.scala
index 81a0a779b..70dae7e9c 100644
--- a/src/main/scala/ir/Program.scala
+++ b/src/main/scala/ir/Program.scala
@@ -106,7 +106,8 @@ class Procedure(
     var address: Option[Int],
     var blocks: ArrayBuffer[Block],
     var in: ArrayBuffer[Parameter],
-    var out: ArrayBuffer[Parameter]
+    var out: ArrayBuffer[Parameter],
+    val externalFunction: Boolean
 ) {
 
   def calls: Set[Procedure] = blocks.flatMap(_.calls).toSet
diff --git a/src/main/scala/translating/BAPToIR.scala b/src/main/scala/translating/BAPToIR.scala
index 0f43a545c..c976c44f2 100644
--- a/src/main/scala/translating/BAPToIR.scala
+++ b/src/main/scala/translating/BAPToIR.scala
@@ -9,7 +9,7 @@ import scala.collection.mutable
 import scala.collection.mutable.Map
 import scala.collection.mutable.ArrayBuffer
 
-class BAPToIR(var program: BAPProgram, mainAddress: Int) {
+class BAPToIR(var program: BAPProgram, mainAddress: Int, externalFunctions: Set[ExternalFunction]) {
 
   private val nameToProcedure: mutable.Map[String, Procedure] = mutable.Map()
   private val labelToBlock: mutable.Map[String, Block] = mutable.Map()
@@ -17,6 +17,7 @@ class BAPToIR(var program: BAPProgram, mainAddress: Int) {
   def translate: Program = {
     var mainProcedure: Option[Procedure] = None
     val procedures: ArrayBuffer[Procedure] = ArrayBuffer()
+    val externalFunctionNames = externalFunctions.map(func => func.name)
     for (s <- program.subroutines) {
       val blocks: ArrayBuffer[Block] = ArrayBuffer()
       for (b <- s.blocks) {
@@ -32,7 +33,7 @@ class BAPToIR(var program: BAPProgram, mainAddress: Int) {
       for (p <- s.out) {
         out.append(p.toIR)
       }
-      val procedure = Procedure(s.name, Some(s.address), blocks, in, out)
+      val procedure = Procedure(s.name, Some(s.address), blocks, in, out, externalFunctionNames.contains(s.name))
       if (s.address == mainAddress) {
         mainProcedure = Some(procedure)
       }
diff --git a/src/main/scala/util/RunUtils.scala b/src/main/scala/util/RunUtils.scala
index 113c35d55..5f7b26ae2 100644
--- a/src/main/scala/util/RunUtils.scala
+++ b/src/main/scala/util/RunUtils.scala
@@ -73,9 +73,9 @@ object RunUtils {
 
     val (externalFunctions, globals, globalOffsets, mainAddress) = loadReadELF(readELFFileName)
 
-    val IRTranslator = BAPToIR(bapProgram, mainAddress)
+    val IRTranslator = BAPToIR(bapProgram, mainAddress, externalFunctions)
     var IRProgram = IRTranslator.translate
-    NonReturningFunctions().transform(IRProgram.procedures)
+    NonReturningFunctions().transform(IRProgram.procedures, externalFunctions)
 
     val specification = loadSpecification(specFileName, IRProgram, globals)
 
@@ -316,7 +316,7 @@ object RunUtils {
     }
 
     def addFakeProcedure(name: String): Unit = {
-      IRProgram.procedures += Procedure(name, None, ArrayBuffer(), ArrayBuffer(), ArrayBuffer())
+      IRProgram.procedures += Procedure(name, None, ArrayBuffer(), ArrayBuffer(), ArrayBuffer(), true)
     }
 
     def resolveAddresses(valueSet: Set[Value]): Set[AddressValue] = {