Add taint to lazy object init and small methods

Author: mamaria-k

include/klee/Core/BranchTypes.h

@@ -25,7 +25,8 @@
   BTYPE(Realloc, 8U)                                                           \
   BTYPE(Free, 9U)                                                              \
   BTYPE(GetVal, 10U)                                                           \
-  BMARK(END, 10U)
+  BTYPE(Taint, 11U)                                                            \
+  BMARK(END, 11U)
 /// \endcond
 
 /** @enum BranchType

include/klee/Expr/Expr.h

@@ -416,6 +416,7 @@ class Expr {
   static ref<ConstantExpr> createFalse();
 
   static ref<ConstantExpr> createEmptyTaint();
+  static ref<ConstantExpr> createTaintBySource(uint64_t source);
   static ref<Expr> combineTaints(const ref<Expr> &taintL,
                                  const ref<Expr> &taintR);
 

lib/Core/AddressSpace.cpp

@@ -9,6 +9,8 @@
 
 #include "AddressSpace.h"
 
+#include <utility>
+
 #include "ExecutionState.h"
 #include "Memory.h"
 #include "TimingSolver.h"
@@ -80,19 +82,21 @@ ObjectPair AddressSpace::findObject(const MemoryObject *mo) const {
              : ObjectPair(nullptr, nullptr);
 }
 
-RefObjectPair AddressSpace::lazyInitializeObject(const MemoryObject *mo) const {
-  return RefObjectPair(mo, new ObjectState(mo, mo->content, mo->type));
+RefObjectPair AddressSpace::lazyInitializeObject(const MemoryObject *mo,
+                                                 ref<Expr> taint) const {
+  return RefObjectPair(
+      mo, new ObjectState(mo, mo->content, mo->type, std::move(taint)));
 }
 
-RefObjectPair
-AddressSpace::findOrLazyInitializeObject(const MemoryObject *mo) const {
+RefObjectPair AddressSpace::findOrLazyInitializeObject(const MemoryObject *mo,
+                                                       ref<Expr> taint) const {
   ObjectPair op = findObject(mo);
   if (op.first && op.second) {
     return RefObjectPair(op.first, op.second);
   }
   assert(!op.first && !op.second);
   if (mo->isLazyInitialized) {
-    return lazyInitializeObject(mo);
+    return lazyInitializeObject(mo, std::move(taint));
   }
   return RefObjectPair(nullptr, nullptr);
 }

lib/Core/AddressSpace.h

@@ -136,8 +136,11 @@ class AddressSpace {
 
   /// Lookup a binding from a MemoryObject.
   ObjectPair findObject(const MemoryObject *mo) const;
-  RefObjectPair lazyInitializeObject(const MemoryObject *mo) const;
-  RefObjectPair findOrLazyInitializeObject(const MemoryObject *mo) const;
+  RefObjectPair lazyInitializeObject(const MemoryObject *mo,
+                                     ref<Expr> taint) const;
+  RefObjectPair
+  findOrLazyInitializeObject(const MemoryObject *mo,
+                             ref<Expr> taint = Expr::createEmptyTaint()) const;
 
   /// Copy the concrete values of all managed ObjectStates into the
   /// actual system memory location they were allocated at.

lib/Core/Executor.cpp

@@ -5985,7 +5985,7 @@ bool Executor::resolveMemoryObjects(
             state, basePointer, target, baseTargetType, minObjectSize, sizeExpr,
             false, checkOutOfBounds, UseSymbolicSizeLazyInit);
         RefObjectPair op = state.addressSpace.findOrLazyInitializeObject(
-            idLazyInitialization.get());
+            idLazyInitialization.get(), address->getTaint());
         state.addressSpace.bindObject(op.first, op.second.get());
         mayBeResolvedMemoryObjects.push_back(idLazyInitialization);
       }

lib/Core/Memory.cpp

@@ -106,22 +106,26 @@ std::string MemoryObject::getAllocInfo() const {
 
 /***/
 
-ObjectState::ObjectState(const MemoryObject *mo, const Array *array, KType *dt)
+ObjectState::ObjectState(const MemoryObject *mo, const Array *array, KType *dt,
+                         ref<Expr> defaultTaintValue)
     : copyOnWriteOwner(0), object(mo), valueOS(ObjectStage(array, nullptr)),
       baseOS(ObjectStage(array->size, Expr::createPointer(0), false,
                          Context::get().getPointerWidth())),
-      taintOS(ObjectStage(array->size, Expr::createEmptyTaint(), false, Expr::TaintWidth)),
+      taintOS(
+          ObjectStage(array->size, defaultTaintValue, false, Expr::TaintWidth)),
       lastUpdate(nullptr), size(array->size), dynamicType(dt), readOnly(false) {
   baseOS.initializeToZero();
   taintOS.initializeToZero();
 }
 
-ObjectState::ObjectState(const MemoryObject *mo, KType *dt)
+ObjectState::ObjectState(const MemoryObject *mo, KType *dt,
+                         ref<Expr> defaultTaintValue)
     : copyOnWriteOwner(0), object(mo),
       valueOS(ObjectStage(mo->getSizeExpr(), nullptr)),
       baseOS(ObjectStage(mo->getSizeExpr(), Expr::createPointer(0), false,
                          Context::get().getPointerWidth())),
-      taintOS(ObjectStage(mo->getSizeExpr(), Expr::createEmptyTaint(), false, Expr::TaintWidth)),
+      taintOS(ObjectStage(mo->getSizeExpr(), defaultTaintValue, false,
+                          Expr::TaintWidth)),
       lastUpdate(nullptr), size(mo->getSizeExpr()), dynamicType(dt),
       readOnly(false) {
   baseOS.initializeToZero();

lib/Core/Memory.h

@@ -310,8 +310,10 @@ class ObjectState {
 
   /// Create a new object state for the given memory
   // For objects in memory
-  ObjectState(const MemoryObject *mo, const Array *array, KType *dt);
-  ObjectState(const MemoryObject *mo, KType *dt);
+  ObjectState(const MemoryObject *mo, const Array *array, KType *dt,
+              ref<Expr> defaultTaintValue = Expr::createEmptyTaint());
+  ObjectState(const MemoryObject *mo, KType *dt,
+              ref<Expr> defaultTaintValue = Expr::createEmptyTaint());
 
   // For symbolic objects not in memory (hack)
 

lib/Expr/Expr.cpp

@@ -553,6 +553,10 @@ ref<ConstantExpr> Expr::createEmptyTaint() {
   return ConstantExpr::create(0, Expr::Int64);
 }
 
+ref<ConstantExpr> Expr::createTaintBySource(uint64_t source) {
+  return ConstantExpr::create((1 << source), Expr::Int64);
+}
+
 ref<Expr> Expr::combineTaints(const ref<Expr> &taintL,
                               const ref<Expr> &taintR) {
   if (SelectExpr *sel = dyn_cast<SelectExpr>(taintL)) {
@@ -2991,7 +2995,7 @@ ref<Expr> ConstantPointerExpr::create(const ref<ConstantExpr> &b,
                                  combineTaints(RHS));                          \
     } else {                                                                   \
       auto value = _e_op::create(getValue(), RHS->getValue());                 \
-      if (getTaint()->isZero() && RHS->getTaint()->isZero()) {                   \
+      if (getTaint()->isZero() && RHS->getTaint()->isZero()) {                 \
         return value;                                                          \
       } else {                                                                 \
         return PointerExpr::create(ConstantExpr::create(0, value->getWidth()), \