Fix issue #16

JonathanSalwan · Feb 6, 2015 · 9e0acb8 · 9e0acb8
1 parent 706d0e6
commit 9e0acb8
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 8 deletions.
diff --git a/src/includes/SnapshotEngine.h b/src/includes/SnapshotEngine.h
@@ -6,6 +6,7 @@
 
 #include "pin.H"
 #include "SymbolicEngine.h"
+#include "TaintEngine.h"
 
 #define LOCKED    1
 #define UNLOCKED  !LOCKED
@@ -23,8 +24,9 @@ class SnapshotEngine{
     /* Status of the snapshot engine */
     BOOL locked;
 
-    SymbolicEngine *snapshotSymEngine;
-    CONTEXT        pinCtx;
+    SymbolicEngine  *snapshotSymEngine;
+    TaintEngine     *snapshotTaintEngine;
+    CONTEXT         pinCtx;
 
 
   public:
@@ -35,8 +37,8 @@ class SnapshotEngine{
     VOID addModification(UINT64 address, UINT8 byte);
     VOID disableSnapshot();
     VOID resetEngine();
-    VOID restoreSnapshot(SymbolicEngine *currentSymEngine, CONTEXT *ctx);
-    VOID takeSnapshot(const SymbolicEngine &currentSymEngine, CONTEXT *ctx);
+    VOID restoreSnapshot(SymbolicEngine *currentSymEngine, TaintEngine *currentTaintEngine, CONTEXT *ctx);
+    VOID takeSnapshot(const SymbolicEngine &currentSymEngine, const TaintEngine &currentTaintEngine, CONTEXT *ctx);
 
 };
 

diff --git a/src/ir/notImplemented.cpp b/src/ir/notImplemented.cpp
@@ -8,6 +8,12 @@ VOID notImplemented(std::string insDis, ADDRINT insAddr)
   if (_analysisStatus == LOCKED || insAddr > LIB_MAPING_MEMORY)
     return;
 
+  /* Craft the Tritinst without element */
+  Tritinst *inst = new Tritinst(insAddr, insDis);
+
+  /* Add the Tritinst in the trace */
+  trace->addInstruction(inst);
+
   displayTrace(insAddr, insDis, "n/a", !TAINTED);
 }
 
diff --git a/src/snapshotEngine/snapshotEngine.cpp b/src/snapshotEngine/snapshotEngine.cpp
@@ -26,23 +26,26 @@ VOID SnapshotEngine::addModification(UINT64 mem, UINT8 byte)
 
 
 /* Enable the snapshot engine. */
-VOID SnapshotEngine::takeSnapshot(const SymbolicEngine &currentSymEngine, CONTEXT *ctx)
+VOID SnapshotEngine::takeSnapshot(const SymbolicEngine &currentSymEngine, const TaintEngine &currentTaintEngine, CONTEXT *ctx)
 {
   /* 1 - Unlock the engine */
   this->locked = UNLOCKED;
 
   /* 2 - Save current symbolic engine state */
   this->snapshotSymEngine = new SymbolicEngine(currentSymEngine);
 
-  /* 3 - Save Pin registers context */
+  /* 3 - Save current taint engine state */
+  this->snapshotTaintEngine = new TaintEngine(currentTaintEngine);
+
+  /* 4 - Save Pin registers context */
   PIN_SaveContext(ctx, &this->pinCtx);
 
   std::cout << "[snapshot]" << std::endl;
 }
 
 
 /* Restore the snapshot. */
-VOID SnapshotEngine::restoreSnapshot(SymbolicEngine *currentSymEngine, CONTEXT *ctx)
+VOID SnapshotEngine::restoreSnapshot(SymbolicEngine *currentSymEngine, TaintEngine *currentTaintEngine, CONTEXT *ctx)
 {
   /* 1 - Restore all memory modification. */
   list< std::pair<UINT64, UINT8> >::iterator i;
@@ -54,7 +57,10 @@ VOID SnapshotEngine::restoreSnapshot(SymbolicEngine *currentSymEngine, CONTEXT *
   /* 2 - Restore current symbolic engine state */
   *currentSymEngine = *this->snapshotSymEngine;
 
-  /* 3 - Restore Pin registers context */
+  /* 3 - Restore current taint engine state */
+  *currentTaintEngine = *this->snapshotTaintEngine;
+
+  /* 4 - Restore Pin registers context */
   PIN_SaveContext(&this->pinCtx, ctx);
 
   std::cout << "[restore snapshot]" << std::endl;