diff --git a/src/main/scala/ir/Parent.scala b/src/main/scala/ir/Parent.scala index 755ac33dd..9e679f14b 100644 --- a/src/main/scala/ir/Parent.scala +++ b/src/main/scala/ir/Parent.scala @@ -1,6 +1,8 @@ package ir import util.Logger +import scala.io.AnsiColor + trait HasParent[T]: /* If a node is reachable from the IL then it *must* have a parent defined. This will only be null until @@ -9,12 +11,17 @@ trait HasParent[T]: All IL structures must set the parent of the child to itself, when a child is added to itself. */ private var _parent: Option[T] = None - private var last_parent: Option[T] = None + private var last_parent: Option[(T, String, String, Int)] = None def parent: T = { - if (!hasParent) { - Logger.error(s"Trying to get the parent of a node that is detached from the progam: $this. Node was last attached to: $last_parent") + _parent match { + case Some(p) => p + case None => { + val msg = s"Trying to get the parent of a node that is detached from the progam: $this. " + + s"Node was last attached to: ${last_parent.map(_._1)}, detached at ${last_parent.map(x => x._2 + x._3 + "@" + x._4)} " + Logger.error(msg) + throw new RuntimeException(msg) + } } - _parent.get } def hasParent: Boolean = _parent.isDefined @@ -45,11 +52,13 @@ trait HasParent[T]: * Remove this element's parent and update any IL control-flow links implied by this relation. * Is idempotent. */ - final def deParent(): Unit = if _parent.isDefined then { - unlinkParent() - if _parent.isDefined then (last_parent = _parent) - _parent = None + final def deParent()(implicit line: sourcecode.Line, file: sourcecode.FileName, name: sourcecode.Name): Unit = { + if _parent.isDefined then { + unlinkParent() + if _parent.isDefined then (last_parent = _parent.map(x => (x, file.value, name.value, line.value))) + _parent = None } + } /** * Set this element's parent and update any IL control-flow links implied by this relation. diff --git a/src/main/scala/ir/Program.scala b/src/main/scala/ir/Program.scala index 93ac1321c..08328ed0a 100644 --- a/src/main/scala/ir/Program.scala +++ b/src/main/scala/ir/Program.scala @@ -330,37 +330,12 @@ class Block private ( def jump: Jump = _jump - def jump_=(j: Call): Unit = { - if (j ne _jump) { - _jump.deParent() - _jump match { - case c: IndirectCall => c.returnTarget.filter(b=>b.kind.isInstanceOf[AfterCall]).map(b => parent.removeBlocks(b)) - case c: DirectCall => c.returnTarget.filter(b=>b.kind.isInstanceOf[AfterCall]).map(b => parent.removeBlocks(b)) - } - _jump = j - _jump.parent = this - j match { - case c: IndirectCall => c.returnTarget = Some(parent.addBlocks(Block.afterCall(c))) - case c: DirectCall => c.returnTarget = Some(parent.addBlocks(Block.afterCall(c))) - } - } - } def jump_=(j: Jump): Unit = { if (j ne _jump) { _jump.deParent() - //_jump match { - // case c: IndirectCall => c.returnTarget.filter(b=>b.kind.isInstanceOf[AfterCall]).map(b => parent.removeBlocks(b)) - // case c: DirectCall => c.returnTarget.filter(b=>b.kind.isInstanceOf[AfterCall]).map(b => parent.removeBlocks(b)) - // case _ => () - //} _jump = j _jump.parent = this - //j match { - // case c: IndirectCall => c.returnTarget = Some(parent.addBlocks(Block.afterCall(c))) - // case c: DirectCall => c.returnTarget = Some(parent.addBlocks(Block.afterCall(c))) - // case _ => () - //} } } diff --git a/src/main/scala/util/RunUtils.scala b/src/main/scala/util/RunUtils.scala index f5dc45f72..45b9717f2 100644 --- a/src/main/scala/util/RunUtils.scala +++ b/src/main/scala/util/RunUtils.scala @@ -222,13 +222,12 @@ object RunUtils { config.analysisResultsPath.foreach(s => writeToFile(printAnalysisResults(IRProgram, cfg, vsaResult), s"${s}_vsa$iteration.txt")) Logger.info("[!] Resolving CFG") - val (newIR, modified): (Program, Boolean) = resolveCFG(cfg, vsaResult, IRProgram) + val (newIR, modified): (Program, Boolean) = resolveCFG(cfg, vsaResult, IRProgram) - unparented = IRProgram.collect { + unparented = newIR.collect { case c: Command if !c.hasParent => c case b: Block if !b.hasParent => b } - assert(unparented.isEmpty) if (modified) { Logger.info(s"[!] Analysing again (iter $iteration)") @@ -444,16 +443,25 @@ object RunUtils { if (targets.size == 1) { modified = true - val newCall = DirectCall(targets.head, indirectCall.returnTarget, indirectCall.label) + + indirectCall.parent.parent.removeBlocks(indirectCall.returnTarget) + val newCall = DirectCall(targets.head, None, indirectCall.label) + newCall.parent = indirectCall.parent + newCall.returnTarget = Some(indirectCall.parent.parent.addBlocks(Block.afterCall(newCall))) + block.replaceJump(newCall) } else if (targets.size > 1) { modified = true val procedure = c.parent.data val newBlocks = ArrayBuffer[Block]() + indirectCall.parent.parent.removeBlocks(indirectCall.returnTarget) for (t <- targets) { val assume = Assume(BinaryExpr(BVEQ, indirectCall.target, BitVecLiteral(t.address.get, 64))) val newLabel: String = block.label + t.name - val directCall = DirectCall(t, indirectCall.returnTarget) + val directCall = DirectCall(t, None) + directCall.parent = indirectCall.parent + directCall.returnTarget = Some(indirectCall.parent.parent.addBlocks(Block.afterCall(directCall))) + newBlocks.append(Block.regular(newLabel, None, ArrayBuffer(assume), directCall)) } procedure.addBlocks(newBlocks) @@ -480,7 +488,7 @@ object RunUtils { case globalAddress: GlobalAddress => if (nameExists(globalAddress.name)) { functionNames += globalAddress - Logger.info(s"RESOLVED: Call to Global address ${globalAddress.name} resolved.") + Logger.info(s"RESOLVED: Call to Global address ${globalAddress.name} rt statuesolved.") } else { addFakeProcedure(globalAddress.name) functionNames += globalAddress