[WIP] Reworking implementation of operation masking sanitizer

resolve-cveassert/src/ArithmeticSanitizer.cpp

@@ -15,7 +15,7 @@
 
 #include "CVEAssert.hpp"
 #include "IRUtils.hpp"
-#include "Vulnerability.hpp"
+#include "Remediation.hpp"
 
 #include <deque>
 #include <memory>
@@ -90,8 +90,7 @@ static void widenIntOverflow(Function *F) {
   }
 }
 
-void sanitizeDivideByZero(Function *F,
-                          Vulnerability::RemediationStrategies strategy) {
+void sanitizeDivideByZero(Function *F, RemediationStrategies strategy) {
   Module *M = F->getParent();
   auto &Ctx = M->getContext();
   auto usize_ty = Type::getInt64Ty(Ctx);
@@ -100,15 +99,15 @@ void sanitizeDivideByZero(Function *F,
   std::vector<Instruction *> worklist;
 
   switch (strategy) {
-  case Vulnerability::RemediationStrategies::CONTINUE:
-  case Vulnerability::RemediationStrategies::EXIT:
-  case Vulnerability::RemediationStrategies::RECOVER:
+  case RemediationStrategies::CONTINUE:
+  case RemediationStrategies::EXIT:
+  case RemediationStrategies::RECOVER:
     break;
 
   default:
     llvm::errs() << "[CVEAssert] Error: sanitizeDivideByZero does not support "
                  << " remediation strategy defaulting to continue strategy!\n";
-    strategy = Vulnerability::RemediationStrategies::CONTINUE;
+    strategy = RemediationStrategies::CONTINUE;
     break;
   }
 
@@ -255,8 +254,7 @@ void sanitizeDivideByZero(Function *F,
   }
 }
 
-void sanitizeIntOverflow(Function *F,
-                         Vulnerability::RemediationStrategies strategy) {
+void sanitizeIntOverflow(Function *F, RemediationStrategies strategy) {
   std::vector<Instruction *> worklist;
   Module *M = F->getParent();
   auto &Ctx = M->getContext();
@@ -265,20 +263,20 @@ void sanitizeIntOverflow(Function *F,
   IRBuilder<> builder(Ctx);
 
   switch (strategy) {
-  case Vulnerability::RemediationStrategies::WIDEN:
+  case RemediationStrategies::WIDEN:
     return widenIntOverflow(F);
 
-  case Vulnerability::RemediationStrategies::RECOVER:
-  case Vulnerability::RemediationStrategies::EXIT:
-  case Vulnerability::RemediationStrategies::WRAP:
-  case Vulnerability::RemediationStrategies::SAT:
+  case RemediationStrategies::RECOVER:
+  case RemediationStrategies::EXIT:
+  case RemediationStrategies::WRAP:
+  case RemediationStrategies::SAT:
     break;
 
   default:
     llvm::errs()
         << "[CVEAssert] Error: sanitizeIntOverflow does not support "
            "remediation strategy specified defaulting to wrap strategy!\n";
-    strategy = Vulnerability::RemediationStrategies::WRAP;
+    strategy = RemediationStrategies::WRAP;
     break;
   }
 
@@ -415,7 +413,7 @@ void sanitizeIntOverflow(Function *F,
     builder.CreateBr(joinResultBB);
 
     builder.SetInsertPoint(&*joinResultBB->begin());
-    if (strategy == Vulnerability::RemediationStrategies::SAT) {
+    if (strategy == RemediationStrategies::SAT) {
       binary_inst->replaceAllUsesWith(satResult);
     } else {
       binary_inst->replaceAllUsesWith(safeResult);
@@ -425,8 +423,7 @@ void sanitizeIntOverflow(Function *F,
   }
 }
 
-void sanitizeBitShift(Function *F,
-                      Vulnerability::RemediationStrategies strategy) {
+void sanitizeBitShift(Function *F, RemediationStrategies strategy) {
   Module *M = F->getParent();
   auto &Ctx = M->getContext();
   auto usize_ty = Type::getInt64Ty(Ctx);
@@ -435,14 +432,14 @@ void sanitizeBitShift(Function *F,
   std::vector<Instruction *> worklist;
 
   switch (strategy) {
-  case Vulnerability::RemediationStrategies::EXIT:
-  case Vulnerability::RemediationStrategies::RECOVER:
+  case RemediationStrategies::EXIT:
+  case RemediationStrategies::RECOVER:
     break;
 
   default:
     llvm::errs() << "[CVEAssert] Error: sanitizeBitShift does not support "
                  << " remediation strategy defaulting to EXIT strategy!\n";
-    strategy = Vulnerability::RemediationStrategies::EXIT;
+    strategy = RemediationStrategies::EXIT;
     break;
   }
 

resolve-cveassert/src/ArithmeticSanitizer.hpp

@@ -5,11 +5,8 @@
 
 #pragma once
 
-#include "Vulnerability.hpp"
+#include "Remediation.hpp"
 #include "llvm/IR/Function.h"
-void sanitizeDivideByZero(llvm::Function *F,
-                          Vulnerability::RemediationStrategies strategy);
-void sanitizeIntOverflow(llvm::Function *F,
-                         Vulnerability::RemediationStrategies strategy);
-void sanitizeBitShift(llvm::Function *F,
-                      Vulnerability::RemediationStrategies strategy);
+void sanitizeDivideByZero(llvm::Function *F, RemediationStrategies strategy);
+void sanitizeIntOverflow(llvm::Function *F, RemediationStrategies strategy);
+void sanitizeBitShift(llvm::Function *F, RemediationStrategies strategy);

resolve-cveassert/src/BoundsCheck.cpp

@@ -18,7 +18,7 @@
 
 #include "CVEAssert.hpp"
 #include "IRUtils.hpp"
-#include "Vulnerability.hpp"
+#include "Remediation.hpp"
 
 #include <map>
 #include <unordered_set>
@@ -162,8 +162,7 @@ static Function *getOrCreateAccessOk(Module *M, BoundsClass cls) {
 }
 
 static Function *getOrCreateBoundsCheckLoadSanitizer(
-    Function *F, Type *ty, Vulnerability::RemediationStrategies strategy,
-    BoundsClass cls) {
+    Function *F, Type *ty, RemediationStrategies strategy, BoundsClass cls) {
   std::string handlerName =
       "__cve_bound_ld_" + getLLVMType(ty) + "_" + classTag(cls);
   Module *M = F->getParent();
@@ -223,8 +222,7 @@ static Function *getOrCreateBoundsCheckLoadSanitizer(
 }
 
 static Function *getOrCreateBoundsCheckStoreSanitizer(
-    Function *F, Type *ty, Vulnerability::RemediationStrategies strategy,
-    BoundsClass cls) {
+    Function *F, Type *ty, RemediationStrategies strategy, BoundsClass cls) {
   std::string handlerName =
       "__cve_bound_st_" + getLLVMType(ty) + "_" + classTag(cls);
   Module *M = F->getParent();
@@ -287,9 +285,10 @@ static Function *getOrCreateBoundsCheckStoreSanitizer(
   return resolveStoreFn;
 }
 
-static Function *getOrCreateBoundsCheckMemcpySanitizer(
-    Function *F, Vulnerability::RemediationStrategies strategy,
-    BoundsClass srcCls, BoundsClass dstCls) {
+static Function *
+getOrCreateBoundsCheckMemcpySanitizer(Function *F,
+                                      RemediationStrategies strategy,
+                                      BoundsClass srcCls, BoundsClass dstCls) {
   std::string handlerName =
       std::string("__cve_memcpy_") + classTag(srcCls) + "_" + classTag(dstCls);
   Module *M = F->getParent();
@@ -357,9 +356,10 @@ static Function *getOrCreateBoundsCheckMemcpySanitizer(
   return resolveMemmoveFn;
 }
 
-static Function *getOrCreateBoundsCheckMemmoveSanitizer(
-    Function *F, Vulnerability::RemediationStrategies strategy,
-    BoundsClass srcCls, BoundsClass dstCls) {
+static Function *
+getOrCreateBoundsCheckMemmoveSanitizer(Function *F,
+                                       RemediationStrategies strategy,
+                                       BoundsClass srcCls, BoundsClass dstCls) {
   std::string handlerName =
       std::string("__cve_memmove_") + classTag(srcCls) + "_" + classTag(dstCls);
   Module *M = F->getParent();
@@ -429,8 +429,7 @@ static Function *getOrCreateBoundsCheckMemmoveSanitizer(
 }
 
 static Function *getOrCreateBoundsCheckMemsetSanitizer(
-    Function *F, Vulnerability::RemediationStrategies strategy,
-    BoundsClass cls) {
+    Function *F, RemediationStrategies strategy, BoundsClass cls) {
   std::string handlerName = std::string("__cve_memset_") + classTag(cls);
   Module *M = F->getParent();
   LLVMContext &Ctx = M->getContext();
@@ -633,8 +632,7 @@ void instrumentGep(Function *F) {
   }
 }
 
-void instrumentMemcpy(Function *F,
-                      Vulnerability::RemediationStrategies strategy) {
+void instrumentMemcpy(Function *F, RemediationStrategies strategy) {
   LLVMContext &Ctx = F->getContext();
   IRBuilder<> builder(Ctx);
   std::vector<Instruction *> memcpyList;
@@ -691,8 +689,7 @@ void instrumentMemcpy(Function *F,
   }
 }
 
-void instrumentMemset(Function *F,
-                      Vulnerability::RemediationStrategies strategy) {
+void instrumentMemset(Function *F, RemediationStrategies strategy) {
   LLVMContext &Ctx = F->getContext();
   IRBuilder<> builder(Ctx);
   std::vector<Instruction *> memsetList;
@@ -760,8 +757,7 @@ void instrumentMemset(Function *F,
   }
 }
 
-void instrumentMemmove(Function *F,
-                       Vulnerability::RemediationStrategies strategy) {
+void instrumentMemmove(Function *F, RemediationStrategies strategy) {
   LLVMContext &Ctx = F->getContext();
   IRBuilder<> builder(Ctx);
   std::vector<Instruction *> memmoveList;
@@ -818,25 +814,24 @@ void instrumentMemmove(Function *F,
   }
 }
 
-void instrumentLoadStore(Function *F,
-                         Vulnerability::RemediationStrategies strategy) {
+void instrumentLoadStore(Function *F, RemediationStrategies strategy) {
   LLVMContext &Ctx = F->getContext();
   IRBuilder<> builder(Ctx);
 
   std::vector<LoadInst *> loadList;
   std::vector<StoreInst *> storeList;
 
   switch (strategy) {
-  case Vulnerability::RemediationStrategies::CONTINUE:
-  case Vulnerability::RemediationStrategies::EXIT:
-  case Vulnerability::RemediationStrategies::RECOVER:
+  case RemediationStrategies::CONTINUE:
+  case RemediationStrategies::EXIT:
+  case RemediationStrategies::RECOVER:
     break;
 
   default:
     llvm::errs() << "[CVEAssert] Error: instrumentLoadStore does not support "
                     "remediation strategy "
                  << "defaulting to continue strategy!\n";
-    strategy = Vulnerability::RemediationStrategies::CONTINUE;
+    strategy = RemediationStrategies::CONTINUE;
     break;
   }
 
@@ -903,8 +898,7 @@ void instrumentLoadStore(Function *F,
   }
 }
 
-void sanitizeMemInstBounds(Function *F,
-                           Vulnerability::RemediationStrategies strategy) {
+void sanitizeMemInstBounds(Function *F, RemediationStrategies strategy) {
   instrumentGep(F);
   instrumentMemcpy(F, strategy);
   instrumentMemmove(F, strategy);

resolve-cveassert/src/BoundsCheck.hpp

@@ -5,15 +5,10 @@
 
 #pragma once
 
-#include "Vulnerability.hpp"
+#include "Remediation.hpp"
 #include "llvm/IR/Function.h"
-void instrumentLoadStore(llvm::Function *F,
-                         Vulnerability::RemediationStrategies strategy);
-void instrumentMemcpy(llvm::Function *F,
-                      Vulnerability::RemediationStrategies strategy);
-void instrumentMemmove(llvm::Function *F,
-                       Vulnerability::RemediationStrategies strategy);
-void instrumentMemset(llvm::Function *F,
-                      Vulnerability::RemediationStrategies strategy);
-void sanitizeMemInstBounds(llvm::Function *F,
-                           Vulnerability::RemediationStrategies strategy);
+void instrumentLoadStore(llvm::Function *F, RemediationStrategies strategy);
+void instrumentMemcpy(llvm::Function *F, RemediationStrategies strategy);
+void instrumentMemmove(llvm::Function *F, RemediationStrategies strategy);
+void instrumentMemset(llvm::Function *F, RemediationStrategies strategy);
+void sanitizeMemInstBounds(llvm::Function *F, RemediationStrategies strategy);

resolve-cveassert/src/CVEAssert.cpp

@@ -35,6 +35,7 @@
 #include "InstrumentAllocators.hpp"
 #include "NullPointerSanitizer.hpp"
 #include "OperationMasking.hpp"
+#include "Remediation.hpp"
 #include "Vulnerability.hpp"
 
 using namespace llvm;
@@ -114,8 +115,7 @@ struct LabelCVEPass : public PassInfoMixin<LabelCVEPass> {
     vulnerabilities = Vulnerability::parseVulnerabilityFile();
   }
 
-  void applyAutomaticSanitizers(Function &F,
-                                Vulnerability::RemediationStrategies strategy) {
+  void applyAutomaticSanitizers(Function &F, RemediationStrategies strategy) {
     /// applies all automatic sanitizers (operation masking excluded)
     sanitizeFreeOfNonHeap(&F, strategy);
     sanitizeMemInstBounds(&F, strategy);
@@ -189,17 +189,17 @@ struct LabelCVEPass : public PassInfoMixin<LabelCVEPass> {
     out << F;
     out << "[CVEAssert] === Inserted Sanitizer Helpers === \n";
 
-    if (vuln.UndesirableFunction.has_value()) {
+    if (vuln.Operation.has_value()) {
       /* NOTE: We are using '0' as a temporary this will be updated future PRs
        */
-      sanitizeUndesirableOperationInFunction(&F, *vuln.UndesirableFunction, 0);
+      sanitizeContract(&F, *vuln.Operation, 0);
       result = PreservedAnalyses::none();
-      out << "[CVEAssert] === Post Sanitization of Undesirable Operation IR "
+      out << "[CVEAssert] === Post Sanitization of Masked Operation IR "
              "=== \n";
       out << F;
     }
 
-    if (vuln.Strategy == Vulnerability::RemediationStrategies::NONE) {
+    if (vuln.Strategy == RemediationStrategies::NONE) {
       out << "[CVEAssert] NONE strategy selected for " << vuln.TargetFileName
           << ":" << vuln.TargetFunctionName << "...\n";
       out << "[CVEAssert] Skipping remediation\n";
@@ -345,7 +345,7 @@ struct LabelCVEPass : public PassInfoMixin<LabelCVEPass> {
 
     for (auto &vuln : vulns) {
       // Also skip instrumentation for skipped vulnerabilities
-      if (vuln.Strategy == Vulnerability::RemediationStrategies::NONE) {
+      if (vuln.Strategy == RemediationStrategies::NONE) {
         continue;
       }
 
@@ -416,7 +416,7 @@ struct LabelCVEPass : public PassInfoMixin<LabelCVEPass> {
     std::vector<Vulnerability> moduleVulns;
 
     for (auto &vuln : vulnerabilities) {
-      if (vuln.Output == Vulnerability::RemediationOutput::PATCH) {
+      if (vuln.Output == RemediationOutput::PATCH) {
         patchVulns.push_back(vuln);
       } else {
         moduleVulns.push_back(vuln);

resolve-cveassert/src/Contract.hpp

@@ -0,0 +1,31 @@
+/*
+ *   Copyright (c) 2025 Riverside Research.
+ *   LGPL-3; See LICENSE.txt in the repo root for details.
+ */
+
+#pragma once
+
+#include "Remediation.hpp"
+#include "llvm/Support/JSON.h"
+
+#include <vector>
+
+enum class PredicateKind {
+  InBounds,
+  NotEqual,
+  NotNull,
+  NonZero,
+};
+
+// Predicates tell the compiler what
+// must be true before executing the operation
+struct Predicate {
+  PredicateKind kind;
+  unsigned arg0;
+  unsigned arg1;
+};
+
+struct Contract {
+  std::vector<Predicate> preconditions;
+  RemediationStrategies strategy;
+};