From 8fac07a12adcd968cd4b6a1c1067e4596ad40b1a Mon Sep 17 00:00:00 2001
From: Peter Collingbourne <peter@pcc.me.uk>
Date: Thu, 30 Apr 2020 16:42:17 -0700
Subject: [PATCH] scudo: Exclude previous tag when retagging freed memory.

This means that immediate use after free will be detected 100% of
the time.

Differential Revision: https://reviews.llvm.org/D79216
---
 compiler-rt/lib/scudo/standalone/combined.h |  4 +++-
 compiler-rt/lib/scudo/standalone/memtag.h   | 19 ++++++++++---------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/compiler-rt/lib/scudo/standalone/combined.h b/compiler-rt/lib/scudo/standalone/combined.h
index 112aefd..583b47a 100644
--- a/compiler-rt/lib/scudo/standalone/combined.h
+++ b/compiler-rt/lib/scudo/standalone/combined.h
@@ -983,7 +983,9 @@ private:
     if (UNLIKELY(NewHeader.ClassId && useMemoryTagging())) {
       u8 PrevTag = extractTag(loadTag(reinterpret_cast<uptr>(Ptr)));
       uptr TaggedBegin, TaggedEnd;
-      setRandomTag(Ptr, Size, &TaggedBegin, &TaggedEnd);
+      // Exclude the previous tag so that immediate use after free is detected
+      // 100% of the time.
+      setRandomTag(Ptr, Size, 1UL << PrevTag, &TaggedBegin, &TaggedEnd);
       storeDeallocationStackMaybe(Ptr, PrevTag);
     }
     // If the quarantine is disabled, the actual size of a chunk is 0 or larger
diff --git a/compiler-rt/lib/scudo/standalone/memtag.h b/compiler-rt/lib/scudo/standalone/memtag.h
index 18dae2b..6f347f4 100644
--- a/compiler-rt/lib/scudo/standalone/memtag.h
+++ b/compiler-rt/lib/scudo/standalone/memtag.h
@@ -93,8 +93,8 @@ class ScopedDisableMemoryTagChecks {
   }
 };
 
-inline void setRandomTag(void *Ptr, uptr Size, uptr *TaggedBegin,
-                         uptr *TaggedEnd) {
+inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
+                         uptr *TaggedBegin, uptr *TaggedEnd) {
   void *End;
   __asm__ __volatile__(
       R"(
@@ -102,7 +102,7 @@ inline void setRandomTag(void *Ptr, uptr Size, uptr *TaggedBegin,
 
     // Set a random tag for Ptr in TaggedPtr. This needs to happen even if
     // Size = 0 so that TaggedPtr ends up pointing at a valid address.
-    irg %[TaggedPtr], %[Ptr]
+    irg %[TaggedPtr], %[Ptr], %[ExcludeMask]
     mov %[Cur], %[TaggedPtr]
 
     // Skip the loop if Size = 0. We don't want to do any tagging in this case.
@@ -120,9 +120,9 @@ inline void setRandomTag(void *Ptr, uptr Size, uptr *TaggedBegin,
 
   2:
   )"
-      : [ TaggedPtr ] "=&r"(*TaggedBegin), [ Cur ] "=&r"(*TaggedEnd),
-        [ End ] "=&r"(End)
-      : [ Ptr ] "r"(Ptr), [ Size ] "r"(Size)
+      :
+      [TaggedPtr] "=&r"(*TaggedBegin), [Cur] "=&r"(*TaggedEnd), [End] "=&r"(End)
+      : [Ptr] "r"(Ptr), [Size] "r"(Size), [ExcludeMask] "r"(ExcludeMask)
       : "memory");
 }
 
@@ -138,7 +138,7 @@ inline void *prepareTaggedChunk(void *Ptr, uptr Size, uptr BlockEnd) {
                        : "memory");
 
   uptr TaggedBegin, TaggedEnd;
-  setRandomTag(Ptr, Size, &TaggedBegin, &TaggedEnd);
+  setRandomTag(Ptr, Size, 0, &TaggedBegin, &TaggedEnd);
 
   // Finally, set the tag of the granule past the end of the allocation to 0,
   // to catch linear overflows even if a previous larger allocation used the
@@ -225,10 +225,11 @@ struct ScopedDisableMemoryTagChecks {
   ScopedDisableMemoryTagChecks() {}
 };
 
-inline void setRandomTag(void *Ptr, uptr Size, uptr *TaggedBegin,
-                         uptr *TaggedEnd) {
+inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
+                         uptr *TaggedBegin, uptr *TaggedEnd) {
   (void)Ptr;
   (void)Size;
+  (void)ExcludeMask;
   (void)TaggedBegin;
   (void)TaggedEnd;
   UNREACHABLE("memory tagging not supported");
-- 
2.7.4