Fix bugs, change annotation config and calc AnnotationData for Executor

mamaria-k · mamaria-k · commit bddcfc0670b4 · 2024-05-26T02:13:09.000+03:00
diff --git a/configs/annotations.json b/configs/annotations.json
diff --git a/include/klee/Core/Interpreter.h b/include/klee/Core/Interpreter.h
@@ -10,8 +10,8 @@
 #define KLEE_INTERPRETER_H
 
 #include "TerminationTypes.h"
-#include "klee/Module/Annotation.h"
 
+#include "klee/Module/AnnotationsData.h"
 #include "klee/Module/SarifReport.h"
 
 #include <cstdint>
@@ -122,17 +122,13 @@ class Interpreter {
     ModuleOptions(const std::string &_LibraryDir,
                   const std::string &_EntryPoint, const std::string &_OptSuffix,
                   const std::string &_MainCurrentName,
-                  const std::string &_MainNameAfterMock,
-                  const std::string &_AnnotationsFile,
-                  const std::string &_TaintAnnotationsFile, bool _Optimize,
+                  const std::string &_MainNameAfterMock, bool _Optimize,
                   bool _Simplify, bool _CheckDivZero, bool _CheckOvershift,
                   bool _AnnotateOnlyExternal, bool _WithFPRuntime,
                   bool _WithPOSIXRuntime)
         : LibraryDir(_LibraryDir), EntryPoint(_EntryPoint),
           OptSuffix(_OptSuffix), MainCurrentName(_MainCurrentName),
-          MainNameAfterMock(_MainNameAfterMock),
-          AnnotationsFile(_AnnotationsFile),
-          TaintAnnotationsFile(_TaintAnnotationsFile), Optimize(_Optimize),
+          MainNameAfterMock(_MainNameAfterMock), Optimize(_Optimize),
           Simplify(_Simplify), CheckDivZero(_CheckDivZero),
           CheckOvershift(_CheckOvershift),
           AnnotateOnlyExternal(_AnnotateOnlyExternal),
@@ -157,6 +153,8 @@ class Interpreter {
     MockPolicy Mock;
     MockStrategyKind MockStrategy;
     MockMutableGlobalsPolicy MockMutableGlobals;
+    std::string AnnotationsFile;
+    std::string TaintAnnotationsFile;
 
     InterpreterOptions(std::optional<SarifReport> Paths)
         : MakeConcreteSymbolic(false), Guidance(GuidanceKind::NoGuidance),
diff --git a/include/klee/Core/MockBuilder.h b/include/klee/Core/MockBuilder.h
@@ -38,7 +38,7 @@ class MockBuilder {
   std::set<std::string> &mainModuleFunctions;
   std::set<std::string> &mainModuleGlobals;
 
-  AnnotationsData annotationsData;
+  const AnnotationsData &annotationsData;
 
   void initMockModule();
   void buildMockMain();
@@ -86,7 +86,8 @@ class MockBuilder {
               std::vector<std::pair<std::string, std::string>> &redefinitions,
               InterpreterHandler *interpreterHandler,
               std::set<std::string> &mainModuleFunctions,
-              std::set<std::string> &mainModuleGlobals);
+              std::set<std::string> &mainModuleGlobals,
+              const AnnotationsData &annotationsData);
 
   std::unique_ptr<llvm::Module> build();
   void buildAllocSource(llvm::Value *prev, llvm::Type *elemType,
diff --git a/include/klee/Core/TerminationTypes.h b/include/klee/Core/TerminationTypes.h
@@ -75,7 +75,8 @@ enum class StateTerminationClass : std::uint8_t {
   TTMARK(EARLYALGORITHM, 72U)                                                  \
   TTYPE(SilentExit, 80U, "")                                                   \
   TTMARK(EARLYUSER, 80U)                                                       \
-  TTMARK(END, 80U)
+  TTMARK(END, 80U)                                                             \
+  TTYPE(Taint, 90U, "taint.err")
 
 ///@brief Reason an ExecutionState got terminated.
 enum class StateTerminationType : std::uint8_t {
diff --git a/include/klee/Module/AnnotationsData.h b/include/klee/Module/AnnotationsData.h
@@ -10,8 +10,9 @@ struct AnnotationsData final {
   AnnotationsMap annotations;
   TaintAnnotation taintAnnotation;
 
-  explicit AnnotationsData(const std::string &annotationsFile,
-                           const std::string &taintAnnotationsFile);
+  explicit AnnotationsData(
+      const std::string &annotationsFile = std::string(),
+      const std::string &taintAnnotationsFile = std::string());
   virtual ~AnnotationsData();
 };
 
diff --git a/include/klee/Module/TaintAnnotation.h b/include/klee/Module/TaintAnnotation.h
@@ -13,16 +13,17 @@ using json = nlohmann::json;
 
 namespace klee {
 
+// TODO: set<size_t> to set<pair<size_t, size)t>>, where pair is sink and target
 using TaintSinksSourcesMap = std::map<size_t, std::set<size_t>>;
 
 class TaintAnnotation final {
 public:
   TaintSinksSourcesMap sinksToSources;
   std::map<std::string, size_t> sinks;
   std::map<std::string, size_t> sources;
+  // TODO: add map from size_t target to string
 
-  explicit TaintAnnotation();
-  explicit TaintAnnotation(const std::string &taintAnnotationsFile);
+  explicit TaintAnnotation(const std::string &path);
   virtual ~TaintAnnotation();
 };
 
diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
@@ -505,7 +505,9 @@ const std::unordered_set<Intrinsic::ID> Executor::modelledFPIntrinsics = {
 
 Executor::Executor(LLVMContext &ctx, const InterpreterOptions &opts,
                    InterpreterHandler *ih)
-    : Interpreter(opts), interpreterHandler(ih), searcher(nullptr),
+    : Interpreter(opts),
+      annotationsData(opts.AnnotationsFile, opts.TaintAnnotationsFile),
+      interpreterHandler(ih), searcher(nullptr),
       externalDispatcher(new ExternalDispatcher(ctx)), statsTracker(0),
       pathWriter(0), symPathWriter(0),
       specialFunctionHandler(0), timers{time::Span(TimerInterval)},
@@ -636,7 +638,8 @@ llvm::Module *Executor::setModule(
       !opts.AnnotationsFile.empty()) {
     MockBuilder mockBuilder(kmodule->module.get(), opts, interpreterOpts,
                             ignoredExternals, redefinitions, interpreterHandler,
-                            mainModuleFunctions, mainModuleGlobals);
+                            mainModuleFunctions, mainModuleGlobals,
+                            annotationsData);
     std::unique_ptr<llvm::Module> mockModule = mockBuilder.build();
 
     std::vector<std::unique_ptr<llvm::Module>> mockModules;
@@ -5022,6 +5025,22 @@ void Executor::terminateStateOnTargetError(ExecutionState &state,
       state, new ErrorEvent(locationOf(state), terminationType, messaget));
 }
 
+// TODO: add taint target errors to taint-annotations.json and change function
+void Executor::terminateStateOnTaintError(ExecutionState &state, size_t sink) {
+  //  reportStateOnTargetError(state, error);
+
+  const auto &sinks = annotationsData.taintAnnotation.sinks;
+  auto taintSink = std::find_if(sinks.begin(), sinks.end(),
+                                [&](const std::pair<std::string, size_t> &i) {
+                                  return i.second == sink;
+                                });
+  if (taintSink == std::end(sinks)) {
+    klee_error("Unknown sink index");
+  }
+  terminateStateOnProgramError(state, taintSink->first + " taint error",
+                               StateTerminationType::Taint);
+}
+
 void Executor::terminateStateOnError(ExecutionState &state,
                                      const llvm::Twine &messaget,
                                      StateTerminationType terminationType,
diff --git a/lib/Core/Executor.h b/lib/Core/Executor.h
@@ -31,6 +31,7 @@
 #include "klee/Expr/Constraints.h"
 #include "klee/Expr/SourceBuilder.h"
 #include "klee/Expr/SymbolicSource.h"
+#include "klee/Module/AnnotationsData.h"
 #include "klee/Module/Cell.h"
 #include "klee/Module/KInstruction.h"
 #include "klee/Module/KModule.h"
@@ -127,6 +128,8 @@ class Executor : public Interpreter {
   RNG theRNG;
 
 private:
+  AnnotationsData annotationsData;
+
   int *errno_addr;
 
   size_t maxNewWriteableOSSize = 0;
@@ -631,6 +634,8 @@ class Executor : public Interpreter {
   /// Then just call `terminateStateOnError`
   void terminateStateOnTargetError(ExecutionState &state, ReachWithError error);
 
+  void terminateStateOnTaintError(ExecutionState &state, size_t sink);
+
   /// Call error handler and terminate state in case of program errors
   /// (e.g. free()ing globals, out-of-bound accesses)
   void terminateStateOnProgramError(ExecutionState &state,
diff --git a/lib/Core/MockBuilder.cpp b/lib/Core/MockBuilder.cpp
@@ -89,14 +89,14 @@ MockBuilder::MockBuilder(
     std::vector<std::pair<std::string, std::string>> &redefinitions,
     InterpreterHandler *interpreterHandler,
     std::set<std::string> &mainModuleFunctions,
-    std::set<std::string> &mainModuleGlobals)
+    std::set<std::string> &mainModuleGlobals,
+    const AnnotationsData &annotationsData)
     : userModule(initModule), ctx(initModule->getContext()), opts(opts),
       interpreterOptions(interpreterOptions),
       ignoredExternals(ignoredExternals), redefinitions(redefinitions),
       interpreterHandler(interpreterHandler),
       mainModuleFunctions(mainModuleFunctions),
-      mainModuleGlobals(mainModuleGlobals),
-      annotationsData(opts.AnnotationsFile, opts.TaintAnnotationsFile) {}
+      mainModuleGlobals(mainModuleGlobals), annotationsData(annotationsData) {}
 
 std::unique_ptr<llvm::Module> MockBuilder::build() {
   initMockModule();
@@ -535,7 +535,7 @@ MockBuilder::buildCallKleeTaintFunction(const std::string &functionName,
       {llvm::Type::getInt8PtrTy(mockModule->getContext()),
        llvm::Type::getInt64Ty(mockModule->getContext())},
       false);
-  auto kleeAddTaintCallee =
+  auto kleeTaintFunctionCallee =
       mockModule->getOrInsertFunction(functionName, kleeTaintFunctionType);
 
   //  //TODO: that's not all:
@@ -579,13 +579,15 @@ MockBuilder::buildCallKleeTaintFunction(const std::string &functionName,
   if (!source->getType()->isPointerTy() && !source->getType()->isArrayTy()) {
     beginPtr = builder->CreateAlloca(source->getType());
     builder->CreateStore(source, beginPtr);
+    beginPtr = builder->CreateBitCast(
+        beginPtr, llvm::Type::getInt8PtrTy(mockModule->getContext()));
   } else {
     beginPtr = builder->CreateBitCast(
         source, llvm::Type::getInt8PtrTy(mockModule->getContext()));
   }
 
   return builder->CreateCall(
-      kleeAddTaintCallee,
+      kleeTaintFunctionCallee,
       {beginPtr, llvm::ConstantInt::get(mockModule->getContext(),
                                         llvm::APInt(64, taint, false))});
 }
@@ -847,6 +849,14 @@ void MockBuilder::buildAnnotationForExternalFunctionReturn(
   std::string retName = "ret_" + func->getName().str();
   llvm::Value *retValuePtr = builder->CreateAlloca(returnType, nullptr);
 
+  //  TODO: fix strange type ("fopen" mock, store instruction)
+  //  if (func->getName() == "fopen") {
+  //    buildCallKleeMakeSymbolic("klee_make_mock", retValuePtr, returnType,
+  //                              func->getName().str());
+  //    llvm::Value *retValue = builder->CreateLoad(returnType, retValuePtr,
+  //    retName); builder->CreateRet(retValue); return;
+  //  }
+
   if (returnType->isPointerTy() && (allocSourcePtr || mustInitNull)) {
     processingValue(retValuePtr, returnType, allocSourcePtr,
                     mustInitNull || maybeInitNull);
diff --git a/lib/Core/SpecialFunctionHandler.cpp b/lib/Core/SpecialFunctionHandler.cpp
@@ -135,6 +135,7 @@ static SpecialFunctionHandler::HandlerInfo handlerInfo[] = {
     add("klee_check_taint_source", handleCheckTaintSource, true),
     add("klee_check_taint_sink", handleCheckTaintSink, true),
     add("klee_taint_sink_hit", handleTaintSinkHit, false),
+
 #ifdef SUPPORT_KLEE_EH_CXX
     add("_klee_eh_Unwind_RaiseException_impl", handleEhUnwindRaiseExceptionImpl,
         false),
@@ -1220,8 +1221,6 @@ void SpecialFunctionHandler::handleFAbs(ExecutionState &state,
   executor.bindLocal(target, state, result);
 }
 
-// TODO: update definitions after adding taint memory
-
 void SpecialFunctionHandler::handleAddTaint(klee::ExecutionState &state,
                                             klee::KInstruction *target,
                                             std::vector<ref<Expr>> &arguments) {
@@ -1231,8 +1230,6 @@ void SpecialFunctionHandler::handleAddTaint(klee::ExecutionState &state,
                                        "klee_add_taint(void*, size_t)");
     return;
   }
-
-  //  executor.executeMemoryOperation(state, true)
 }
 
 void SpecialFunctionHandler::handleClearTaint(
@@ -1244,8 +1241,6 @@ void SpecialFunctionHandler::handleClearTaint(
                                        "klee_clear_taint(void*, size_t)");
     return;
   }
-
-  //  executor.executeMemoryOperation(state, true)
 }
 
 void SpecialFunctionHandler::handleCheckTaintSource(
@@ -1257,9 +1252,6 @@ void SpecialFunctionHandler::handleCheckTaintSource(
                "klee_check_taint_source(void*, size_t)");
     return;
   }
-
-  //  ref<Expr> result = ConstantExpr::create(true, Expr::Bool);
-  //  executor.bindLocal(target, state, result);
 }
 
 void SpecialFunctionHandler::handleCheckTaintSink(
@@ -1285,6 +1277,12 @@ void SpecialFunctionHandler::handleTaintSinkHit(
     return;
   }
 
-  //  executor.terminateStateOnError(state, "FormatString",
-  //  StateTerminationType::User);
+  char *end = nullptr;
+  size_t sink = strtoul(arguments[0]->toString().c_str(), &end, 10);
+  if (*end != '\0' || errno == ERANGE) {
+    executor.terminateStateOnUserError(
+        state, "Incorrect argument 0 to klee_taint_sink_hit(size_t)");
+  }
+
+  executor.terminateStateOnTaintError(state, sink);
 }
diff --git a/lib/Module/AnnotationsData.cpp b/lib/Module/AnnotationsData.cpp
@@ -2,12 +2,9 @@
 
 namespace klee {
 
-klee::AnnotationsData::AnnotationsData(
-    const std::string &annotationsFile,
-    const std::string &taintAnnotationsFile) {
-  taintAnnotation = taintAnnotationsFile.empty()
-                        ? TaintAnnotation()
-                        : TaintAnnotation(taintAnnotationsFile);
+klee::AnnotationsData::AnnotationsData(const std::string &annotationsFile,
+                                       const std::string &taintAnnotationsFile)
+    : taintAnnotation(taintAnnotationsFile) {
   annotations = parseAnnotations(annotationsFile);
 }
 
diff --git a/lib/Module/TaintAnnotation.cpp b/lib/Module/TaintAnnotation.cpp
@@ -5,18 +5,18 @@
 
 namespace klee {
 
-TaintAnnotation::TaintAnnotation() = default;
+TaintAnnotation::TaintAnnotation(const std::string &path) {
+  if (path.empty()) {
+    return;
+  }
 
-TaintAnnotation::TaintAnnotation(const std::string &taintAnnotationsFilePath) {
-  std::ifstream taintAnnotationsFile(taintAnnotationsFilePath);
+  std::ifstream taintAnnotationsFile(path);
   if (!taintAnnotationsFile.good()) {
-    klee_error("Taint annotation: Opening %s failed.",
-               taintAnnotationsFilePath.c_str());
+    klee_error("Taint annotation: Opening %s failed.", path.c_str());
   }
   json taintAnnotationsJson = json::parse(taintAnnotationsFile, nullptr, false);
   if (taintAnnotationsJson.is_discarded()) {
-    klee_error("Taint annotation: Parsing JSON %s failed.",
-               taintAnnotationsFilePath.c_str());
+    klee_error("Taint annotation: Parsing JSON %s failed.", path.c_str());
   }
 
   std::set<std::string> sourcesStrs;
diff --git a/tools/klee/main.cpp b/tools/klee/main.cpp
@@ -15,6 +15,7 @@
 #include "klee/Core/Context.h"
 #include "klee/Core/Interpreter.h"
 #include "klee/Core/TargetedExecutionReporter.h"
+#include "klee/Module/AnnotationsData.h"
 #include "klee/Module/LocationInfo.h"
 #include "klee/Module/SarifReport.h"
 #include "klee/Module/TargetForest.h"
@@ -404,15 +405,16 @@ cl::opt<MockMutableGlobalsPolicy> MockMutableGlobals(
 
 cl::opt<std::string>
     AnnotationsFile("annotations", cl::desc("Path to the annotation JSON file"),
-                    cl::value_desc("path file"), cl::cat(MockCat));
+                    cl::init(std::string()), cl::value_desc("path file"),
+                    cl::cat(MockCat));
 
 cl::opt<bool> AnnotateOnlyExternal(
     "annotate-only-external",
     cl::desc("Ignore annotations for defined function (default=false)"),
     cl::init(false), cl::cat(MockCat));
 
 cl::opt<std::string>
-    TaintAnnotationsFile("taint-annotations",
+    TaintAnnotationsFile("taint-annotations", cl::init(std::string()),
                          cl::desc("Path to the taint annotations JSON file"),
                          cl::value_desc("path file"), cl::cat(MockCat));
 
@@ -2179,8 +2181,6 @@ int main(int argc, char **argv, char **envp) {
       LibraryDir, EntryPoint, opt_suffix,
       /*MainCurrentName=*/EntryPoint,
       /*MainNameAfterMock=*/"__klee_mock_wrapped_main",
-      /*AnnotationsFile=*/AnnotationsFile,
-      /*TaintAnnotationsFile=*/TaintAnnotationsFile,
       /*Optimize=*/OptimizeModule,
       /*Simplify*/ SimplifyModule,
       /*CheckDivZero=*/CheckDivZero,
@@ -2390,6 +2390,8 @@ int main(int argc, char **argv, char **envp) {
   IOpts.Mock = Mock;
   IOpts.MockStrategy = MockStrategy;
   IOpts.MockMutableGlobals = MockMutableGlobals;
+  IOpts.AnnotationsFile = AnnotationsFile;
+  IOpts.TaintAnnotationsFile = TaintAnnotationsFile;
 
   std::unique_ptr<Interpreter> interpreter(
       Interpreter::create(ctx, IOpts, handler.get()));
