Add taint offset and fix

Author: mamaria-k

configs/annotations.json

@@ -268,9 +268,6 @@
       [
         "TaintSink::FormatString",
         "TaintSink::SensitiveDataLeak"
-      ],
-      [
-        "TaintSink::SensitiveDataLeak"
       ]
     ],
     "properties": []
@@ -2337,5 +2334,16 @@
       ]
     ],
     "properties": []
+  },
+  "vprintf_s": {
+    "name": "vprintf_s",
+    "annotation": [
+      [],
+      [],
+      [
+        "TaintSink::FormatString"
+      ]
+    ],
+    "properties": []
   }
 }

include/klee/Expr/Expr.h

@@ -1739,7 +1739,7 @@ class PointerExpr : public NonConstantExpr {
 
   bool isKnownValue() const { return getBase()->isZero(); }
 
-  ref<ConstantExpr> combineTaints(const ref<PointerExpr> &RHS) {
+  ref<Expr> combineTaints(const ref<PointerExpr> &RHS) {
     return Expr::combineTaints(getTaint(), RHS->getTaint());
   }
 

lib/Core/SpecialFunctionHandler.cpp

@@ -1300,8 +1300,15 @@ void SpecialFunctionHandler::handleAddTaint(klee::ExecutionState &state,
 
   uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
 //  printf("klee_add_taint source: %zu\n", taintSource);
-  executor.executeChangeTaintSource(
-      state, target, executor.makePointer(arguments[0]), taintSource, true);
+
+  ref<PointerExpr> pointer = executor.makePointer(arguments[0]);
+  if (auto *p = dyn_cast<PointerExpr>(arguments[0])) {
+    if (p->isKnownValue()) {
+      pointer =
+          PointerExpr::create(p->getValue(), p->getValue(), p->getTaint());
+    }
+  }
+  executor.executeChangeTaintSource(state, target, pointer, taintSource, true);
 }
 
 void SpecialFunctionHandler::handleClearTaint(
@@ -1316,8 +1323,15 @@ void SpecialFunctionHandler::handleClearTaint(
 
   uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
 //  printf("klee_clear_taint source: %zu\n", taintSource);
-  executor.executeChangeTaintSource(
-      state, target, executor.makePointer(arguments[0]), taintSource, false);
+
+  ref<PointerExpr> pointer = executor.makePointer(arguments[0]);
+  if (auto *p = dyn_cast<PointerExpr>(arguments[0])) {
+    if (p->isKnownValue()) {
+      pointer =
+          PointerExpr::create(p->getValue(), p->getValue(), p->getTaint());
+    }
+  }
+  executor.executeChangeTaintSource(state, target, pointer, taintSource, false);
 }
 
 void SpecialFunctionHandler::handleCheckTaintSource(
@@ -1332,8 +1346,15 @@ void SpecialFunctionHandler::handleCheckTaintSource(
 
   uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
 //  printf("klee_check_taint_source source: %zu\n", taintSource);
-  executor.executeCheckTaintSource(
-      state, target, executor.makePointer(arguments[0]), taintSource);
+
+  ref<PointerExpr> pointer = executor.makePointer(arguments[0]);
+  if (auto *p = dyn_cast<PointerExpr>(arguments[0])) {
+    if (p->isKnownValue()) {
+      pointer =
+          PointerExpr::create(p->getValue(), p->getValue(), p->getTaint());
+    }
+  }
+  executor.executeCheckTaintSource(state, target, pointer, taintSource);
 }
 
 void SpecialFunctionHandler::handleGetTaintHits(
@@ -1348,8 +1369,15 @@ void SpecialFunctionHandler::handleGetTaintHits(
 
   uint64_t taintSink = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
 //  printf("klee_get_taint_hits sink: %zu\n", taintSink);
-  executor.executeGetTaintHits(state, target,
-                               executor.makePointer(arguments[0]), taintSink);
+
+  ref<PointerExpr> pointer = executor.makePointer(arguments[0]);
+  if (auto *p = dyn_cast<PointerExpr>(arguments[0])) {
+    if (p->isKnownValue()) {
+      pointer =
+          PointerExpr::create(p->getValue(), p->getValue(), p->getTaint());
+    }
+  }
+  executor.executeGetTaintHits(state, target, pointer, taintSink);
 }
 
 void SpecialFunctionHandler::handleTaintHit(klee::ExecutionState &state,

lib/Module/Annotation.cpp

@@ -140,10 +140,6 @@ Free::Free(const std::string &str) : Unknown(str) {
 Kind Free::getKind() const { return Kind::Free; }
 
 Taint::Taint(const std::string &str) : Unknown(str) {
-  if (!rawOffset.empty()) {
-    klee_error("Annotation Taint: Incorrect offset format, must be empty");
-  }
-
   taintType = rawValue.substr(0, rawValue.find(':'));
   // TODO: in the future, support typeless annotations (meaning all types)
   if (taintType.empty()) {
@@ -166,7 +162,7 @@ TaintOutput::TaintOutput(const std::string &str) : Taint(str) {}
 Kind TaintOutput::getKind() const { return Kind::TaintOutput; }
 
 /*
- * Format: TaintPropagation::{type}:{data}
+ * Format: TaintPropagation:{offset}:{type}:{data}
  */
 
 TaintPropagation::TaintPropagation(const std::string &str) : Taint(str) {
@@ -201,7 +197,7 @@ TaintPropagation::TaintPropagation(const std::string &str) : Taint(str) {
 Kind TaintPropagation::getKind() const { return Kind::TaintPropagation; }
 
 /*
- * Format: TaintSink::{type}
+ * Format: TaintSink:{offset}:{type}
  */
 
 TaintSink::TaintSink(const std::string &str) : Taint(str) {}