[HWASAN] Support short granules in __hwasan_test_shadow

Reviewed By: thurston

Differential Revision: https://reviews.llvm.org/D149430

compiler-rt/lib/hwasan/hwasan.cpp

@@ -445,16 +445,32 @@ void __hwasan_print_shadow(const void *p, uptr sz) {
 sptr __hwasan_test_shadow(const void *p, uptr sz) {
   if (sz == 0)
     return -1;
-  tag_t ptr_tag = GetTagFromPointer((uptr)p);
-  uptr ptr_raw = UntagAddr(reinterpret_cast<uptr>(p));
+  uptr ptr = reinterpret_cast<uptr>(p);
+  tag_t ptr_tag = GetTagFromPointer(ptr);
+  uptr ptr_raw = UntagAddr(ptr);
   uptr shadow_first = MemToShadow(ptr_raw);
-  uptr shadow_last = MemToShadow(ptr_raw + sz - 1);
-  for (uptr s = shadow_first; s <= shadow_last; ++s)
-    if (*(tag_t *)s != ptr_tag) {
-      sptr offset = ShadowToMem(s) - ptr_raw;
+  uptr shadow_last = MemToShadow(ptr_raw + sz);
+  for (uptr s = shadow_first; s < shadow_last; ++s) {
+    if (UNLIKELY(*(tag_t *)s != ptr_tag)) {
+      uptr short_size =
+          ShortTagSize(*(tag_t *)s, AddTagToPointer(ShadowToMem(s), ptr_tag));
+      sptr offset = ShadowToMem(s) - ptr_raw + short_size;
       return offset < 0 ? 0 : offset;
     }
-  return -1;
+  }
+
+  uptr end = ptr + sz;
+  uptr tail_sz = end & (kShadowAlignment - 1);
+  if (!tail_sz)
+    return -1;
+
+  uptr short_size =
+      ShortTagSize(*(tag_t *)shadow_last, end & ~(kShadowAlignment - 1));
+  if (LIKELY(tail_sz <= short_size))
+    return -1;
+
+  sptr offset = sz - tail_sz + short_size;
+  return offset < 0 ? 0 : offset;
 }
 
 u16 __sanitizer_unaligned_load16(const uu16 *p) {

compiler-rt/lib/hwasan/hwasan_checks.h

@@ -125,8 +125,22 @@ __attribute__((always_inline)) static void SigTrap(uptr p, uptr size) {
   // __builtin_unreachable();
 }
 
+__attribute__((always_inline, nodebug)) static inline uptr ShortTagSize(
+    tag_t mem_tag, uptr ptr) {
+  DCHECK(IsAligned(ptr, kShadowAlignment));
+  tag_t ptr_tag = GetTagFromPointer(ptr);
+  if (ptr_tag == mem_tag)
+    return kShadowAlignment;
+  if (!mem_tag || mem_tag >= kShadowAlignment)
+    return 0;
+  if (*(u8 *)(ptr | (kShadowAlignment - 1)) != ptr_tag)
+    return 0;
+  return mem_tag;
+}
+
 __attribute__((always_inline, nodebug)) static inline bool
 PossiblyShortTagMatches(tag_t mem_tag, uptr ptr, uptr sz) {
+  DCHECK(IsAligned(ptr, kShadowAlignment));
   tag_t ptr_tag = GetTagFromPointer(ptr);
   if (ptr_tag == mem_tag)
     return true;

compiler-rt/test/hwasan/TestCases/test_shadow.c

@@ -0,0 +1,41 @@
+// RUN: %clang_hwasan %s -o %t && %run %t
+
+#include <assert.h>
+#include <sanitizer/hwasan_interface.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+int main() {
+  __hwasan_enable_allocator_tagging();
+  for (int sz = 0; sz < 64; ++sz) {
+    fprintf(stderr, "sz: %d\n", sz);
+    char *x = (char *)malloc(sz);
+    do {
+      // Empty range is always OK.
+      for (int b = -16; b < sz + 32; ++b)
+        assert(__hwasan_test_shadow(x + b, 0) == -1);
+
+      int real_sz = sz ? sz : 1;
+      // Unlucky case when we cant distinguish between tag and short granule size.
+      if (__hwasan_tag_pointer(x, real_sz % 16) == x)
+        break;
+
+      // Underflow - the first byte is bad.
+      for (int b = -16; b < 0; ++b)
+        assert(__hwasan_test_shadow(x + b, real_sz) == 0);
+
+      // Inbound ranges.
+      for (int b = 0; b < real_sz; ++b)
+        for (int e = b; e <= real_sz; ++e)
+          assert(__hwasan_test_shadow(x + b, e - b) == -1);
+
+      // Overflow - the first byte after the buffer is bad.
+      for (int b = 0; b <= real_sz; ++b)
+        for (int e = real_sz + 1; e <= real_sz + 64; ++e)
+          assert(__hwasan_test_shadow(x + b, e - b) == (real_sz - b));
+
+    } while (0);
+    free(x);
+  }
+  return 0;
+}