From 7e2c123b71cb57203ca586fbbece5613012435f1 Mon Sep 17 00:00:00 2001 From: Alistair Michael Date: Wed, 31 Jan 2024 15:57:13 +1000 Subject: [PATCH] aftercall rename --- src/main/scala/ir/IRCursor.scala | 8 ++--- src/main/scala/ir/Parent.scala | 13 +++++++-- src/main/scala/ir/Program.scala | 37 ++++++++++++++++++------ src/main/scala/ir/Visitor.scala | 4 +-- src/main/scala/translating/BAPToIR.scala | 1 + src/main/scala/util/RunUtils.scala | 18 +++++++++++- 6 files changed, 63 insertions(+), 18 deletions(-) diff --git a/src/main/scala/ir/IRCursor.scala b/src/main/scala/ir/IRCursor.scala index b92edbd9c..8d516c752 100644 --- a/src/main/scala/ir/IRCursor.scala +++ b/src/main/scala/ir/IRCursor.scala @@ -72,7 +72,7 @@ trait IntraProcIRCursor extends IRWalk[CFGPosition, CFGPosition] { case j: Jump => Set(j.parent.statements.lastOption.getOrElse(j.parent)) case b: Block => b.kind match { - case CallReturn(from) => Set(from) + case AfterCall(from) => Set(from) case Entry(proc) => Set(proc) case _ => b.incomingJumps.asInstanceOf[Set[CFGPosition]] } @@ -100,7 +100,7 @@ trait IntraProcBlockIRCursor extends IRWalk[CFGPosition, Block] { case b: Block => b.kind match { case Entry(_) => Set.empty - case CallReturn(from) => Set(from.parent) + case AfterCall(from) => Set(from.parent) case _ => b.incomingJumps.map(_.parent) } case j: Command => pred(j.parent) @@ -139,7 +139,7 @@ trait InterProcIRCursor extends IRWalk[CFGPosition, CFGPosition] { case c: Procedure => c.incomingCalls().toSet.asInstanceOf[Set[CFGPosition]] case b: Block => b.kind match { - case CallReturn(DirectCall(target, returnTarget, _)) => target.returnBlock.map(_.jump).toSet + case AfterCall(DirectCall(target, returnTarget, _)) => target.returnBlock.map(_.jump).toSet case _ => Set() } case _ => IntraProcIRCursor.pred(pos) @@ -160,7 +160,7 @@ trait InterProcBlockIRCursor extends IRWalk[CFGPosition, Block] { pos match { case b: Block => b.kind match { - case CallReturn(from) => from.parent.parent.returnBlock.toSet + case AfterCall(from) => from.parent.parent.returnBlock.toSet case Entry(of) => of.incomingCalls().map(_.parent).toSet case _ => b.prevBlocks.toSet } diff --git a/src/main/scala/ir/Parent.scala b/src/main/scala/ir/Parent.scala index 46557f5c9..755ac33dd 100644 --- a/src/main/scala/ir/Parent.scala +++ b/src/main/scala/ir/Parent.scala @@ -1,4 +1,5 @@ -package ir +package ir +import util.Logger trait HasParent[T]: /* @@ -8,7 +9,14 @@ 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 - def parent: T = _parent.get + private var last_parent: Option[T] = 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.get + } + def hasParent: Boolean = _parent.isDefined def parent_=(value: T): Unit = setParent(value) @@ -39,6 +47,7 @@ trait HasParent[T]: */ final def deParent(): Unit = if _parent.isDefined then { unlinkParent() + if _parent.isDefined then (last_parent = _parent) _parent = None } diff --git a/src/main/scala/ir/Program.scala b/src/main/scala/ir/Program.scala index c4f51e42f..376c7353e 100644 --- a/src/main/scala/ir/Program.scala +++ b/src/main/scala/ir/Program.scala @@ -200,18 +200,23 @@ class Procedure private ( def returnBlock: Option[Block] = _returnBlock def returnBlock_=(value: Block): Unit = { - removeBlocks(_returnBlock) - _returnBlock = Some(addBlocks(value)) + if (!returnBlock.contains(value)) { + removeBlocks(_returnBlock) + _returnBlock = Some(addBlocks(value)) + } } def entryBlock: Option[Block] = _entryBlock def entryBlock_=(value: Block): Unit = { - removeBlocks(_entryBlock) - _entryBlock = Some(addBlocks(value)) + if (!entryBlock.contains(value)) { + removeBlocks(_entryBlock) + _entryBlock = Some(addBlocks(value)) + } } def addBlocks(block: Block): Block = { + block.parent = this if (!_blocks.contains(block)) { block.parent = this _blocks.add(block) @@ -296,7 +301,7 @@ sealed trait BlockKind case class Regular() extends BlockKind /* Block is the fallthrough / return target of a call. */ -case class CallReturn(from: Call) extends BlockKind +case class AfterCall(from: Call) extends BlockKind /* Block is the single return point for a procedure */ case class Return(from: Procedure) extends BlockKind @@ -327,13 +332,24 @@ class Block private ( 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 _ => () + //} } } def replaceJump(j: Jump) = { jump = j + assert(jump.parent eq this) this } @@ -439,12 +455,15 @@ object Block: new Block(Regular(), label, address, IntrusiveList.empty, GoTo(Seq(), Some(label + "_unknown"))) } - def callReturn(from: Call) : Block = { - val jump = from match + def afterCall(from: Call) : Block = { + val jump = from match case d: DirectCall => GoTo(d.returnTarget.toSet) case c: IndirectCall => GoTo(c.returnTarget.toSet) - - new Block(CallReturn(from), from.parent.label + "_basil_callreturn", None, Seq(), jump) + + jump.parent = from.parent + val ac = Block(AfterCall(from), from.parent.label + "_basil_aftercall", None, Seq(), jump) + ac.parent = from.parent.parent + ac } def procedureReturn(from: Procedure): Block = { diff --git a/src/main/scala/ir/Visitor.scala b/src/main/scala/ir/Visitor.scala index 06b78ccf6..618117f28 100644 --- a/src/main/scala/ir/Visitor.scala +++ b/src/main/scala/ir/Visitor.scala @@ -428,7 +428,7 @@ class AddCallReturnBlocks extends Visitor { * Add a dummy call return block with no statements after each call. */ override def visitDirectCall(node: DirectCall): Jump = { - val b = node.parent.parent.addBlocks(Block.callReturn(node)) + val b = node.parent.parent.addBlocks(Block.afterCall(node)) node.returnTarget = Some(b) node } @@ -436,7 +436,7 @@ class AddCallReturnBlocks extends Visitor { override def visitIndirectCall(node: IndirectCall): Jump = { // skip return nodes if !(node.parent.kind.isInstanceOf[Return]) then { - val b = node.parent.parent.addBlocks(Block.callReturn(node)) + val b = node.parent.parent.addBlocks(Block.afterCall(node)) node.returnTarget = Some(b) } node diff --git a/src/main/scala/translating/BAPToIR.scala b/src/main/scala/translating/BAPToIR.scala index a1e674744..e1de424ca 100644 --- a/src/main/scala/translating/BAPToIR.scala +++ b/src/main/scala/translating/BAPToIR.scala @@ -48,6 +48,7 @@ class BAPToIR(var program: BAPProgram, mainAddress: Int) { val (jump, newBlocks) = translate(b.jumps, block) procedure.addBlocks(newBlocks) block.replaceJump(jump) + assert(jump.hasParent) } // Set entry block to the block with the same address as the procedure or the first in sequence diff --git a/src/main/scala/util/RunUtils.scala b/src/main/scala/util/RunUtils.scala index f4b77e509..43ca8b725 100644 --- a/src/main/scala/util/RunUtils.scala +++ b/src/main/scala/util/RunUtils.scala @@ -96,10 +96,10 @@ object RunUtils { val renamer = Renamer(reserved) val returnUnifier = ConvertToSingleProcedureReturn() val callReturner = AddCallReturnBlocks() + IRProgram = callReturner.visitProgram(IRProgram) IRProgram = externalRemover.visitProgram(IRProgram) IRProgram = renamer.visitProgram(IRProgram) IRProgram = returnUnifier.visitProgram(IRProgram) - IRProgram = callReturner.visitProgram(IRProgram) q.loading.dumpIL.foreach(s => writeToFile(serialiseIL(IRProgram), s"$s-before-analysis.il")) @@ -157,12 +157,21 @@ object RunUtils { Logger.info(externalAddresses) Logger.info("Subroutine Addresses:") Logger.info(subroutines) + var unparented = IRProgram.collect { + case c: Command if !c.hasParent => c + case b: Block if !b.hasParent => b + } + assert(unparented.isEmpty) val mergedSubroutines = subroutines ++ externalAddresses val cfg = ProgramCfgFactory().fromIR(IRProgram) + unparented = IRProgram.collect { + case c: IndirectCall if (!c.hasParent) => c + } + assert(unparented.isEmpty) val domain = computeDomain(IntraProcIRCursor, IRProgram.procedures) Logger.info("[!] Running Constant Propagation") @@ -214,6 +223,13 @@ object RunUtils { Logger.info("[!] Resolving CFG") val (newIR, modified): (Program, Boolean) = resolveCFG(cfg, vsaResult, IRProgram) + + unparented = IRProgram.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)") return analyse(newIR, externalFunctions, globals, globalOffsets, config, iteration + 1)