1 files changed, 69 insertions, 2 deletions

diff --git a/lib/asan/asan_poisoning.cc b/lib/asan/asan_poisoning.cc
index 3b9d9f6795ca..dc5749243569 100644
--- a/lib/asan/asan_poisoning.cc
+++ b/lib/asan/asan_poisoning.cc
@@ -13,17 +13,19 @@
 //===----------------------------------------------------------------------===//
 
 #include "asan_interceptors.h"
-#include "asan_interface.h"
 #include "asan_internal.h"
 #include "asan_mapping.h"
+#include "sanitizer/asan_interface.h"
+#include "sanitizer_common/sanitizer_libc.h"
 
 namespace __asan {
 
 void PoisonShadow(uptr addr, uptr size, u8 value) {
+  if (!flags()->poison_heap) return;
   CHECK(AddrIsAlignedByGranularity(addr));
   CHECK(AddrIsAlignedByGranularity(addr + size));
   uptr shadow_beg = MemToShadow(addr);
-  uptr shadow_end = MemToShadow(addr + size);
+  uptr shadow_end = MemToShadow(addr + size - SHADOW_GRANULARITY) + 1;
   CHECK(REAL(memset) != 0);
   REAL(memset)((void*)shadow_beg, value, shadow_end - shadow_beg);
 }
@@ -32,6 +34,7 @@ void PoisonShadowPartialRightRedzone(uptr addr,
                                      uptr size,
                                      uptr redzone_size,
                                      u8 value) {
+  if (!flags()->poison_heap) return;
   CHECK(AddrIsAlignedByGranularity(addr));
   u8 *shadow = (u8*)MemToShadow(addr);
   for (uptr i = 0; i < redzone_size;
@@ -151,3 +154,67 @@ void __asan_unpoison_memory_region(void const volatile *addr, uptr size) {
 bool __asan_address_is_poisoned(void const volatile *addr) {
   return __asan::AddressIsPoisoned((uptr)addr);
 }
+
+uptr __asan_region_is_poisoned(uptr beg, uptr size) {
+  if (!size) return 0;
+  uptr end = beg + size;
+  if (!AddrIsInMem(beg)) return beg;
+  if (!AddrIsInMem(end)) return end;
+  uptr aligned_b = RoundUpTo(beg, SHADOW_GRANULARITY);
+  uptr aligned_e = RoundDownTo(end, SHADOW_GRANULARITY);
+  uptr shadow_beg = MemToShadow(aligned_b);
+  uptr shadow_end = MemToShadow(aligned_e);
+  // First check the first and the last application bytes,
+  // then check the SHADOW_GRANULARITY-aligned region by calling
+  // mem_is_zero on the corresponding shadow.
+  if (!__asan::AddressIsPoisoned(beg) &&
+      !__asan::AddressIsPoisoned(end - 1) &&
+      (shadow_end <= shadow_beg ||
+       __sanitizer::mem_is_zero((const char *)shadow_beg,
+                                shadow_end - shadow_beg)))
+    return 0;
+  // The fast check failed, so we have a poisoned byte somewhere.
+  // Find it slowly.
+  for (; beg < end; beg++)
+    if (__asan::AddressIsPoisoned(beg))
+      return beg;
+  UNREACHABLE("mem_is_zero returned false, but poisoned byte was not found");
+  return 0;
+}
+
+// This is a simplified version of __asan_(un)poison_memory_region, which
+// assumes that left border of region to be poisoned is properly aligned.
+static void PoisonAlignedStackMemory(uptr addr, uptr size, bool do_poison) {
+  if (size == 0) return;
+  uptr aligned_size = size & ~(SHADOW_GRANULARITY - 1);
+  PoisonShadow(addr, aligned_size,
+               do_poison ? kAsanStackUseAfterScopeMagic : 0);
+  if (size == aligned_size)
+    return;
+  s8 end_offset = (s8)(size - aligned_size);
+  s8* shadow_end = (s8*)MemToShadow(addr + aligned_size);
+  s8 end_value = *shadow_end;
+  if (do_poison) {
+    // If possible, mark all the bytes mapping to last shadow byte as
+    // unaddressable.
+    if (end_value > 0 && end_value <= end_offset)
+      *shadow_end = (s8)kAsanStackUseAfterScopeMagic;
+  } else {
+    // If necessary, mark few first bytes mapping to last shadow byte
+    // as addressable
+    if (end_value != 0)
+      *shadow_end = Max(end_value, end_offset);
+  }
+}
+
+void __asan_poison_stack_memory(uptr addr, uptr size) {
+  if (flags()->verbosity > 0)
+    Report("poisoning: %p %zx\n", (void*)addr, size);
+  PoisonAlignedStackMemory(addr, size, true);
+}
+
+void __asan_unpoison_stack_memory(uptr addr, uptr size) {
+  if (flags()->verbosity > 0)
+    Report("unpoisoning: %p %zx\n", (void*)addr, size);
+  PoisonAlignedStackMemory(addr, size, false);
+}