[DFSan] Add `zeroext` attribute for callbacks with 8bit shadow variable arguments

[tangyouling]

Summary

Add zeroext attribute for below callbacks' first parameter
(8bit shadow variable arguments) to conform to many platforms'
ABI calling convention and some compiler behavior.

• __dfsan_load_callback
• __dfsan_store_callback
• __dfsan_cmp_callback
• __dfsan_conditional_callback
• __dfsan_conditional_callback_origin
• __dfsan_reaches_function_callback
• __dfsan_reaches_function_callback_origin

The type of these callbacks' first parameter is u8 (see the
definition of dfsan_label). First, many platforms' ABI
requires unsigned integer data types (except unsigned int)
are zero-extended when stored in general-purpose register.
Second, the problem is that compiler optimization may assume
the arguments are zero-extended and, if not, misbehave, e.g.
it uses an i8 argument to index into a jump table. If the
argument has non-zero high bits, the output executable may
crash at run-time. So we need to add the zeroext attribute
when declaring and calling them.

[tangyouling]

add zeroext attribute patch with help from @SixWeining, thanks.

[browneee]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

299–306

Move this into D140690

All the other ZExt fixes in this file are unrelated to this.

1053–1055

I think the 8bit shadow variable arguments should be zext on any platform (this applies to all the ZExt changes in this file).

1129–1134

Move this into D140690

All the other ZExt fixes in this file are unrelated to this.

1435

Please use a separate AttributeList for each call, like in initializeRuntimeFunctions().

llvm/test/Instrumentation/DataFlowSanitizer/callback-loongarch64.ll

1

Move file to D140690.

[tangyouling]

Address @browneee comments.

[tangyouling]

Add some missing tests.

[MaskRay]

Similar to D76624 (SystemZ for msan).

[MaskRay]

Some platforms' ABI including LoongArch and RISCV requires unsigned integer data types (except unsigned int) are zero-extended when stored in general-purpose register. So we need to add the zeroext attribute when declaring and calling them.

This is not precise. First, most architectures zero-extend unsigned integer data types... So you change "some" to "many" and drop LoongArch/RISCV.
Second, the problem is that compiler optimization may assume the arguments are zero-extended and, if not, misbehave, e.g. it uses an i8 argument to index into a jump table. If the argument has non-zero high bits, the output executable may crash at run-time.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

 // NOLINTNEXTLINE(readability-identifier-naming)
 static const MemoryMapParams Linux_X86_64_MemoryMapParams = {
     0,              // AndMask (not used)
     0x500000000000, // XorMask
     0,              // ShadowBase (not used)
     0x100000000000, // OriginBase
 };
 
+// loongarch64 Linux
+// NOLINTNEXTLINE(readability-identifier-naming)
+static const MemoryMapParams Linux_LoongArch64_MemoryMapParams = {
+    0,              // AndMask (not used)
+    0x500000000000, // XorMask
+    0,              // ShadowBase (not used)
+    0x100000000000, // OriginBase
+};
+
 namespace {
 
 class DFSanABIList {
   std::unique_ptr<SpecialCaseList> SCL;
 
 public:
   DFSanABIList() = default;
 

 
 void DFSanFunction::addConditionalCallbacksIfEnabled(Instruction &I,
                                                      Value *Condition) {
   if (!ClConditionalCallbacks) {
     return;
   }
   IRBuilder<> IRB(&I);
   Value *CondShadow = getShadow(Condition);
+  CallInst *CI;
   if (DFS.shouldTrackOrigins()) {
     Value *CondOrigin = getOrigin(Condition);
-    IRB.CreateCall(DFS.DFSanConditionalCallbackOriginFn,
+    CI = IRB.CreateCall(DFS.DFSanConditionalCallbackOriginFn,
                         {CondShadow, CondOrigin});
   } else {
-    IRB.CreateCall(DFS.DFSanConditionalCallbackFn, {CondShadow});
+    CI = IRB.CreateCall(DFS.DFSanConditionalCallbackFn, {CondShadow});
   }
+  Triple TargetTriple(I.getModule()->getTargetTriple());
+  if (TargetTriple.getArch() == Triple::loongarch64)
+    CI->addParamAttr(0, Attribute::ZExt);
 }
 
 void DFSanFunction::addReachesFunctionCallbacksIfEnabled(IRBuilder<> &IRB,
                                                          Instruction &I,
                                                          Value *Data) {
   if (!ClReachesFunctionCallbacks) {
     return;
   }
void DFSanFunction::addReachesFunctionCallbacksIfEnabled(IRBuilder<> &IRB,
   if (DFS.shouldTrackOrigins()) {
     Value *DataOrigin = getOrigin(Data);
     args = { DataShadow, DataOrigin, FilePathPtr, CILine, FunctionNamePtr };
     CB = IRB.CreateCall(DFS.DFSanReachesFunctionCallbackOriginFn, args);
   } else {
     args = { DataShadow, FilePathPtr, CILine, FunctionNamePtr };
     CB = IRB.CreateCall(DFS.DFSanReachesFunctionCallbackFn, args);
   }
+  Triple TargetTriple(I.getModule()->getTargetTriple());
+  if (TargetTriple.getArch() == Triple::loongarch64)
+    CB->addParamAttr(0, Attribute::ZExt);
   CB->setDebugLoc(dbgloc);
 }
 
 Type *DataFlowSanitizer::getShadowTy(Type *OrigTy) {
   if (!OrigTy->isSized())
     return PrimitiveShadowTy;
   if (isa<IntegerType>(OrigTy))
     return PrimitiveShadowTy;

 }
 
 bool DataFlowSanitizer::initializeModule(Module &M) {
   Triple TargetTriple(M.getTargetTriple());
   const DataLayout &DL = M.getDataLayout();
 
   if (TargetTriple.getOS() != Triple::Linux)
     report_fatal_error("unsupported operating system");
-  if (TargetTriple.getArch() != Triple::x86_64)
-    report_fatal_error("unsupported architecture");
+  if (TargetTriple.getArch() == Triple::x86_64)
     MapParams = &Linux_X86_64_MemoryMapParams;
-
+  else if (TargetTriple.getArch() == Triple::loongarch64)
+    MapParams = &Linux_LoongArch64_MemoryMapParams;
+  else
+    report_fatal_error("unsupported architecture");
   Mod = &M;
   Ctx = &M.getContext();
   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);
DFSanRuntimeFunctions.insert(
       DFSanMemShadowOriginConditionalExchangeFn.getCallee()
           ->stripPointerCasts());
   DFSanRuntimeFunctions.insert(
       DFSanMaybeStoreOriginFn.getCallee()->stripPointerCasts());
 }
 
 // Initializes event callback functions and declare them in the module
 void DataFlowSanitizer::initializeCallbackFunctions(Module &M) {
-  DFSanLoadCallbackFn = Mod->getOrInsertFunction("__dfsan_load_callback",
-                                                 DFSanLoadStoreCallbackFnTy);
-  DFSanStoreCallbackFn = Mod->getOrInsertFunction("__dfsan_store_callback",
-                                                  DFSanLoadStoreCallbackFnTy);
+  Triple TargetTriple(M.getTargetTriple());
+  AttributeList AL;
+  if (TargetTriple.getArch() == Triple::loongarch64)
+    AL = AL.addParamAttribute(M.getContext(), 0, Attribute::ZExt);
+  DFSanLoadCallbackFn = Mod->getOrInsertFunction(
+      "__dfsan_load_callback", DFSanLoadStoreCallbackFnTy, AL);
+  DFSanStoreCallbackFn = Mod->getOrInsertFunction(
+      "__dfsan_store_callback", DFSanLoadStoreCallbackFnTy, AL);
   DFSanMemTransferCallbackFn = Mod->getOrInsertFunction(
       "__dfsan_mem_transfer_callback", DFSanMemTransferCallbackFnTy);
-  DFSanCmpCallbackFn =
-      Mod->getOrInsertFunction("__dfsan_cmp_callback", DFSanCmpCallbackFnTy);
+  DFSanCmpCallbackFn = Mod->getOrInsertFunction("__dfsan_cmp_callback",
+                                                DFSanCmpCallbackFnTy, AL);
 
   DFSanConditionalCallbackFn = Mod->getOrInsertFunction(
-      "__dfsan_conditional_callback", DFSanConditionalCallbackFnTy);
+      "__dfsan_conditional_callback", DFSanConditionalCallbackFnTy, AL);
   DFSanConditionalCallbackOriginFn =
       Mod->getOrInsertFunction("__dfsan_conditional_callback_origin",
-                               DFSanConditionalCallbackOriginFnTy);
-  DFSanReachesFunctionCallbackFn = Mod->getOrInsertFunction(
-      "__dfsan_reaches_function_callback", DFSanReachesFunctionCallbackFnTy);
+                               DFSanConditionalCallbackOriginFnTy, AL);
+  DFSanReachesFunctionCallbackFn =
+      Mod->getOrInsertFunction("__dfsan_reaches_function_callback",
+                               DFSanReachesFunctionCallbackFnTy, AL);
   DFSanReachesFunctionCallbackOriginFn =
       Mod->getOrInsertFunction("__dfsan_reaches_function_callback_origin",
-                               DFSanReachesFunctionCallbackOriginFnTy);
+                               DFSanReachesFunctionCallbackOriginFnTy, AL);
 }
 
 void DataFlowSanitizer::injectMetadataGlobals(Module &M) {
   // These variables can be used:
   // - by the runtime (to discover what the shadow width was, during
   //   compilation)
   // - in testing (to avoid hardcoding the shadow width and type but instead
   //   extract them by pattern matching)
void DFSanVisitor::visitLoadInst(LoadInst &LI) {
 
   if (ShouldTrackOrigins) {
     DFSF.setOrigin(&LI, DFSF.combineOrigins(Shadows, Origins, Pos));
   }
 
   if (ClEventCallbacks) {
     IRBuilder<> IRB(Pos);
     Value *Addr8 = IRB.CreateBitCast(LI.getPointerOperand(), DFSF.DFS.Int8Ptr);
+    CallInst *CI =
         IRB.CreateCall(DFSF.DFS.DFSanLoadCallbackFn, {PrimitiveShadow, Addr8});
+    Triple TargetTriple(LI.getModule()->getTargetTriple());
+    if (TargetTriple.getArch() == Triple::loongarch64)
+      CI->addParamAttr(0, Attribute::ZExt);
   }
 
   IRBuilder<> IRB(AfterLi);
   DFSF.addReachesFunctionCallbacksIfEnabled(IRB, LI, &LI);
 }
 
 Value *DFSanFunction::updateOriginIfTainted(Value *Shadow, Value *Origin,
                                             IRBuilder<> &IRB) {
void DFSanVisitor::visitStoreInst(StoreInst &SI) {
   Value *Origin = nullptr;
   if (ShouldTrackOrigins)
     Origin = DFSF.combineOrigins(Shadows, Origins, &SI);
   DFSF.storePrimitiveShadowOrigin(SI.getPointerOperand(), Size, SI.getAlign(),
                                   PrimitiveShadow, Origin, &SI);
   if (ClEventCallbacks) {
     IRBuilder<> IRB(&SI);
     Value *Addr8 = IRB.CreateBitCast(SI.getPointerOperand(), DFSF.DFS.Int8Ptr);
+    CallInst *CI =
         IRB.CreateCall(DFSF.DFS.DFSanStoreCallbackFn, {PrimitiveShadow, Addr8});
+    Triple TargetTriple(SI.getModule()->getTargetTriple());
+    if (TargetTriple.getArch() == Triple::loongarch64)
+      CI->addParamAttr(0, Attribute::ZExt);
   }
 }
 
 void DFSanVisitor::visitCASOrRMW(Align InstAlignment, Instruction &I) {
   assert(isa<AtomicRMWInst>(I) || isa<AtomicCmpXchgInst>(I));
 
   Value *Val = I.getOperand(1);
   const auto &DL = I.getModule()->getDataLayout();

 
 void DFSanVisitor::visitCastInst(CastInst &CI) { visitInstOperands(CI); }
 
 void DFSanVisitor::visitCmpInst(CmpInst &CI) {
   visitInstOperands(CI);
   if (ClEventCallbacks) {
     IRBuilder<> IRB(&CI);
     Value *CombinedShadow = DFSF.getShadow(&CI);
+    CallInst *CallI =
         IRB.CreateCall(DFSF.DFS.DFSanCmpCallbackFn, CombinedShadow);
+    Triple TargetTriple(CI.getModule()->getTargetTriple());
+    if (TargetTriple.getArch() == Triple::loongarch64)
+      CallI->addParamAttr(0, Attribute::ZExt);
   }
 }
 
 void DFSanVisitor::visitLandingPadInst(LandingPadInst &LPI) {
   // We do not need to track data through LandingPadInst.
   //
   // For the C++ exceptions, if a value is thrown, this value will be stored
   // in a memory location provided by __cxa_allocate_exception(...) (on the

llvm/test/Instrumentation/DataFlowSanitizer/callback-loongarch64.ll

+; RUN: opt < %s -passes=dfsan -dfsan-event-callbacks=1 -S | FileCheck %s
+target datalayout = "e-m:e-p:64:64-i64:64-i128:128-n64-S128"
+target triple = "loongarch64-unknown-linux-gnu"
+
+; CHECK: @__dfsan_shadow_width_bits = weak_odr constant i32 [[#SBITS:]]
+; CHECK: @__dfsan_shadow_width_bytes = weak_odr constant i32 [[#SBYTES:]]
+
+define i8 @load8(ptr %p) {
+  ; CHECK: call void @__dfsan_load_callback(i[[#SBITS]] zeroext %[[LABEL:.*]], ptr %p)
+  ; CHECK: %a = load i8, ptr %p
+  ; CHECK: store i[[#SBITS]] %[[LABEL]], ptr @__dfsan_retval_tls
+
+  %a = load i8, ptr %p
+  ret i8 %a
+}
+
+define void @store8(ptr %p, i8 %a) {
+  ; CHECK: store i[[#SBITS]] %[[LABEL:.*]], ptr %{{.*}}
+  ; CHECK: call void @__dfsan_store_callback(i[[#SBITS]] zeroext %[[LABEL]], ptr %p)
+  ; CHECK: store i8 %a, ptr %p
+
+  store i8 %a, ptr %p
+  ret void
+}
+
+define i1 @cmp(i8 %a, i8 %b) {
+  ; CHECK: call void @__dfsan_cmp_callback(i[[#SBITS]] zeroext %[[CMPLABEL:.*]])
+  ; CHECK: %c = icmp ne i8 %a, %b
+  ; CHECK: store i[[#SBITS]] %[[CMPLABEL]], ptr @__dfsan_retval_tls
+
+  %c = icmp ne i8 %a, %b
+  ret i1 %c
+}
+
+; CHECK: declare void @__dfsan_load_callback(i[[#SBITS]] zeroext, ptr)
+; CHECK: declare void @__dfsan_store_callback(i[[#SBITS]] zeroext, ptr)
+; CHECK: declare void @__dfsan_mem_transfer_callback(ptr, i64)
+; CHECK: declare void @__dfsan_cmp_callback(i[[#SBITS]] zeroext)
+; CHECK: declare void @__dfsan_conditional_callback(i[[#SBITS]] zeroext)
+; CHECK: declare void @__dfsan_conditional_callback_origin(i[[#SBITS]] zeroext, i32)
+; CHECK: declare void @__dfsan_reaches_function_callback(i[[#SBITS]] zeroext, ptr, i32, ptr)
+; CHECK: declare void @__dfsan_reaches_function_callback_origin(i[[#SBITS]] zeroext, i32, ptr, i32, ptr)