[DFSan] Add dfsan-combine-taint-lookup-table option as work around for false negatives when dfsan-combine-pointer-labels-on-load=0 and dfsan-combine-offset-labels-on-gep=0 miss data flows through lookup tables.

[browneee]

Summary

Example case:
https://github.com/abseil/abseil-cpp/blob/628a2825f8dc0219964886e7cc3f7f519e3bd950/absl/strings/ascii.h#L182

[browneee]

Rebase.

[browneee]

Rebase

[vitalybuka]

description title is unusually long

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

803

drop {}
and the rest of the patch

1850

in llvm code we usually omit {} on simple cases like this

2111

V->stripPointerCasts() ?

[browneee]

Fix extra braces.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

2111

This is different - it traverses through non-zero GEPs.

compiler-rt/test/dfsan/lookup_table.c

+// RUN: %clang_dfsan %s -mllvm -dfsan-combine-offset-labels-on-gep=false -mllvm -dfsan-combine-pointer-labels-on-load=false -mllvm -dfsan-combine-taint-lookup-table=remap_to_upper -DLOOKUP_TABLE -o %t && %run %t
+// RUN: %clang_dfsan %s -mllvm -dfsan-combine-offset-labels-on-gep=false -mllvm -dfsan-combine-pointer-labels-on-load=false -mllvm -dfsan-combine-taint-lookup-table=no_match -o %t && %run %t
+// RUN: %clang_dfsan %s -mllvm -dfsan-combine-offset-labels-on-gep=false -mllvm -dfsan-combine-pointer-labels-on-load=false -o %t && %run %t
+//
+// REQUIRES: x86_64-target-arch
+
+#include <sanitizer/dfsan_interface.h>
+#include <assert.h>
+
+const char remap_to_upper[256] = {
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', 'A', 'B',
+    'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K',
+    'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T',
+    'U', 'V', 'W', 'X', 'Y', 'Z', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.', '.', '.', '.', '.', '.',
+    '.', '.', '.', '.',
+};
+
+char character_mapping(unsigned char c) {
+  return remap_to_upper[c];
+}
+
+int main(void) {
+  char a = 'b';
+  dfsan_label i_label = 1;
+  dfsan_set_label(i_label, &a, sizeof(a));
+  assert(dfsan_read_label(&a, sizeof(a)) == i_label);
+
+  char b = character_mapping(a);
+  assert(b == 'B');
+
+#ifdef LOOKUP_TABLE
+  assert(dfsan_read_label(&b, sizeof(b)) == i_label);
+#else
+  assert(dfsan_read_label(&b, sizeof(b)) == 0);
+#endif
+  return 0;
+}

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

 #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"

 // Controls whether the pass propagates labels of offsets in GEP instructions.
 static cl::opt<bool> ClCombineOffsetLabelsOnGEP(
     "dfsan-combine-offset-labels-on-gep",
     cl::desc(
         "Combine the label of the offset with the label of the pointer when "
         "doing pointer arithmetic."),
     cl::Hidden, cl::init(true));
 
+static cl::list<std::string> ClCombineTaintLookupTables(
+    "dfsan-combine-taint-lookup-table",
+    cl::desc(
+        "When dfsan-combine-offset-labels-on-gep and/or "
+        "dfsan-combine-pointer-labels-on-load are false, this flag can "
+        "be used to re-enable combining offset and/or pointer taint when "
+        "loading specific constant global variables (i.e. lookup tables)."),
+    cl::Hidden);
+
 static cl::opt<bool> ClDebugNonzeroLabels(
     "dfsan-debug-nonzero-labels",
     cl::desc("Insert calls to __dfsan_nonzero_label on observing a parameter, "
              "load or return with a nonzero label"),
     cl::Hidden);
 
 // Experimental feature that inserts callbacks for certain data events.
 // Currently callbacks are only inserted for loads, stores, memory transfers
class DataFlowSanitizer {
   FunctionCallee DFSanMemOriginTransferFn;
   FunctionCallee DFSanMaybeStoreOriginFn;
   SmallPtrSet<Value *, 16> DFSanRuntimeFunctions;
   MDNode *ColdCallWeights;
   MDNode *OriginStoreWeights;
   DFSanABIList ABIList;
   DenseMap<Value *, Function *> UnwrappedFnMap;
   AttributeMask ReadOnlyNoneAttrs;
+  StringSet<> CombineTaintLookupTableNames;
 
   /// Memory map parameters used in calculation mapping application addresses
   /// to shadow addresses and origin addresses.
   const MemoryMapParams *MapParams;
 
   Value *getShadowOffset(Value *Addr, IRBuilder<> &IRB);
   Value *getShadowAddress(Value *Addr, Instruction *Pos);
   Value *getShadowAddress(Value *Addr, Instruction *Pos, Value *ShadowOffset);
void storeZeroPrimitiveShadow(Value *Addr, uint64_t Size, Align ShadowAlign,
                                 Instruction *Pos);
 
   Align getShadowAlign(Align InstAlignment);
 
   // If ClConditionalCallbacks is enabled, insert a callback after a given
   // branch instruction using the given conditional expression.
   void addConditionalCallbacksIfEnabled(Instruction &I, Value *Condition);
 
+  bool isLookupTableConstant(Value *P);
+
 private:
   /// Collapses the shadow with aggregate type into a single primitive shadow
   /// value.
   template <class AggregateType>
   Value *collapseAggregateShadow(AggregateType *AT, Value *Shadow,
                                  IRBuilder<> &IRB);
 
   Value *collapseToPrimitiveShadow(Value *Shadow, IRBuilder<> &IRB);

 
 DataFlowSanitizer::DataFlowSanitizer(
     const std::vector<std::string> &ABIListFiles) {
   std::vector<std::string> AllABIListFiles(std::move(ABIListFiles));
   llvm::append_range(AllABIListFiles, ClABIListFiles);
   // FIXME: should we propagate vfs::FileSystem to this constructor?
   ABIList.set(
       SpecialCaseList::createOrDie(AllABIListFiles, *vfs::getRealFileSystem()));
+
+  for (StringRef v : ClCombineTaintLookupTables)
+    CombineTaintLookupTableNames.insert(v);
 }
 
 TransformedFunction DataFlowSanitizer::getCustomFunctionType(FunctionType *T) {
   SmallVector<Type *, 4> ArgTypes;
 
   // Some parameters of the custom function being constructed are
   // parameters of T.  Record the mapping from parameters of T to
   // parameters of the custom function, so that parameter attributes
Align DFSanFunction::getShadowAlign(Align InstAlignment) {
   return Align(Alignment.value() * DFS.ShadowWidthBytes);
 }
 
 Align DFSanFunction::getOriginAlign(Align InstAlignment) {
   const Align Alignment = llvm::assumeAligned(InstAlignment.value());
   return Align(std::max(MinOriginAlignment, Alignment));
 }
 
+bool DFSanFunction::isLookupTableConstant(Value *P) {
+  if (GlobalVariable *GV = dyn_cast<GlobalVariable>(P->stripPointerCasts()))
+    if (GV->isConstant() && GV->hasName())
+      return DFS.CombineTaintLookupTableNames.count(GV->getName());
+
+  return false;
+}
+
 bool DFSanFunction::useCallbackLoadLabelAndOrigin(uint64_t Size,
                                                   Align InstAlignment) {
   // When enabling tracking load instructions, we always use
   // __dfsan_load_label_and_origin to reduce code size.
   if (ClTrackOrigins == 2)
     return true;
 
   assert(Size != 0);
static AtomicOrdering addAcquireOrdering(AtomicOrdering AO) {
   case AtomicOrdering::AcquireRelease:
     return AtomicOrdering::AcquireRelease;
   case AtomicOrdering::SequentiallyConsistent:
     return AtomicOrdering::SequentiallyConsistent;
   }
   llvm_unreachable("Unknown ordering");
 }
 
+Value *StripPointerGEPsAndCasts(Value *V) {
+  if (!V->getType()->isPointerTy())
+    return V;
+
+  // DFSan pass should be running on valid IR, but we'll
+  // keep a seen set to ensure there are no issues.
+  SmallPtrSet<const Value *, 4> Visited;
+  Visited.insert(V);
+  do {
+    if (auto *GEP = dyn_cast<GEPOperator>(V)) {
+      V = GEP->getPointerOperand();
+    } else if (Operator::getOpcode(V) == Instruction::BitCast) {
+      V = cast<Operator>(V)->getOperand(0);
+      if (!V->getType()->isPointerTy())
+        return V;
+    } else if (isa<GlobalAlias>(V)) {
+      V = cast<GlobalAlias>(V)->getAliasee();
+    }
+  } while (Visited.insert(V).second);
+
+  return V;
+}
+
 void DFSanVisitor::visitLoadInst(LoadInst &LI) {
   auto &DL = LI.getModule()->getDataLayout();
   uint64_t Size = DL.getTypeStoreSize(LI.getType());
   if (Size == 0) {
     DFSF.setShadow(&LI, DFSF.DFS.getZeroShadow(&LI));
     DFSF.setOrigin(&LI, DFSF.DFS.ZeroOrigin);
     return;
   }
void DFSanVisitor::visitLoadInst(LoadInst &LI) {
   Value *PrimitiveShadow, *Origin;
   std::tie(PrimitiveShadow, Origin) =
       DFSF.loadShadowOrigin(LI.getPointerOperand(), Size, LI.getAlign(), Pos);
   const bool ShouldTrackOrigins = DFSF.DFS.shouldTrackOrigins();
   if (ShouldTrackOrigins) {
     Shadows.push_back(PrimitiveShadow);
     Origins.push_back(Origin);
   }
-  if (ClCombinePointerLabelsOnLoad) {
+  if (ClCombinePointerLabelsOnLoad ||
+      DFSF.isLookupTableConstant(
+          StripPointerGEPsAndCasts(LI.getPointerOperand()))) {
     Value *PtrShadow = DFSF.getShadow(LI.getPointerOperand());
     PrimitiveShadow = DFSF.combineShadows(PrimitiveShadow, PtrShadow, Pos);
     if (ShouldTrackOrigins) {
       Shadows.push_back(PtrShadow);
       Origins.push_back(DFSF.getOrigin(LI.getPointerOperand()));
     }
   }
   if (!DFSF.DFS.isZeroShadow(PrimitiveShadow))
void DFSanVisitor::visitLandingPadInst(LandingPadInst &LPI) {
   //
   // The second element in the pair result of the LandingPadInst is a
   // register value, but it is for a type ID and should never be tainted.
   DFSF.setShadow(&LPI, DFSF.DFS.getZeroShadow(&LPI));
   DFSF.setOrigin(&LPI, DFSF.DFS.ZeroOrigin);
 }
 
 void DFSanVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
-  if (ClCombineOffsetLabelsOnGEP) {
+  if (ClCombineOffsetLabelsOnGEP ||
+      DFSF.isLookupTableConstant(
+          StripPointerGEPsAndCasts(GEPI.getPointerOperand()))) {
     visitInstOperands(GEPI);
     return;
   }
 
   // Only propagate shadow/origin of base pointer value but ignore those of
   // offset operands.
   Value *BasePointer = GEPI.getPointerOperand();
   DFSF.setShadow(&GEPI, DFSF.getShadow(BasePointer));

llvm/test/Instrumentation/DataFlowSanitizer/lookup_table.ll

+; RUN: opt < %s -dfsan -dfsan-combine-pointer-labels-on-load=false -dfsan-combine-offset-labels-on-gep=false -dfsan-combine-taint-lookup-table=lookup_table_a -S | FileCheck %s --check-prefixes=CHECK,LOOKUP_A
+; RUN: opt < %s -dfsan -dfsan-combine-pointer-labels-on-load=false -dfsan-combine-offset-labels-on-gep=false -S | FileCheck %s --check-prefixes=CHECK,NO_LOOKUP_A
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK: @__dfsan_arg_tls = external thread_local(initialexec) global [[TLS_ARR:\[100 x i64\]]]
+; CHECK: @__dfsan_retval_tls = external thread_local(initialexec) global [[TLS_ARR]]
+; CHECK: @__dfsan_shadow_width_bits = weak_odr constant i32 [[#SBITS:]]
+; CHECK: @__dfsan_shadow_width_bytes = weak_odr constant i32 [[#SBYTES:]]
+
+@lookup_table_a = external local_unnamed_addr constant [256 x i8], align 16
+@lookup_table_b = external local_unnamed_addr constant [256 x i8], align 16
+
+define i8 @load_lookup_table_a(i8 %p) {
+  ; CHECK-LABEL:           @load_lookup_table_a.dfsan
+  ; CHECK-NEXT:            %[[#PS:]] = load i[[#SBITS]], i[[#SBITS]]* bitcast ([100 x i64]* @__dfsan_arg_tls to i[[#SBITS]]*), align [[ALIGN:2]]
+  ; CHECK-NEXT:            %c = zext i8 %p to i64
+  ; CHECK-NEXT:            %b = getelementptr inbounds [256 x i8], [256 x i8]* @lookup_table_a, i64 0, i64 %c
+  ; CHECK-NEXT:            %a = load i8, i8* %b, align 1
+  ; Propagates p shadow when lookup_table_a flag is provided, otherwise propagates 0 shadow
+  ; LOOKUP_A-NEXT:         store i[[#SBITS]] %[[#PS]], i[[#SBITS]]* bitcast ([100 x i64]* @__dfsan_retval_tls to i[[#SBITS]]*), align [[ALIGN]]
+  ; NO_LOOKUP_A-NEXT:      store i[[#SBITS]] 0, i[[#SBITS]]* bitcast ([100 x i64]* @__dfsan_retval_tls to i[[#SBITS]]*), align [[ALIGN]]
+  ; CHECK-NEXT:            ret i8 %a
+
+  %c = zext i8 %p to i64
+  %b = getelementptr inbounds [256 x i8], [256 x i8]* @lookup_table_a, i64 0, i64 %c
+  %a = load i8, i8* %b
+  ret i8 %a
+}
+
+define i8 @load_lookup_table_b(i8 %p) {
+  ; CHECK-LABEL:           @load_lookup_table_b.dfsan
+  ; CHECK-NEXT:            %[[#PS:]] = load i[[#SBITS]], i[[#SBITS]]* bitcast ([100 x i64]* @__dfsan_arg_tls to i[[#SBITS]]*), align 2
+  ; CHECK-NEXT:            %c = zext i8 %p to i64
+  ; CHECK-NEXT:            %b = getelementptr inbounds [256 x i8], [256 x i8]* @lookup_table_b, i64 0, i64 %c
+  ; CHECK-NEXT:            %a = load i8, i8* %b, align 1
+  ; Propagates 0 shadow
+  ; CHECK-NEXT:            store i[[#SBITS]] 0, i[[#SBITS]]* bitcast ([100 x i64]* @__dfsan_retval_tls to i[[#SBITS]]*), align [[ALIGN]]
+  ; CHECK-NEXT:            ret i8 %a
+
+  %c = zext i8 %p to i64
+  %b = getelementptr inbounds [256 x i8], [256 x i8]* @lookup_table_b, i64 0, i64 %c
+  %a = load i8, i8* %b, align 1
+  ret i8 %a
+}