[Asan] Ensure unpoisonning doesn't get inlined unnecessarily due to small holes in the mask

[saudi]

Summary

In the case of allocas introduced by the Address Sanitizer, the shadow bytes may contain sparse zeros, missing opportunities for grouping into __asan_set_shadow_00 calls, causing inlining despite the "asan-max-inline-poisoning-size" argument value.

This fix ignores small holes in the shadow bytes when unpoisonning.

[MaskRay]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

2956

llvm style prefers pre-increment.

[rsundahl]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

2964

This function is probably not the place to guess that it's ok to ignore the shadow mask because it's assumed that it's being called for un-poisoning. While this may be the case, the shadow mask and shadow data provide high fidelity control over behavior and assuming the intent at this low level is problematic. I realize that at a line 2888 the same assumption is made during inlining, but moving this up to the caller and cleaning up these two functions would be far superior.

llvm/test/Instrumentation/AddressSanitizer/asan-optimize-inline-poisoning.ll

2

Your test when run with -asan-max-inline-poisoning-size=0 still inlines but it doesn't look like it is due to your changes (unless datalayout is contrived for the test and doesn't occur otherwise.) It's more likely that my testing for https://reviews.llvm.org/D136197 was inadequate.

[saudi]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

2964

I can take a look and try to move the logic to a higher level.
The copyToShadow declaration has a comment attached:

// Copies bytes from ShadowBytes into shadow memory for indexes where
// ShadowMask is not zero. If ShadowMask[i] is zero, we assume that
// ShadowBytes[i] is constantly zero and doesn't need to be overwritten.

Would the solution be to:

1. remove from copyToShadow / copyToShadowInline the assumption about zero (and remove the second sentence from the comment I mentionned above)
2. Add a method bridgeUnpoisonningGaps that would optimize the shadow buffer so that the generated instructions are more optimized

?

I'm a bit stuck, as this may make bridgeUnpoisonningGaps assume too much about the optimization details.

What do you think?

llvm/test/Instrumentation/AddressSanitizer/asan-optimize-inline-poisoning.ll

2

I tried this on trunk, but the output didn't seem to contain inlined poisonning. see https://llvm.godbolt.org/z/qWr3P54Ex

The generated stores don't look related with [un]poisonning, as they don't write to shadow memory.

Or did I miss something?

[rsundahl]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

2964

An optimization may be a much better place to capture the idea and your proposal seems good provided more senior reviewers are supportive . I have looked at this loopy/call/inline code enough to wish it simpler and if the shadow mask wasn't broken up into these smaller fragments then the code could be simpler and more readable. (This effectively would be equivalent to doing a local pre-scan of the shadowMask and setting it to 1 wherever the shadowData was 0, and then falling into the loop. Might be a good way to test it first.)

llvm/test/Instrumentation/AddressSanitizer/asan-optimize-inline-poisoning.ll

2

I think you'll find the stores at line 69 and line 139 are inline accesses to the shadow memory. That said, I don't think that this is a consequence of your change but more a result of the logic in the loop that's already there.

[vitalybuka]

Is this still relevant?

[saudi]

In D143565#4620489, @vitalybuka wrote:

Is this still relevant?

Hello, yes it still is to us, we are using this in our fork. I didn't have time to address the comments. and I would need some time also to get back into it, understanding the specifics and how to refactor this.

The real-life context that made this optimization necessary:

• Windows x64 target, with EH enabled
• a large function with many blocks, most of which contained local variables often with small gaps between them due to alignment etc. (it was a catch2 framework unittest, where a TEST_CASE contained a big hierarchy SECTION)

With EH enabled, a large number of cleanup pads were generated, causing a lot of variables deinitialization code.
Having small gaps between variables, the shadow poisoning would be split into many inlined parts

The result was very large code size output by asan (>100MB object files), most of which was poisoning code.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

void FunctionStackPoisoner::copyToShadow(ArrayRef<uint8_t> ShadowMask,
   copyToShadow(ShadowMask, ShadowBytes, 0, ShadowMask.size(), IRB, ShadowBase);
 }
 
 void FunctionStackPoisoner::copyToShadow(ArrayRef<uint8_t> ShadowMask,
                                          ArrayRef<uint8_t> ShadowBytes,
                                          size_t Begin, size_t End,
                                          IRBuilder<> &IRB, Value *ShadowBase) {
   assert(ShadowMask.size() == ShadowBytes.size());
+  const size_t MaxHoleSize =
+      std::min<size_t>(sizeof(uint64_t), ASan.LongSize / 8);
   size_t Done = Begin;
   for (size_t i = Begin, j = Begin + 1; i < End; i = j++) {
     if (!ShadowMask[i]) {
       assert(!ShadowBytes[i]);
       continue;
     }
     uint8_t Val = ShadowBytes[i];
     if (!AsanSetShadowFunc[Val])
       continue;
 
     // Skip same values.
     for (; j < End && ShadowMask[j] && Val == ShadowBytes[j]; ++j) {
     }
 
+    // Special case for zero value: allow a few consecutive zeros in the middle
+    // of the mask. It avoids splitting the buffer into many inlined chunks.
+    if (!Val) {
+      size_t HoleSize = 0;
+      for (; j < End && !ShadowBytes[j] && HoleSize <= MaxHoleSize; ++j) {
+        if (!ShadowMask[j])
+          HoleSize++;
+        else
+          HoleSize = 0;
+      }
+      // Ignore trailing zeros in the mask.
+      j -= HoleSize;
+    }
+
     if (j - i >= ClMaxInlinePoisoningSize) {
       copyToShadowInline(ShadowMask, ShadowBytes, Done, i, IRB, ShadowBase);
       IRB.CreateCall(AsanSetShadowFunc[Val],
                      {IRB.CreateAdd(ShadowBase, ConstantInt::get(IntptrTy, i)),
                       ConstantInt::get(IntptrTy, j - i)});
       Done = j;
     }
   }
 

llvm/test/Instrumentation/AddressSanitizer/asan-optimize-inline-poisoning.ll

+; RUN: opt < %s  -passes=asan -asan-max-inline-poisoning-size=8 -S | FileCheck %s --check-prefix=CHECK
+
+target datalayout = "e-m:w-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.20.0"
+
+%struct.S = type { ptr, ptr }
+
+define void @allocas() sanitize_address {
+; CHECK-LABEL: void @allocas()
+entry:
+  %agg.tmp = alloca %struct.S, align 8
+  %agg.tmp1 = alloca %struct.S, align 8
+  %agg.tmp2 = alloca %struct.S, align 8
+  %agg.tmp3 = alloca %struct.S, align 8
+  %agg.tmp4 = alloca %struct.S, align 8
+  %agg.tmp5 = alloca %struct.S, align 8
+  %agg.tmp6 = alloca %struct.S, align 8
+  %agg.tmp7 = alloca %struct.S, align 8
+  %agg.tmp8 = alloca %struct.S, align 8
+  %agg.tmp9 = alloca %struct.S, align 8
+  %agg.tmp10 = alloca %struct.S, align 8
+  %agg.tmp11 = alloca %struct.S, align 8
+  %agg.tmp12 = alloca %struct.S, align 8
+  %agg.tmp13 = alloca %struct.S, align 8
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp)
+  call void @use_copy(ptr noundef nonnull %agg.tmp)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp1)
+  call void @use_copy(ptr noundef nonnull %agg.tmp1)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp2)
+  call void @use_copy(ptr noundef nonnull %agg.tmp2)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp3)
+  call void @use_copy(ptr noundef nonnull %agg.tmp3)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp4)
+  call void @use_copy(ptr noundef nonnull %agg.tmp4)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp5)
+  call void @use_copy(ptr noundef nonnull %agg.tmp5)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp6)
+  call void @use_copy(ptr noundef nonnull %agg.tmp6)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp7)
+  call void @use_copy(ptr noundef nonnull %agg.tmp7)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp8)
+  call void @use_copy(ptr noundef nonnull %agg.tmp8)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp9)
+  call void @use_copy(ptr noundef nonnull %agg.tmp9)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp10)
+  call void @use_copy(ptr noundef nonnull %agg.tmp10)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp11)
+  call void @use_copy(ptr noundef nonnull %agg.tmp11)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp12)
+  call void @use_copy(ptr noundef nonnull %agg.tmp12)
+  call void @get(ptr nonnull sret(%struct.S) align 8 %agg.tmp13)
+  call void @use_copy(ptr noundef nonnull %agg.tmp13)
+  br label %return
+return:
+; CHECK: return:
+; CHECK: call void @__asan_set_shadow_f5(i64 %{{.*}}, i64 64)
+; CHECK: call void @__asan_set_shadow_00(i64 %{{.*}}, i64 60)
+; CHECK: ret void
+  ret void
+}
+
+declare void @use_copy(ptr noundef)
+declare void @get(ptr sret(%struct.S) align 8, ...)