[asan] Add support of lifetime poisoning into ComputeASanStackFrameLayout

[vitalybuka]

Summary

We are going to combine poisoning of red zones and scope poisoning.
This functionality is still disabled on instrumentation side by setting LifetimeSize to 0.
PR27453

[eugenis]

LGTM

lib/Transforms/Utils/ASanStackFrameLayout.cpp

84

You need a stronger assert: lifetime aligned up is not larger than size. Right?

[eugenis]

lib/Transforms/Utils/ASanStackFrameLayout.cpp

84

OK, this assert is fine.

include/llvm/Transforms/Utils/ASanStackFrameLayout.h

 
 class AllocaInst;
 
 // These magic constants should be the same as in
 // in asan_internal.h from ASan runtime in compiler-rt.
 static const int kAsanStackLeftRedzoneMagic = 0xf1;
 static const int kAsanStackMidRedzoneMagic = 0xf2;
 static const int kAsanStackRightRedzoneMagic = 0xf3;
+static const int kAsanStackUseAfterScopeMagic = 0xf8;
 
 // Input/output data struct for ComputeASanStackFrameLayout.
 struct ASanStackVariableDescription {
-  const char *Name;  // Name of the variable that will be displayed by asan
+  const char *Name;    // Name of the variable that will be displayed by asan
                        // if a stack-related bug is reported.
-  uint64_t Size;     // Size of the variable in bytes.
-  size_t Alignment;  // Alignment of the variable (power of 2).
-  AllocaInst *AI;    // The actual AllocaInst.
-  size_t Offset;     // Offset from the beginning of the frame;
+  uint64_t Size;       // Size of the variable in bytes.
+  size_t LifetimeSize; // Size in bytes to use for lifetime analysis check.
+                       // Will be rounded up to Granularity.
+  size_t Alignment;    // Alignment of the variable (power of 2).
+  AllocaInst *AI;      // The actual AllocaInst.
+  size_t Offset;       // Offset from the beginning of the frame;
                        // set by ComputeASanStackFrameLayout.
 };
 
 // Output data struct for ComputeASanStackFrameLayout.
 struct ASanStackFrameLayout {
   // Frame description, see DescribeAddressIfStack in ASan runtime.
   SmallString<64> DescriptionString;
   // The contents of the shadow memory for the stack frame that we need
   // to set at function entry.

lib/Transforms/Instrumentation/AddressSanitizer.cpp

void FunctionStackPoisoner::poisonStack() {
   // If we have a call to llvm.localescape, keep it in the entry block.
   if (LocalEscapeCall) LocalEscapeCall->moveBefore(InsBefore);
 
   SmallVector<ASanStackVariableDescription, 16> SVD;
   SVD.reserve(AllocaVec.size());
   for (AllocaInst *AI : AllocaVec) {
     ASanStackVariableDescription D = {AI->getName().data(),
                                       ASan.getAllocaSizeInBytes(*AI),
-                                      AI->getAlignment(), AI, 0};
+                                      0,
+                                      AI->getAlignment(),
+                                      AI,
+                                      0};
     SVD.push_back(D);
   }
   // Minimal header size (left redzone) is 4 pointers,
   // i.e. 32 bytes on 64-bit platforms and 16 bytes in 32-bit platforms.
   size_t MinHeaderSize = ASan.LongSize / 2;
   ASanStackFrameLayout L;
   ComputeASanStackFrameLayout(SVD, 1ULL << Mapping.Scale, MinHeaderSize, &L);
   DEBUG(dbgs() << L.DescriptionString << " --- " << L.FrameSize << "\n");

lib/Transforms/Utils/ASanStackFrameLayout.cpp

for (size_t i = 0; i < NumVars; i++) {
     size_t Alignment = std::max(Granularity, Vars[i].Alignment);
     (void)Alignment;  // Used only in asserts.
     size_t Size = Vars[i].Size;
     const char *Name = Vars[i].Name;
     assert((Alignment & (Alignment - 1)) == 0);
     assert(Layout->FrameAlignment >= Alignment);
     assert((Offset % Alignment) == 0);
     assert(Size > 0);
+    assert(Vars[i].LifetimeSize <= Size);
     StackDescription << " " << Offset << " " << Size << " " << strlen(Name)
                      << " " << Name;
     size_t NextAlignment = IsLast ? Granularity
                    : std::max(Granularity, Vars[i + 1].Alignment);
     size_t SizeWithRedzone = VarAndRedzoneSize(Vars[i].Size, NextAlignment);
-    SB.insert(SB.end(), Size / Granularity, 0);
+    size_t LifetimeShadowSize =
+        (Vars[i].LifetimeSize + Granularity - 1) / Granularity;
+    SB.insert(SB.end(), LifetimeShadowSize, kAsanStackUseAfterScopeMagic);
+    if (Size / Granularity >= LifetimeShadowSize) {
+      SB.insert(SB.end(), Size / Granularity - LifetimeShadowSize, 0);
       if (Size % Granularity)
         SB.insert(SB.end(), Size % Granularity);
+    }
     SB.insert(SB.end(), (SizeWithRedzone - Size) / Granularity,
         IsLast ? kAsanStackRightRedzoneMagic
         : kAsanStackMidRedzoneMagic);
     Vars[i].Offset = Offset;
     Offset += SizeWithRedzone;
+    assert(Offset == SB.size() * Granularity);
   }
   if (Offset % MinHeaderSize) {
     size_t ExtraRedzone = MinHeaderSize - (Offset % MinHeaderSize);
     SB.insert(SB.end(), ExtraRedzone / Granularity,
               kAsanStackRightRedzoneMagic);
     Offset += ExtraRedzone;
   }
   Layout->DescriptionString = StackDescription.str();
   Layout->FrameSize = Offset;
   assert((Layout->FrameSize % MinHeaderSize) == 0);
   assert(Layout->FrameSize / Granularity == Layout->ShadowBytes.size());
 }
 
 } // llvm namespace

unittests/Transforms/Utils/ASanStackFrameLayoutTest.cpp

 static std::string
 ShadowBytesToString(ArrayRef<uint8_t> ShadowBytes) {
   std::ostringstream os;
   for (size_t i = 0, n = ShadowBytes.size(); i < n; i++) {
     switch (ShadowBytes[i]) {
       case kAsanStackLeftRedzoneMagic:    os << "L"; break;
       case kAsanStackRightRedzoneMagic:   os << "R"; break;
       case kAsanStackMidRedzoneMagic:     os << "M"; break;
+      case kAsanStackUseAfterScopeMagic:
+        os << "S";
+        break;
       default:                            os << (unsigned)ShadowBytes[i];
     }
   }
   return os.str();
 }
 
 static void TestLayout(SmallVector<ASanStackVariableDescription, 10> Vars,
                        size_t Granularity, size_t MinHeaderSize,
                        const std::string &ExpectedDescr,
                        const std::string &ExpectedShadow) {
   ASanStackFrameLayout L;
   ComputeASanStackFrameLayout(Vars, Granularity, MinHeaderSize, &L);
   EXPECT_EQ(ExpectedDescr, L.DescriptionString);
   EXPECT_EQ(ExpectedShadow, ShadowBytesToString(L.ShadowBytes));
 }
 
 TEST(ASanStackFrameLayout, Test) {
 #define VEC1(a) SmallVector<ASanStackVariableDescription, 10>(1, a)
 #define VEC(a)                                                                 \
   SmallVector<ASanStackVariableDescription, 10>(a, a + sizeof(a) / sizeof(a[0]))
 
-#define VAR(name, size, alignment)                                             \
+#define VAR(name, size, lifetime, alignment)                                   \
   ASanStackVariableDescription name##size##_##alignment = {                    \
     #name #size "_" #alignment,                                                \
     size,                                                                      \
+    lifetime,                                                                  \
     alignment,                                                                 \
     0,                                                                         \
     0                                                                          \
   }
 
-  VAR(a, 1, 1);
-  VAR(p, 1, 32);
-  VAR(p, 1, 256);
-  VAR(a, 2, 1);
-  VAR(a, 3, 1);
-  VAR(a, 4, 1);
-  VAR(a, 7, 1);
-  VAR(a, 8, 1);
-  VAR(a, 9, 1);
-  VAR(a, 16, 1);
-  VAR(a, 41, 1);
-  VAR(a, 105, 1);
+  VAR(a, 1, 0, 1);
+  VAR(p, 1, 0, 32);
+  VAR(p, 1, 0, 256);
+  VAR(a, 2, 0, 1);
+  VAR(a, 3, 0, 1);
+  VAR(a, 4, 0, 1);
+  VAR(a, 7, 0, 1);
+  VAR(a, 8, 8, 1);
+  VAR(a, 9, 0, 1);
+  VAR(a, 16, 0, 1);
+  VAR(a, 41, 9, 1);
+  VAR(a, 105, 103, 1);
 
   TestLayout(VEC1(a1_1), 8, 16, "1 16 1 4 a1_1", "LL1R");
   TestLayout(VEC1(a1_1), 64, 64, "1 64 1 4 a1_1", "L1");
   TestLayout(VEC1(p1_32), 8, 32, "1 32 1 5 p1_32", "LLLL1RRR");
   TestLayout(VEC1(p1_32), 8, 64, "1 64 1 5 p1_32", "LLLLLLLL1RRRRRRR");
 
   TestLayout(VEC1(a1_1), 8, 32, "1 32 1 4 a1_1", "LLLL1RRR");
   TestLayout(VEC1(a2_1), 8, 32, "1 32 2 4 a2_1", "LLLL2RRR");
   TestLayout(VEC1(a3_1), 8, 32, "1 32 3 4 a3_1", "LLLL3RRR");
   TestLayout(VEC1(a4_1), 8, 32, "1 32 4 4 a4_1", "LLLL4RRR");
   TestLayout(VEC1(a7_1), 8, 32, "1 32 7 4 a7_1", "LLLL7RRR");
-  TestLayout(VEC1(a8_1), 8, 32, "1 32 8 4 a8_1", "LLLL0RRR");
+  TestLayout(VEC1(a8_1), 8, 32, "1 32 8 4 a8_1", "LLLLSRRR");
   TestLayout(VEC1(a9_1), 8, 32, "1 32 9 4 a9_1", "LLLL01RR");
   TestLayout(VEC1(a16_1), 8, 32, "1 32 16 5 a16_1", "LLLL00RR");
   TestLayout(VEC1(p1_256), 8, 32, "1 256 1 6 p1_256",
              "LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL1RRR");
-  TestLayout(VEC1(a41_1), 8, 32, "1 32 41 5 a41_1", "LLLL000001RRRRRR");
+  TestLayout(VEC1(a41_1), 8, 32, "1 32 41 5 a41_1", "LLLLSS0001RRRRRR");
   TestLayout(VEC1(a105_1), 8, 32, "1 32 105 6 a105_1",
-             "LLLL00000000000001RRRRRR");
+             "LLLLSSSSSSSSSSSSS1RRRRRR");
 
   {
     ASanStackVariableDescription t[] = {a1_1, p1_256};
-    TestLayout(VEC(t), 8, 32,
-               "2 256 1 6 p1_256 272 1 4 a1_1",
-               "LLLLLLLL" "LLLLLLLL" "LLLLLLLL" "LLLLLLLL" "1M1R");
+    TestLayout(VEC(t), 8, 32, "2 256 1 6 p1_256 272 1 4 a1_1",
+               "LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL1M1R");
   }
 
   {
     ASanStackVariableDescription t[] = {a1_1, a16_1, a41_1};
-    TestLayout(VEC(t), 8, 32,
-               "3 32 1 4 a1_1 48 16 5 a16_1 80 41 5 a41_1",
-               "LLLL" "1M00" "MM00" "0001" "RRRR");
+    TestLayout(VEC(t), 8, 32, "3 32 1 4 a1_1 48 16 5 a16_1 80 41 5 a41_1",
+               "LLLL1M00MMSS0001RRRR");
   }
 #undef VEC1
 #undef VEC
 #undef VAR
 }