diff --git a/src/includes/SnapshotEngine.h b/src/includes/SnapshotEngine.h index 0d1433ed2..1650f9fa6 100644 --- 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 ¤tSymEngine, CONTEXT *ctx); + VOID restoreSnapshot(SymbolicEngine *currentSymEngine, TaintEngine *currentTaintEngine, CONTEXT *ctx); + VOID takeSnapshot(const SymbolicEngine ¤tSymEngine, const TaintEngine ¤tTaintEngine, CONTEXT *ctx); }; diff --git a/src/ir/notImplemented.cpp b/src/ir/notImplemented.cpp index 9adf0bbb3..fde73353b 100644 --- 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 index edca239ab..4830bf599 100644 --- a/src/snapshotEngine/snapshotEngine.cpp +++ b/src/snapshotEngine/snapshotEngine.cpp @@ -26,7 +26,7 @@ VOID SnapshotEngine::addModification(UINT64 mem, UINT8 byte) /* Enable the snapshot engine. */ -VOID SnapshotEngine::takeSnapshot(const SymbolicEngine ¤tSymEngine, CONTEXT *ctx) +VOID SnapshotEngine::takeSnapshot(const SymbolicEngine ¤tSymEngine, const TaintEngine ¤tTaintEngine, CONTEXT *ctx) { /* 1 - Unlock the engine */ this->locked = UNLOCKED; @@ -34,7 +34,10 @@ VOID SnapshotEngine::takeSnapshot(const SymbolicEngine ¤tSymEngine, CONTEX /* 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; @@ -42,7 +45,7 @@ VOID SnapshotEngine::takeSnapshot(const SymbolicEngine ¤tSymEngine, CONTEX /* 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 >::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;