hwasan: Pad arrays with non-1 size correctly.

[pcc]

Summary

Spotted by eugenis.

[pcc]

• Add test

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

bool HWAddressSanitizer::sanitizeFunction(Function &F) {
   // uninteresting allocas from hiding in instrumented alloca's padding and so
   // that we have enough space to store real tags for short granules.
   DenseMap<AllocaInst *, AllocaInst *> AllocaToPaddedAllocaMap;
   for (AllocaInst *AI : AllocasToInstrument) {
     uint64_t Size = getAllocaSizeInBytes(*AI);
     uint64_t AlignedSize = alignTo(Size, Mapping.getAllocaAlignment());
     AI->setAlignment(std::max(AI->getAlignment(), 16u));
     if (Size != AlignedSize) {
+      Type *AllocatedType = AI->getAllocatedType();
+      if (AI->isArrayAllocation()) {
+        uint64_t ArraySize =
+            cast<ConstantInt>(AI->getArraySize())->getZExtValue();
+        AllocatedType = ArrayType::get(AllocatedType, ArraySize);
+      }
       Type *TypeWithPadding = StructType::get(
-          AI->getAllocatedType(), ArrayType::get(Int8Ty, AlignedSize - Size));
+          AllocatedType, ArrayType::get(Int8Ty, AlignedSize - Size));
       auto *NewAI = new AllocaInst(
           TypeWithPadding, AI->getType()->getAddressSpace(), nullptr, "", AI);
       NewAI->takeName(AI);
       NewAI->setAlignment(AI->getAlignment());
       NewAI->setUsedWithInAlloca(AI->isUsedWithInAlloca());
       NewAI->setSwiftError(AI->isSwiftError());
       NewAI->copyMetadata(*AI);
-      Value *Zero = ConstantInt::get(Int32Ty, 0);
-      auto *GEP = GetElementPtrInst::Create(TypeWithPadding, NewAI,
-                                            {Zero, Zero}, "", AI);
-      AI->replaceAllUsesWith(GEP);
+      auto *Bitcast = new BitCastInst(NewAI, AI->getType(), "", AI);
+      AI->replaceAllUsesWith(Bitcast);
       AllocaToPaddedAllocaMap[AI] = NewAI;
     }
   }
 
   if (!AllocaToPaddedAllocaMap.empty()) {
     for (auto &BB : F)
       for (auto &Inst : BB)
         if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst))

llvm/test/Instrumentation/HWAddressSanitizer/alloca-with-calls.ll

 ; Test alloca instrumentation when tags are generated by HWASan function.
 ;
 ; RUN: opt < %s -hwasan -hwasan-generate-tags-with-calls -S | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-android"
 
 declare void @use32(i32*)
 
 define void @test_alloca() sanitize_hwaddress {
 ; CHECK-LABEL: @test_alloca(
-; CHECK: %[[GEP:[^ ]*]] = getelementptr { i32, [12 x i8] }, { i32, [12 x i8] }* %x, i32 0, i32 0
+; CHECK: %[[BC:[^ ]*]] = bitcast { i32, [12 x i8] }* %x to i32*
 ; CHECK: %[[T1:[^ ]*]] = call i8 @__hwasan_generate_tag()
 ; CHECK: %[[A:[^ ]*]] = zext i8 %[[T1]] to i64
-; CHECK: %[[B:[^ ]*]] = ptrtoint i32* %[[GEP]] to i64
+; CHECK: %[[B:[^ ]*]] = ptrtoint i32* %[[BC]] to i64
 ; CHECK: %[[C:[^ ]*]] = shl i64 %[[A]], 56
 ; CHECK: or i64 %[[B]], %[[C]]
 
 entry:
   %x = alloca i32, align 4
   call void @use32(i32* nonnull %x)
   ret void
 }

llvm/test/Instrumentation/HWAddressSanitizer/alloca.ll

 define void @test_alloca() sanitize_hwaddress {
 ; CHECK-LABEL: @test_alloca(
 ; CHECK: %[[FP:[^ ]*]] = call i8* @llvm.frameaddress(i32 0)
 ; CHECK: %[[A:[^ ]*]] = ptrtoint i8* %[[FP]] to i64
 ; CHECK: %[[B:[^ ]*]] = lshr i64 %[[A]], 20
 ; CHECK: %[[BASE_TAG:[^ ]*]] = xor i64 %[[A]], %[[B]]
 
 ; CHECK: %[[X:[^ ]*]] = alloca { i32, [12 x i8] }, align 16
-; CHECK: %[[X_GEP:[^ ]*]] = getelementptr { i32, [12 x i8] }, { i32, [12 x i8] }* %[[X]], i32 0, i32 0
+; CHECK: %[[X_BC:[^ ]*]] = bitcast { i32, [12 x i8] }* %[[X]] to i32*
 ; CHECK: %[[X_TAG:[^ ]*]] = xor i64 %[[BASE_TAG]], 0
-; CHECK: %[[X1:[^ ]*]] = ptrtoint i32* %[[X_GEP]] to i64
+; CHECK: %[[X1:[^ ]*]] = ptrtoint i32* %[[X_BC]] to i64
 ; CHECK: %[[C:[^ ]*]] = shl i64 %[[X_TAG]], 56
 ; CHECK: %[[D:[^ ]*]] = or i64 %[[X1]], %[[C]]
 ; CHECK: %[[X_HWASAN:[^ ]*]] = inttoptr i64 %[[D]] to i32*
 
 ; CHECK: %[[X_TAG2:[^ ]*]] = trunc i64 %[[X_TAG]] to i8
-; CHECK: %[[E:[^ ]*]] = ptrtoint i32* %[[X_GEP]] to i64
+; CHECK: %[[E:[^ ]*]] = ptrtoint i32* %[[X_BC]] to i64
 ; CHECK: %[[F:[^ ]*]] = lshr i64 %[[E]], 4
 ; DYNAMIC-SHADOW: %[[X_SHADOW:[^ ]*]] = getelementptr i8, i8* %.hwasan.shadow, i64 %[[F]]
 ; ZERO-BASED-SHADOW: %[[X_SHADOW:[^ ]*]] = inttoptr i64 %[[F]] to i8*
 ; CHECK: %[[X_SHADOW_GEP:[^ ]*]] = getelementptr i8, i8* %[[X_SHADOW]], i32 0
 ; CHECK: store i8 4, i8* %[[X_SHADOW_GEP]]
-; CHECK: %[[X_I8:[^ ]*]] = bitcast i32* %[[X_GEP]] to i8*
+; CHECK: %[[X_I8:[^ ]*]] = bitcast i32* %[[X_BC]] to i8*
 ; CHECK: %[[X_I8_GEP:[^ ]*]] = getelementptr i8, i8* %[[X_I8]], i32 15
 ; CHECK: store i8 %[[X_TAG2]], i8* %[[X_I8_GEP]]
 ; CHECK: call void @use32(i32* nonnull %[[X_HWASAN]])
 
 ; UAR-TAGS: %[[BASE_TAG_COMPL:[^ ]*]] = xor i64 %[[BASE_TAG]], 255
 ; UAR-TAGS: %[[X_TAG_UAR:[^ ]*]] = trunc i64 %[[BASE_TAG_COMPL]] to i8
-; CHECK: %[[E2:[^ ]*]] = ptrtoint i32* %[[X_GEP]] to i64
+; CHECK: %[[E2:[^ ]*]] = ptrtoint i32* %[[X_BC]] to i64
 ; CHECK: %[[F2:[^ ]*]] = lshr i64 %[[E2]], 4
 ; DYNAMIC-SHADOW: %[[X_SHADOW2:[^ ]*]] = getelementptr i8, i8* %.hwasan.shadow, i64 %[[F2]]
 ; ZERO-BASED-SHADOW: %[[X_SHADOW2:[^ ]*]] = inttoptr i64 %[[F2]] to i8*
 ; NO-UAR-TAGS: call void @llvm.memset.p0i8.i64(i8* align 1 %[[X_SHADOW2]], i8 0, i64 1, i1 false)
 ; UAR-TAGS: call void @llvm.memset.p0i8.i64(i8* align 1 %[[X_SHADOW2]], i8 %[[X_TAG_UAR]], i64 1, i1 false)
 ; CHECK: ret void
 
 
 entry:
   %x = alloca i32, align 4
   call void @use32(i32* nonnull %x)
   ret void
 }

llvm/test/Instrumentation/HWAddressSanitizer/kernel-alloca.ll

 define void @test_alloca() sanitize_hwaddress {
 ; CHECK-LABEL: @test_alloca(
 ; CHECK: %[[FP:[^ ]*]] = call i8* @llvm.frameaddress(i32 0)
 ; CHECK: %[[A:[^ ]*]] = ptrtoint i8* %[[FP]] to i64
 ; CHECK: %[[B:[^ ]*]] = lshr i64 %[[A]], 20
 ; CHECK: %[[BASE_TAG:[^ ]*]] = xor i64 %[[A]], %[[B]]
 
 ; CHECK: %[[X:[^ ]*]] = alloca { i32, [12 x i8] }, align 16
-; CHECK: %[[X_GEP:[^ ]*]] = getelementptr { i32, [12 x i8] }, { i32, [12 x i8] }* %[[X]], i32 0, i32 0
+; CHECK: %[[X_BC:[^ ]*]] = bitcast { i32, [12 x i8] }* %[[X]] to i32*
 ; CHECK: %[[X_TAG:[^ ]*]] = xor i64 %[[BASE_TAG]], 0
-; CHECK: %[[X1:[^ ]*]] = ptrtoint i32* %[[X_GEP]] to i64
+; CHECK: %[[X1:[^ ]*]] = ptrtoint i32* %[[X_BC]] to i64
 ; CHECK: %[[C:[^ ]*]] = shl i64 %[[X_TAG]], 56
 ; CHECK: %[[D:[^ ]*]] = or i64 %[[C]], 72057594037927935
 ; CHECK: %[[E:[^ ]*]] = and i64 %[[X1]], %[[D]]
 ; CHECK: %[[X_HWASAN:[^ ]*]] = inttoptr i64 %[[E]] to i32*
 
 entry:
   %x = alloca i32, align 4
   call void @use32(i32* nonnull %x)
   ret void
 }