S2EExecutor: it is now safe to throw exceptions in onFork()

The patch ensures that the execution is not left in an
inconsistent state.

Signed-off-by: Vitaly Chipounov <vitaly.chipounov@epfl.ch>

Conflicts:
	klee/lib/Core/Executor.cpp
	qemu/s2e/S2EExecutor.cpp

klee/include/klee/Executor.h

@@ -295,6 +295,13 @@ class Executor : public Interpreter {
   /// so that the branched state gets it as well.
   virtual void notifyBranch(ExecutionState &state);
 
+  /// When the fork is complete and state properly updated,
+  /// notify the S2EExecutor, so that it can generate an onFork event.
+  /// Sending notification after the fork completed
+  /// allows plugins to kill states and exit to the CPU loop safely.
+  virtual void notifyFork(ExecutionState &originalState, ref<Expr> &condition,
+                  Executor::StatePair &targets);
+
   /// Add the given (boolean) condition as a constraint on state. This
   /// function is a wrapper around the state's addConstraint function
   /// which also manages manages propogation of implied values,

klee/lib/Core/Executor.cpp

@@ -928,6 +928,13 @@ bool Executor::resolveSpeculativeState(ExecutionState &state)
 }
 
 
+void Executor::notifyFork(ExecutionState &originalState, ref<Expr> &condition,
+                          Executor::StatePair &targets)
+{
+    //Should not get here
+    assert(false && "Must go through S2E");
+}
+
 Executor::StatePair 
 Executor::fork(ExecutionState &current, ref<Expr> condition, bool isInternal) {
   condition = simplifyExpr(current, condition);
@@ -1724,6 +1731,8 @@ void Executor::executeInstruction(ExecutionState &state, KInstruction *ki) {
         transferToBasicBlock(bi->getSuccessor(0), bi->getParent(), *branches.first);
       if (branches.second)
         transferToBasicBlock(bi->getSuccessor(1), bi->getParent(), *branches.second);
+
+      notifyFork(state, cond, branches);
     }
     break;
   }
@@ -1891,7 +1900,9 @@ void Executor::executeInstruction(ExecutionState &state, KInstruction *ki) {
         bool success = solver->getValue(*free, v, value);
         assert(success && "FIXME: Unhandled solver failure");
         (void) success;
-        StatePair res = fork(*free, EqExpr::create(v, value), true);
+        ref<Expr> cond = EqExpr::create(v, value);
+        StatePair res = fork(*free, cond, true);
+        notifyFork(*free, cond, res);
         if (res.first) {
           uint64_t addr = value->getZExtValue();
           if (legalFunctions.count(addr)) {
@@ -3184,8 +3195,10 @@ void Executor::executeAlloc(ExecutionState &state,
       example = tmp;
     }
 
-    StatePair fixedSize = fork(state, EqExpr::create(example, size), true);
-
+    ref<Expr> cond = EqExpr::create(example, size);
+    StatePair fixedSize = fork(state, cond, true);
+    notifyFork(state, cond, fixedSize);
+
     if (fixedSize.second) { 
       // Check for exactly two values
       ref<ConstantExpr> tmp;
@@ -3237,7 +3250,9 @@ void Executor::executeAlloc(ExecutionState &state,
 void Executor::executeFree(ExecutionState &state,
                            ref<Expr> address,
                            KInstruction *target) {
-  StatePair zeroPointer = fork(state, Expr::createIsZero(address), true);
+  ref<Expr> cond = Expr::createIsZero(address);
+  StatePair zeroPointer = fork(state, cond, true);
+
   if (zeroPointer.first) {
     if (target)
       bindLocal(target, *zeroPointer.first, Expr::createPointer(0));
@@ -3266,6 +3281,9 @@ void Executor::executeFree(ExecutionState &state,
       }
     }
   }
+
+
+  notifyFork(state, cond, zeroPointer);
 }
 
 void Executor::resolveExact(ExecutionState &state,
@@ -3283,6 +3301,7 @@ void Executor::resolveExact(ExecutionState &state,
     ref<Expr> inBounds = EqExpr::create(p, it->first->getBaseExpr());
 
     StatePair branches = fork(*unbound, inBounds, true);
+    notifyFork(*unbound, inBounds, branches);
 
     if (branches.first)
       results.push_back(std::make_pair(*it, branches.first));
@@ -3349,6 +3368,9 @@ void Executor::executeMemoryOperation(ExecutionState &state,
             //The forked state will have to re-execute the memory op
             branches.second->pc = branches.second->prevPC;
           }
+
+          notifyFork(state, condition, branches);
+
       } else {
           solver->setTimeout(stpTimeout);
           if (!state.addressSpace.resolveOne(state, solver, address, op, success)) {
@@ -3496,6 +3518,8 @@ void Executor::executeMemoryOperation(ExecutionState &state,
       }
     }
 
+    notifyFork(state, inBounds, branches);
+
     unbound = branches.second;
     if (!unbound)
       break;

qemu/s2e/S2EExecutor.cpp

@@ -575,11 +575,12 @@ void S2EExecutor::handleForkAndConcretize(Executor* executor,
             //Will have to reexecute handleForkAndConcretize in the speculative state
             sp.second->pc = sp.second->prevPC;
         }
+        s2eExecutor->bindLocal(target, *state, concreteAddress);
+        s2eExecutor->notifyFork(*state, condition, sp);
     } else {
         state->addConstraint(condition);
+        s2eExecutor->bindLocal(target, *state, concreteAddress);
     }
-
-    s2eExecutor->bindLocal(target, *state, concreteAddress);
 }
 
 void S2EExecutor::handleMakeSymbolic(Executor* executor,
@@ -2030,6 +2031,34 @@ void S2EExecutor::deleteState(klee::ExecutionState *state)
     m_deletedStates.push_back(static_cast<S2EExecutionState*>(state));
 }
 
+void S2EExecutor::notifyFork(ExecutionState &originalState, ref<Expr> &condition,
+                             Executor::StatePair &targets)
+{
+    if (targets.first == NULL || targets.second == NULL) {
+        return;
+    }
+
+    std::vector<S2EExecutionState*> newStates(2);
+    std::vector<ref<Expr> > newConditions(2);
+
+    S2EExecutionState *state = static_cast<S2EExecutionState*>(&originalState);
+    newStates[0] = static_cast<S2EExecutionState*>(targets.first);
+    newStates[1] = static_cast<S2EExecutionState*>(targets.second);
+
+    newConditions[0] = condition;
+    newConditions[1] = klee::NotExpr::create(condition);
+
+    try {
+        m_s2e->getCorePlugin()->onStateFork.emit(state, newStates, newConditions);
+    } catch (CpuExitException e) {
+        if (state->stack.size() != 1) {
+            state->m_needFinalizeTBExec = true;
+            state->m_forkAborted = true;
+        }
+        throw e;
+    }
+}
+
 void S2EExecutor::doStateFork(S2EExecutionState *originalState,
                  const vector<S2EExecutionState*>& newStates,
                  const vector<ref<Expr> >& newConditions)
@@ -2059,10 +2088,7 @@ void S2EExecutor::doStateFork(S2EExecutionState *originalState,
         foreach(const StackFrame& fr, originalState->stack) {
             m_s2e->getDebugStream() << fr.kf->function->getName().str() << '\n';
         }
-    }
-
-    m_s2e->getCorePlugin()->onStateFork.emit(originalState,
-                                             newStates, newConditions);
+    }   
 }
 
 S2EExecutor::StatePair S2EExecutor::fork(ExecutionState &current,

qemu/s2e/S2EExecutor.h

@@ -286,8 +286,8 @@ class S2EExecutor : public klee::Executor
                        S2EExecutionState* newState);
 
     void doStateFork(S2EExecutionState *originalState,
-                     const std::vector<S2EExecutionState*>& newStates,
-                     const std::vector<klee::ref<klee::Expr> >& conditions);
+                        const std::vector<S2EExecutionState*>& newStates,
+                        const std::vector<klee::ref<klee::Expr> >& conditions);
 
     void doLoadBalancing();
 
@@ -320,6 +320,9 @@ class S2EExecutor : public klee::Executor
 
     void notifyBranch(klee::ExecutionState &state);
 
+    void notifyFork(klee::ExecutionState &originalState, klee::ref<klee::Expr> &condition,
+                    StatePair &targets);
+
     /** Kills the specified state and raises an exception to exit the cpu loop */
     virtual void terminateState(klee::ExecutionState &state);