[HWASAN] Support short granules in __hwasan_test_shadow

[vitalybuka]

update

[thurston]

compiler-rt/lib/hwasan/hwasan.cpp

469–473

if (LIKELY(tail_sz <= short_size))
  return -1;

sptr offset = sz - tail_sz + short_size;
return offset < 0 ? 0 : offset;

The condition that tail_sz <= short_size makes it impossible for 'offset < 0' (other than if overflow happens? - that should have stricter handling than simply returning a zero offset)

[vitalybuka]

compiler-rt/lib/hwasan/hwasan.cpp

469–473

Note that can be "sz < tail_sz" when unaligned begin and end in the same granule.

compiler-rt/lib/hwasan/hwasan.cpp

for (uptr s = shadow_first; s <= shadow_last; ++s) {
     else
       Printf("  %zx: %02x\n", granule_addr, mem_tag);
   }
 }
 
 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;
     }
+  }
+
+  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) {
   return *p;
 }
 u32 __sanitizer_unaligned_load32(const uu32 *p) {
   return *p;
 }

compiler-rt/lib/hwasan/hwasan_checks.h

asm volatile(
       "addiw x0, x0, %2\n" ::"r"(x10),
       "r"(x11), "I"(0x40 + SigTrapEncoding(EA, AT)));
 #else
   __builtin_trap();
 #endif
   // __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;
   if (mem_tag >= kShadowAlignment)
     return false;
   if ((ptr & (kShadowAlignment - 1)) + sz > mem_tag)
     return false;
   return *(u8 *)(ptr | (kShadowAlignment - 1)) == ptr_tag;

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

+// 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;
+}