From 6f84441b47ca4e9bfe17216d6e00aec44a0d1194 Mon Sep 17 00:00:00 2001
From: Maria <mashak1407@gmail.com>
Date: Mon, 10 Jun 2024 15:02:23 +0300
Subject: [PATCH] Add taint offset and fix

---
 configs/annotations.json            | 14 +++++++--
 include/klee/Expr/Expr.h            |  2 +-
 lib/Core/SpecialFunctionHandler.cpp | 44 +++++++++++++++++++++++------
 lib/Module/Annotation.cpp           |  8 ++----
 4 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/configs/annotations.json b/configs/annotations.json
index 1df2c2f62c..f8870e2775 100644
--- a/configs/annotations.json
+++ b/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": []
   }
 }
\ No newline at end of file
diff --git a/include/klee/Expr/Expr.h b/include/klee/Expr/Expr.h
index 20a31ec1d3..8f683b16c7 100644
--- 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
index 50960177d3..6ad491edf5 100644
--- a/lib/Core/SpecialFunctionHandler.cpp
+++ b/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,
diff --git a/lib/Module/Annotation.cpp b/lib/Module/Annotation.cpp
index 6a40c6c6fa..42cb68d0d0 100644
--- a/lib/Module/Annotation.cpp
+++ b/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) {}