Add taint offset and fix

UnitTestBot · Jun 10, 2024 · cf083b6 · cf083b6
1 parent 84af116
commit cf083b6
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 21 deletions.
diff --git a/include/klee/Expr/Expr.h b/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());
   }
 

diff --git a/lib/Core/SpecialFunctionHandler.cpp b/lib/Core/SpecialFunctionHandler.cpp
@@ -1299,9 +1299,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(
@@ -1315,9 +1321,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(
@@ -1331,9 +1343,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(
@@ -1347,9 +1365,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,
@@ -1364,6 +1388,5 @@ void SpecialFunctionHandler::handleTaintHit(klee::ExecutionState &state,
 
   uint64_t taintHits = dyn_cast<ConstantExpr>(arguments[0])->getZExtValue();
   size_t taintSink = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
-//  printf("klee_taint_hit hits: %zu sink: %zu\n", taintHits, taintSink);
   executor.terminateStateOnTargetTaintError(state, taintHits, taintSink);
 }
diff --git a/lib/Module/Annotation.cpp b/lib/Module/Annotation.cpp
@@ -140,12 +140,7 @@ 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()) {
     klee_error("Annotation Taint: Incorrect value format, must has taint type");
   }
@@ -166,7 +161,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 +196,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) {}