[compiler-rt] [sanitizer] Test for __sanitizer_get_dtls_size (PR #108349)

[Vitaly Buka via llvm-commits]

https://github.com/vitalybuka updated https://github.com/llvm/llvm-project/pull/108349

>From dda9e86e32ada5c97099e2e179c8b780529f1fc5 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka at google.com>
Date: Thu, 12 Sep 2024 01:38:14 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20change?=
 =?UTF-8?q?s=20to=20main=20this=20commit=20is=20based=20on?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4

[skip ci]
---
 .../sanitizer_common_interface.inc            |  1 +
 .../sanitizer_interface_internal.h            |  5 ++++
 .../sanitizer_tls_get_addr.cpp                | 28 +++++++++++--------
 .../lib/sanitizer_common/weak_symbols.txt     |  1 +
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common_interface.inc b/compiler-rt/lib/sanitizer_common/sanitizer_common_interface.inc
index 91be9e9dcc4551..66744aa021e62b 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common_interface.inc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common_interface.inc
@@ -22,6 +22,7 @@ INTERFACE_FUNCTION(__sanitizer_verify_double_ended_contiguous_container)
 INTERFACE_WEAK_FUNCTION(__sanitizer_on_print)
 INTERFACE_WEAK_FUNCTION(__sanitizer_report_error_summary)
 INTERFACE_WEAK_FUNCTION(__sanitizer_sandbox_on_notify)
+INTERFACE_WEAK_FUNCTION(__sanitizer_get_dtls_size)
 // Sanitizer weak hooks
 INTERFACE_WEAK_FUNCTION(__sanitizer_weak_hook_memcmp)
 INTERFACE_WEAK_FUNCTION(__sanitizer_weak_hook_strcmp)
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_interface_internal.h b/compiler-rt/lib/sanitizer_common/sanitizer_interface_internal.h
index cd0d45e2f3fab1..093d30af5d4817 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_interface_internal.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_interface_internal.h
@@ -49,6 +49,11 @@ __sanitizer_sandbox_on_notify(__sanitizer_sandbox_arguments *args);
 SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE void
 __sanitizer_report_error_summary(const char *error_summary);
 
+// Returns size of dynamically allocated block. This function can be overridden by
+// the client.
+SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE __sanitizer::uptr
+__sanitizer_get_dtls_size(const void *tls_begin);
+
 SANITIZER_INTERFACE_ATTRIBUTE void __sanitizer_cov_dump();
 SANITIZER_INTERFACE_ATTRIBUTE void __sanitizer_dump_coverage(
     const __sanitizer::uptr *pcs, const __sanitizer::uptr len);
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_tls_get_addr.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_tls_get_addr.cpp
index e5839f6f497474..173b3e139570c0 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_tls_get_addr.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_tls_get_addr.cpp
@@ -110,15 +110,16 @@ SANITIZER_WEAK_ATTRIBUTE
 const void *__sanitizer_get_allocated_begin(const void *p);
 }
 
-static bool GetDTLSRange(uptr &tls_beg, uptr &tls_size) {
-  const void *start = __sanitizer_get_allocated_begin((void *)tls_beg);
+SANITIZER_INTERFACE_WEAK_DEF(uptr, __sanitizer_get_dtls_size,
+                             const void *tls_begin) {
+  const void *start = __sanitizer_get_allocated_begin(tls_begin);
   if (!start)
-    return false;
-  tls_beg = (uptr)start;
-  tls_size = __sanitizer_get_allocated_size(start);
+    return 0;
+  CHECK_EQ(start, tls_begin);
+  uptr tls_size = __sanitizer_get_allocated_size(start);
   VReport(2, "__tls_get_addr: glibc DTLS suspected; tls={%p,0x%zx}\n",
-          (void *)tls_beg, tls_size);
-  return true;
+          tls_begin, tls_size);
+  return tls_size;
 }
 
 DTLS::DTV *DTLS_on_tls_get_addr(void *arg_void, void *res,
@@ -142,10 +143,12 @@ DTLS::DTV *DTLS_on_tls_get_addr(void *arg_void, void *res,
     // creation.
     VReport(2, "__tls_get_addr: static tls: %p\n", (void *)tls_beg);
     tls_size = 0;
-  } else if (!GetDTLSRange(tls_beg, tls_size)) {
-    VReport(2, "__tls_get_addr: Can't guess glibc version\n");
-    // This may happen inside the DTOR of main thread, so just ignore it.
-    tls_size = 0;
+  } else {
+    tls_size = __sanitizer_get_dtls_size(reinterpret_cast<const void*>(tls_beg));
+    if (!tls_size) {
+      VReport(2, "__tls_get_addr: Can't guess glibc version\n");
+      // This may happen inside the DTOR of main thread, so just ignore it.
+    }
   }
   dtv->beg = tls_beg;
   dtv->size = tls_size;
@@ -160,6 +163,9 @@ bool DTLSInDestruction(DTLS *dtls) {
 }
 
 #else
+SANITIZER_INTERFACE_WEAK_DEF(uptr, __sanitizer_get_dtls_size, const void *) {
+  return 0;
+}
 DTLS::DTV *DTLS_on_tls_get_addr(void *arg, void *res,
   unsigned long, unsigned long) { return 0; }
 DTLS *DTLS_Get() { return 0; }
diff --git a/compiler-rt/lib/sanitizer_common/weak_symbols.txt b/compiler-rt/lib/sanitizer_common/weak_symbols.txt
index 1eb1ce8d6b9c55..77e7b5d9f702e2 100644
--- a/compiler-rt/lib/sanitizer_common/weak_symbols.txt
+++ b/compiler-rt/lib/sanitizer_common/weak_symbols.txt
@@ -1,4 +1,5 @@
 ___sanitizer_free_hook
+___sanitizer_get_dtls_size
 ___sanitizer_malloc_hook
 ___sanitizer_report_error_summary
 ___sanitizer_sandbox_on_notify