[analyzer] Fix taint propagation by remembering to the location context

Fixes the issue D118987 by mapping the propagation to the callsite's
LocationContext.
This way we can keep track of the in-flight propagations.

Note that empty propagation sets won't be inserted.

Reviewed By: NoQ, Szelethus

Differential Revision: https://reviews.llvm.org/D119128

clang/lib/StaticAnalyzer/Checkers/GenericTaintChecker.cpp

@@ -38,6 +38,8 @@ using namespace clang;
 using namespace ento;
 using namespace taint;
 
+using llvm::ImmutableSet;
+
 namespace {
 
 class GenericTaintChecker;
@@ -434,7 +436,9 @@ template <> struct ScalarEnumerationTraits<TaintConfiguration::VariadicType> {
 /// to the call post-visit. The values are signed integers, which are either
 /// ReturnValueIndex, or indexes of the pointer/reference argument, which
 /// points to data, which should be tainted on return.
-REGISTER_SET_WITH_PROGRAMSTATE(TaintArgsOnPostVisit, ArgIdxTy)
+REGISTER_MAP_WITH_PROGRAMSTATE(TaintArgsOnPostVisit, const LocationContext *,
+                               ImmutableSet<ArgIdxTy>)
+REGISTER_SET_FACTORY_WITH_PROGRAMSTATE(ArgIdxFactory, ArgIdxTy)
 
 void GenericTaintRuleParser::validateArgVector(const std::string &Option,
                                                const ArgVecTy &Args) const {
@@ -685,22 +689,26 @@ void GenericTaintChecker::checkPostCall(const CallEvent &Call,
   // Set the marked values as tainted. The return value only accessible from
   // checkPostStmt.
   ProgramStateRef State = C.getState();
+  const StackFrameContext *CurrentFrame = C.getStackFrame();
 
   // Depending on what was tainted at pre-visit, we determined a set of
   // arguments which should be tainted after the function returns. These are
   // stored in the state as TaintArgsOnPostVisit set.
-  TaintArgsOnPostVisitTy TaintArgs = State->get<TaintArgsOnPostVisit>();
-  if (TaintArgs.isEmpty())
+  TaintArgsOnPostVisitTy TaintArgsMap = State->get<TaintArgsOnPostVisit>();
+
+  const ImmutableSet<ArgIdxTy> *TaintArgs = TaintArgsMap.lookup(CurrentFrame);
+  if (!TaintArgs)
     return;
+  assert(!TaintArgs->isEmpty());
 
   LLVM_DEBUG(for (ArgIdxTy I
-                  : TaintArgs) {
+                  : *TaintArgs) {
     llvm::dbgs() << "PostCall<";
     Call.dump(llvm::dbgs());
     llvm::dbgs() << "> actually wants to taint arg index: " << I << '\n';
   });
 
-  for (ArgIdxTy ArgNum : TaintArgs) {
+  for (ArgIdxTy ArgNum : *TaintArgs) {
     // Special handling for the tainted return value.
     if (ArgNum == ReturnValueIndex) {
       State = addTaint(State, Call.getReturnValue());
@@ -714,7 +722,7 @@ void GenericTaintChecker::checkPostCall(const CallEvent &Call,
   }
 
   // Clear up the taint info from the state.
-  State = State->remove<TaintArgsOnPostVisit>();
+  State = State->remove<TaintArgsOnPostVisit>(CurrentFrame);
   C.addTransition(State);
 }
 
@@ -776,28 +784,33 @@ void GenericTaintRule::process(const GenericTaintChecker &Checker,
   };
 
   /// Propagate taint where it is necessary.
+  auto &F = State->getStateManager().get_context<ArgIdxFactory>();
+  ImmutableSet<ArgIdxTy> Result = F.getEmptySet();
   ForEachCallArg(
-      [this, &State, WouldEscape, &Call](ArgIdxTy I, const Expr *E, SVal V) {
+      [this, WouldEscape, &Call, &Result, &F](ArgIdxTy I, const Expr *E,
+                                              SVal V) {
         if (PropDstArgs.contains(I)) {
           LLVM_DEBUG(llvm::dbgs() << "PreCall<"; Call.dump(llvm::dbgs());
                      llvm::dbgs()
                      << "> prepares tainting arg index: " << I << '\n';);
-          State = State->add<TaintArgsOnPostVisit>(I);
+          Result = F.add(Result, I);
         }
 
         // TODO: We should traverse all reachable memory regions via the
         // escaping parameter. Instead of doing that we simply mark only the
         // referred memory region as tainted.
         if (WouldEscape(V, E->getType())) {
-          LLVM_DEBUG(if (!State->contains<TaintArgsOnPostVisit>(I)) {
+          LLVM_DEBUG(if (!Result.contains(I)) {
             llvm::dbgs() << "PreCall<";
             Call.dump(llvm::dbgs());
             llvm::dbgs() << "> prepares tainting arg index: " << I << '\n';
           });
-          State = State->add<TaintArgsOnPostVisit>(I);
+          Result = F.add(Result, I);
         }
       });
 
+  if (!Result.isEmpty())
+    State = State->set<TaintArgsOnPostVisit>(C.getStackFrame(), Result);
   C.addTransition(State);
 }
 
@@ -888,7 +901,11 @@ void GenericTaintChecker::taintUnsafeSocketProtocol(const CallEvent &Call,
   if (SafeProtocol)
     return;
 
-  C.addTransition(C.getState()->add<TaintArgsOnPostVisit>(ReturnValueIndex));
+  ProgramStateRef State = C.getState();
+  auto &F = State->getStateManager().get_context<ArgIdxFactory>();
+  ImmutableSet<ArgIdxTy> Result = F.add(F.getEmptySet(), ReturnValueIndex);
+  State = State->set<TaintArgsOnPostVisit>(C.getStackFrame(), Result);
+  C.addTransition(State);
 }
 
 /// Checker registration

clang/test/Analysis/taint-checker-callback-order-has-definition.c

@@ -3,9 +3,6 @@
 // RUN:   -mllvm -debug-only=taint-checker \
 // RUN:   2>&1 | FileCheck %s
 
-// FIXME: We should not crash.
-// XFAIL: *
-
 struct _IO_FILE;
 typedef struct _IO_FILE FILE;
 FILE *fopen(const char *fname, const char *mode);
@@ -31,12 +28,8 @@ void top(const char *fname, char *buf) {
   // CHECK-NEXT: PreCall<fgets(buf, 42, fp)> prepares tainting arg index: 1
   // CHECK-NEXT: PreCall<fgets(buf, 42, fp)> prepares tainting arg index: 2
 
-  // FIXME: We should propagate taint from PreCall<fgets> -> PostCall<fgets>.
-  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: -1
-  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: 0
-  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: 1
-  // CHECK-NEXT: PostCall<nested_call()> actually wants to taint arg index: 2
-
-  // FIXME: We should not crash.
-  // CHECK: PLEASE submit a bug report
+  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: -1
+  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: 0
+  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: 1
+  // CHECK-NEXT: PostCall<fgets(buf, 42, fp)> actually wants to taint arg index: 2
 }