From 055ca98fb2a98f81bbe9e419eeea840d256490e9 Mon Sep 17 00:00:00 2001
From: van Hauser <vh@thc.org>
Date: Wed, 8 Apr 2020 12:52:14 +0200
Subject: add missing piece for CTX coverage

---
 llvm_mode/afl-llvm-pass.so.cc | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

(limited to 'llvm_mode/afl-llvm-pass.so.cc')

diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc
index 31d00fec..058ab71f 100644
--- a/llvm_mode/afl-llvm-pass.so.cc
+++ b/llvm_mode/afl-llvm-pass.so.cc
@@ -682,6 +682,22 @@ bool AFLCoverage::runOnModule(Module &M) {
 
       }
 
+      // in CTX mode we have to restore the original context for the caller -
+      // she might be calling other functions which need the correct CTX
+      if (ctx_str) {
+
+        Instruction *Inst = BB.getTerminator();
+        if (isa<ReturnInst>(Inst) || isa<ResumeInst>(Inst)) {
+
+          IRBuilder<> Post_IRB(Inst);
+          StoreInst * RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext);
+          RestoreCtx->setMetadata(M.getMDKindID("nosanitize"),
+                                  MDNode::get(C, None));
+
+        }
+
+      }
+
       inst_blocks++;
 
     }
-- 
cgit 1.4.1