Better memory poisoning helpers

commit_hash:dc1bba1db2b4a5bae15299e926e816920168018f

3 files changed, 61 insertions, 42 deletions

diff --git a/library/cpp/yt/memory/poison-inl.h b/library/cpp/yt/memory/poison-inl.h
index c7563565a89..1625de04847 100644
--- a/library/cpp/yt/memory/poison-inl.h
+++ b/library/cpp/yt/memory/poison-inl.h
@@ -18,12 +18,15 @@ void __asan_poison_memory_region(void const volatile *addr, size_t size);
 void __asan_unpoison_memory_region(void const volatile *addr, size_t size);
 } // extern "C"
 
-Y_FORCE_INLINE void PoisonMemory(TMutableRef ref)
+Y_FORCE_INLINE void PoisonUninitializedMemory(TMutableRef /*ref*/)
+{ }
+
+Y_FORCE_INLINE void PoisonFreedMemory(TMutableRef ref)
 {
     __asan_poison_memory_region(ref.data(), ref.size());
 }
 
-Y_FORCE_INLINE void UnpoisonMemory(TMutableRef ref)
+Y_FORCE_INLINE void RecycleFreedMemory(TMutableRef ref)
 {
     __asan_unpoison_memory_region(ref.data(), ref.size());
 }
@@ -31,26 +34,31 @@ Y_FORCE_INLINE void UnpoisonMemory(TMutableRef ref)
 #elif defined(_msan_enabled_)
 
 extern "C" {
-void __msan_unpoison(const volatile void* a, size_t size);
 void __msan_poison(const volatile void* a, size_t size);
 } // extern "C"
 
-Y_FORCE_INLINE void PoisonMemory(TMutableRef ref)
+Y_FORCE_INLINE void PoisonUninitializedMemory(TMutableRef ref)
 {
     __msan_poison(ref.data(), ref.size());
 }
 
-Y_FORCE_INLINE void UnpoisonMemory(TMutableRef ref)
+Y_FORCE_INLINE void PoisonFreedMemory(TMutableRef ref)
 {
-    __msan_unpoison(ref.data(), ref.size());
+    __msan_poison(ref.data(), ref.size());
 }
 
+Y_FORCE_INLINE void RecycleFreedMemory(TMutableRef /*ref*/)
+{ }
+
 #elif defined(NDEBUG)
 
-Y_FORCE_INLINE void PoisonMemory(TMutableRef /*ref*/)
+Y_FORCE_INLINE void PoisonUninitializedMemory(TMutableRef /*ref*/)
+{ }
+
+Y_FORCE_INLINE void PoisonFreedMemory(TMutableRef /*ref*/)
 { }
 
-Y_FORCE_INLINE void UnpoisonMemory(TMutableRef /*ref*/)
+Y_FORCE_INLINE void RecycleFreedMemory(TMutableRef /*ref*/)
 { }
 
 #endif
diff --git a/library/cpp/yt/memory/poison.cpp b/library/cpp/yt/memory/poison.cpp
index bc4bcad4e02..2d05246df8f 100644
--- a/library/cpp/yt/memory/poison.cpp
+++ b/library/cpp/yt/memory/poison.cpp
@@ -6,57 +6,51 @@ namespace NYT {
 
 namespace {
 
-template <char Byte0, char Byte1, char Byte2, char Byte3, char Byte4, char Byte5, char Byte6, char Byte7>
+template <char Byte0, char Byte1, char Byte2, char Byte3>
 void ClobberMemory(char* __restrict__ ptr, size_t size)
 {
-    while (size >= 8) {
+    while (size >= 4) {
         *ptr++ = Byte0;
         *ptr++ = Byte1;
         *ptr++ = Byte2;
         *ptr++ = Byte3;
-        *ptr++ = Byte4;
-        *ptr++ = Byte5;
-        *ptr++ = Byte6;
-        *ptr++ = Byte7;
-        size -= 8;
+        size -= 4;
     }
 
     switch (size) {
-        case 7:
-            *ptr++ = Byte0;
-            [[fallthrough]];
-        case 6:
-            *ptr++ = Byte1;
-            [[fallthrough]];
-        case 5:
-            *ptr++ = Byte2;
-            [[fallthrough]];
-        case 4:
-            *ptr++ = Byte3;
-            [[fallthrough]];
         case 3:
-            *ptr++ = Byte4;
+            *ptr++ = Byte0;
             [[fallthrough]];
         case 2:
-            *ptr++ = Byte5;
+            *ptr++ = Byte1;
             [[fallthrough]];
         case 1:
-            *ptr++ = Byte6;
+            *ptr++ = Byte2;
     }
 }
 
 } // namespace
 
 #if !defined(NDEBUG) && !defined(_asan_enabled_) && !defined(_msan_enabled_)
-void PoisonMemory(TMutableRef ref)
+
+void PoisonUninitializedMemory(TMutableRef ref)
 {
-    ClobberMemory<'d', 'e', 'a', 'd', 'b', 'e', 'e', 'f'>(ref.data(), ref.size());
+    // BAADBOBA
+    ClobberMemory<'\xba', '\xad', '\xb0', '\xba'>(ref.data(), ref.size());
 }
 
-void UnpoisonMemory(TMutableRef ref)
+void PoisonFreedMemory(TMutableRef ref)
 {
-    ClobberMemory<'c', 'a', 'f', 'e', 'b', 'a', 'b', 'e'>(ref.data(), ref.size());
+    // DEADBEEF
+    ClobberMemory<'\xde', '\xad', '\xbe', '\xef'>(ref.data(), ref.size());
 }
+
+void RecycleFreedMemory(TMutableRef ref)
+{
+    // COOLBIBA
+    ClobberMemory<'\xc0', '\x01', '\xb1', '\xba'>(ref.data(), ref.size());
+}
+
 #endif
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/poison.h b/library/cpp/yt/memory/poison.h
index 9519f3da10e..cdd7eb557a0 100644
--- a/library/cpp/yt/memory/poison.h
+++ b/library/cpp/yt/memory/poison.h
@@ -6,15 +6,32 @@ namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-//! In release builds, does nothing.
-//! In checked builds, clobbers memory with garbage pattern.
-//! In sanitized builds, invokes sanitizer poisoning.
-void PoisonMemory(TMutableRef ref);
+//! Poisons an uninitialized slice of memory.
+/*
+ * In release builds, does nothing.
+ * In checked builds, clobbers memory with a garbage pattern.
+ * In ASAN builds, does nothing.
+ * In MSAN builds, invokes sanitizer poisoning to catch uninit-read.
+ */
+void PoisonUninitializedMemory(TMutableRef ref);
 
-//! In release builds, does nothing.
-//! In checked builds, clobbers memory with (another) garbage pattern.
-//! In sanitized builds, invokes sanitizer unpoisoning.
-void UnpoisonMemory(TMutableRef ref);
+//! Poisons a freed slice of memory.
+/*
+ * In release builds, does nothing.
+ * In checked builds, clobbers memory with a garbage pattern.
+ * In ASAN and MSAN builds, invokes sanitizer poisoning to catch use-after-free.
+ */
+void PoisonFreedMemory(TMutableRef ref);
+
+//! Indicates that a slice of memory that was previously given to #PoisonFreedMemory
+//! has been recycled and can be reused.
+/*!
+ * In release builds, does nothing.
+ * In checked builds, clobbers memory with a garbage pattern.
+ * In ASAN builds, invokes sanitizer unpoisoning.
+ * In MSAN builds, does nothing (the memory remains poisoned to catch uninit-read).
+ */
+void RecycleFreedMemory(TMutableRef ref);
 
 ////////////////////////////////////////////////////////////////////////////////