[dfsan] Add origin address calculation

[stephan.yichao.zhao]

Summary

This is a part of https://reviews.llvm.org/D95835.

[morehouse]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1554

We can avoid this AND operation entirely by including it in the shadow mask.

[stephan.yichao.zhao]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

1554

Yeah. ~0x700...0003 works. It seems that the problem is we wanted to get both shadow and origin addresses. So

1. with the current approach,

offset = addr & ~0x700...00
shadow = offset x 2
origin = offset + origin_base
// When align is < 4, we may need one more origin & 4

2. if we use ~0x700...03

offset_s = addr & ~0x700...00
shadow = offset_s x 2
offset_o = addr & ~0x700...03
origin = offset_o + origin_base

So the case 1 actually uses 1 less IR instruction in most cases.

[morehouse]

LGTM

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

 /// hold the label.  On Linux/x86_64, memory is laid out as follows:
 ///
 /// +--------------------+ 0x800000000000 (top of memory)
 /// | application memory |
 /// +--------------------+ 0x700000008000 (kAppAddr)
 /// |                    |
 /// |       unused       |
 /// |                    |
-/// +--------------------+ 0x200200000000 (kUnusedAddr)
+/// +--------------------+ 0x300200000000 (kUnusedAddr)
 /// |    union table     |
-/// +--------------------+ 0x200000000000 (kUnionTableAddr)
+/// +--------------------+ 0x300000000000 (kUnionTableAddr)
+/// |       origin       |
+/// +--------------------+ 0x200000008000 (kOriginAddr)
 /// |   shadow memory    |
 /// +--------------------+ 0x000000010000 (kShadowAddr)
 /// | reserved by kernel |
 /// +--------------------+ 0x000000000000
 ///
 /// To derive a shadow memory address from an application memory address,
 /// bits 44-46 are cleared to bring the address into the range
 /// [0x000000008000,0x100000000000).  Then the address is shifted left by 1 to

 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
 // This must be consistent with ShadowWidthBits.
 static const Align kShadowTLSAlignment = Align(2);
 
+static const Align kMinOriginAlignment = Align(4);
+
 // The size of TLS variables. These constants must be kept in sync with the ones
 // in dfsan.cpp.
 static const unsigned kArgTLSSize = 800;
 static const unsigned kRetvalTLSSize = 800;
 
 // External symbol to be used when generating the shadow address for
 // architectures with multiple VMAs. Instead of using a constant integer
 // the runtime will set the external mask based on the VMA range.
enum WrapperKind {
     WK_Custom
   };
 
   Module *Mod;
   LLVMContext *Ctx;
   Type *Int8Ptr;
   IntegerType *OriginTy;
   PointerType *OriginPtrTy;
+  ConstantInt *OriginBase;
   /// The shadow type for all primitive types and vector types.
   IntegerType *PrimitiveShadowTy;
   PointerType *PrimitiveShadowPtrTy;
   IntegerType *IntptrTy;
   ConstantInt *ZeroPrimitiveShadow;
   ConstantInt *ShadowPtrMask;
   ConstantInt *ShadowPtrMul;
   Constant *ArgTLS;
class DataFlowSanitizer {
   FunctionCallee DFSanMaybeStoreOriginFn;
   SmallPtrSet<Value *, 16> DFSanRuntimeFunctions;
   MDNode *ColdCallWeights;
   DFSanABIList ABIList;
   DenseMap<Value *, Function *> UnwrappedFnMap;
   AttrBuilder ReadOnlyNoneAttrs;
   bool DFSanRuntimeShadowMask = false;
 
+  Value *getShadowOffset(Value *Addr, IRBuilder<> &IRB);
   Value *getShadowAddress(Value *Addr, Instruction *Pos);
+  std::pair<Value *, Value *>
+  getShadowOriginAddress(Value *Addr, Align InstAlignment, Instruction *Pos);
   bool isInstrumented(const Function *F);
   bool isInstrumented(const GlobalAlias *GA);
   FunctionType *getArgsFunctionType(FunctionType *T);
   FunctionType *getTrampolineFunctionType(FunctionType *T);
   TransformedFunction getCustomFunctionType(FunctionType *T);
   InstrumentedABI getInstrumentedABI();
   WrapperKind getWrapperKind(Function *F);
   void addGlobalNamePrefix(GlobalValue *GV);
bool DataFlowSanitizer::init(Module &M) {
   Int8Ptr = Type::getInt8PtrTy(*Ctx);
   OriginTy = IntegerType::get(*Ctx, OriginWidthBits);
   OriginPtrTy = PointerType::getUnqual(OriginTy);
   PrimitiveShadowTy = IntegerType::get(*Ctx, ShadowWidthBits);
   PrimitiveShadowPtrTy = PointerType::getUnqual(PrimitiveShadowTy);
   IntptrTy = DL.getIntPtrType(*Ctx);
   ZeroPrimitiveShadow = ConstantInt::getSigned(PrimitiveShadowTy, 0);
   ShadowPtrMul = ConstantInt::getSigned(IntptrTy, ShadowWidthBytes);
+  OriginBase = ConstantInt::get(IntptrTy, 0x200000000000LL);
   if (IsX86_64)
     ShadowPtrMask = ConstantInt::getSigned(IntptrTy, ~0x700000000000LL);
   else if (IsMIPS64)
     ShadowPtrMask = ConstantInt::getSigned(IntptrTy, ~0xF000000000LL);
   // AArch64 supports multiple VMAs and the shadow mask is set at runtime.
   else if (IsAArch64)
     DFSanRuntimeShadowMask = true;
   else

 
 void DFSanFunction::setShadow(Instruction *I, Value *Shadow) {
   assert(!ValShadowMap.count(I));
   assert(DFS.shouldTrackFieldsAndIndices() ||
          Shadow->getType() == DFS.PrimitiveShadowTy);
   ValShadowMap[I] = Shadow;
 }
 
-Value *DataFlowSanitizer::getShadowAddress(Value *Addr, Instruction *Pos) {
+Value *DataFlowSanitizer::getShadowOffset(Value *Addr, IRBuilder<> &IRB) {
+  // Returns Addr & shadow_mask
   assert(Addr != RetvalTLS && "Reinstrumenting?");
-  IRBuilder<> IRB(Pos);
   Value *ShadowPtrMaskValue;
   if (DFSanRuntimeShadowMask)
     ShadowPtrMaskValue = IRB.CreateLoad(IntptrTy, ExternalShadowMask);
   else
     ShadowPtrMaskValue = ShadowPtrMask;
-  return IRB.CreateIntToPtr(
-      IRB.CreateMul(
-          IRB.CreateAnd(IRB.CreatePtrToInt(Addr, IntptrTy),
-                        IRB.CreatePtrToInt(ShadowPtrMaskValue, IntptrTy)),
-          ShadowPtrMul),
+  return IRB.CreateAnd(IRB.CreatePtrToInt(Addr, IntptrTy),
+                       IRB.CreatePtrToInt(ShadowPtrMaskValue, IntptrTy));
+}
+
+std::pair<Value *, Value *>
+DataFlowSanitizer::getShadowOriginAddress(Value *Addr, Align InstAlignment,
+                                          Instruction *Pos) {
+  // Returns ((Addr & shadow_mask) + origin_base) & ~4UL
+  IRBuilder<> IRB(Pos);
+  Value *ShadowOffset = getShadowOffset(Addr, IRB);
+  Value *ShadowPtr = IRB.CreateIntToPtr(
+      IRB.CreateMul(ShadowOffset, ShadowPtrMul), PrimitiveShadowPtrTy);
+  Value *OriginPtr = nullptr;
+  if (shouldTrackOrigins()) {
+    Value *OriginLong = IRB.CreateAdd(ShadowOffset, OriginBase);
+    const Align Alignment = llvm::assumeAligned(InstAlignment.value());
+    // When alignment is >= 4, Addr must be aligned to 4, otherwise it is UB.
+    // So Mask is unnecessary.
+    if (Alignment < kMinOriginAlignment) {
+      uint64_t Mask = kMinOriginAlignment.value() - 1;
+      OriginLong = IRB.CreateAnd(OriginLong, ConstantInt::get(IntptrTy, ~Mask));
+    }
+    OriginPtr = IRB.CreateIntToPtr(OriginLong, OriginPtrTy);
+  }
+  return {ShadowPtr, OriginPtr};
+}
+
+Value *DataFlowSanitizer::getShadowAddress(Value *Addr, Instruction *Pos) {
+  // Returns (Addr & shadow_mask) x 2
+  IRBuilder<> IRB(Pos);
+  Value *ShadowOffset = getShadowOffset(Addr, IRB);
+  return IRB.CreateIntToPtr(IRB.CreateMul(ShadowOffset, ShadowPtrMul),
                             PrimitiveShadowPtrTy);
 }
 
 Value *DFSanFunction::combineShadowsThenConvert(Type *T, Value *V1, Value *V2,
                                                 Instruction *Pos) {
   Value *PrimitiveValue = combineShadows(V1, V2, Pos);
   return expandFromPrimitiveShadow(T, PrimitiveValue, Pos);
 }