Internal change

PiperOrigin-RevId: 756459725

Author: protobuf-github-bot

upb/mem/arena.c

@@ -441,6 +441,9 @@ static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc, size_t first_size) {
   upb_Atomic_Init(&a->body.space_allocated, actual_block_size);
   a->body.blocks = NULL;
   a->body.upb_alloc_cleanup = NULL;
+#ifndef __STDC_NO_THREADS__
+  a->head.tid = 0;
+#endif
   UPB_PRIVATE(upb_Xsan_Init)(UPB_XSAN(&a->body));
 
   _upb_Arena_AddBlock(&a->head, mem, first_block_overhead, actual_block_size);
@@ -479,6 +482,9 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
   a->body.block_alloc = _upb_Arena_MakeBlockAlloc(alloc, 1);
   a->head.UPB_PRIVATE(ptr) = (void*)UPB_ALIGN_MALLOC((uintptr_t)(a + 1));
   a->head.UPB_PRIVATE(end) = UPB_PTR_AT(mem, n, char);
+#ifndef __STDC_NO_THREADS__
+  a->head.tid = 0;
+#endif
   UPB_PRIVATE(upb_Xsan_Init)(UPB_XSAN(&a->body));
 #ifdef UPB_TRACING_ENABLED
   upb_Arena_LogInit(&a->head, n);

upb/mem/internal/arena.h

@@ -11,7 +11,11 @@
 #include <stddef.h>
 #include <stdint.h>
 #include <string.h>
+#ifndef __STDC_NO_THREADS__
+#include <threads.h>
+#endif
 
+#include "upb/port/atomic.h"
 #include "upb/port/sanitizers.h"
 
 // Must be last.
@@ -22,13 +26,16 @@
 //
 // We need this because the decoder inlines a upb_Arena for performance but
 // the full struct is not visible outside of arena.c. Yes, I know, it's awful.
-#define UPB_ARENA_SIZE_HACK (9 + (UPB_XSAN_STRUCT_SIZE * 2))
+#define UPB_ARENA_SIZE_HACK (10 + (UPB_XSAN_STRUCT_SIZE * 2))
 
 // LINT.IfChange(upb_Arena)
 
 struct upb_Arena {
   char* UPB_ONLYBITS(ptr);
   char* UPB_ONLYBITS(end);
+#ifndef __STDC_NO_THREADS__
+  UPB_ATOMIC(thrd_t) tid;
+#endif
   UPB_XSAN_MEMBER
 };
 
@@ -38,6 +45,32 @@ struct upb_Arena {
 extern "C" {
 #endif
 
+// enterTid() and exitTid() are used to compare the current thread ID to the
+// thread ID stored in the arena during allocations. If they are not the same
+// during arena malloc/realloc, the arena is being used by another thread and
+// should crash.
+#ifndef __STDC_NO_THREADS__
+static thrd_t enterTid(struct upb_Arena* a) {
+  thrd_t t = thrd_current();
+  thrd_t old = upb_Atomic_Exchange(&a->tid, t, memory_order_relaxed);
+  if (old != t && old != 0) {
+    __builtin_trap();
+  }
+  return old;
+}
+
+static void exitTid(struct upb_Arena* a, uintptr_t tid) {
+  thrd_t old = upb_Atomic_Exchange(&a->tid, tid, memory_order_relaxed);
+  if (old != thrd_current()) {
+    __builtin_trap();
+  }
+}
+#else
+// No-op for non-threaded builds.
+static uintptr_t enterTid(struct upb_Arena* a) { return 0; }
+static void exitTid(struct upb_Arena* a, uintptr_t tid) {}
+#endif
+
 void UPB_PRIVATE(_upb_Arena_SwapIn)(struct upb_Arena* des,
                                     const struct upb_Arena* src);
 void UPB_PRIVATE(_upb_Arena_SwapOut)(struct upb_Arena* des,
@@ -63,21 +96,24 @@ UPB_INLINE bool UPB_PRIVATE(_upb_Arena_IsAligned)(const void* ptr) {
 }
 
 UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size) {
+  uintptr_t tid = enterTid(a);
   UPB_PRIVATE(upb_Xsan_AccessReadWrite)(UPB_XSAN(a));
 
   size_t span = UPB_PRIVATE(_upb_Arena_AllocSpan)(size);
 
   if (UPB_UNLIKELY(UPB_PRIVATE(_upb_ArenaHas)(a) < span)) {
     void* UPB_PRIVATE(_upb_Arena_SlowMalloc)(struct upb_Arena * a, size_t size);
-    return UPB_PRIVATE(_upb_Arena_SlowMalloc)(a, span);
+    void* ret = UPB_PRIVATE(_upb_Arena_SlowMalloc)(a, span);
+    exitTid(a, tid);
+    return ret;
   }
 
   // We have enough space to do a fast malloc.
   void* ret = a->UPB_ONLYBITS(ptr);
   a->UPB_ONLYBITS(ptr) += span;
   UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(ret));
   UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(a->UPB_ONLYBITS(ptr)));
-
+  exitTid(a, tid);
   return UPB_PRIVATE(upb_Xsan_NewUnpoisonedRegion)(UPB_XSAN(a), ret, size);
 }
 
@@ -123,6 +159,7 @@ UPB_API_INLINE bool upb_Arena_TryExtend(struct upb_Arena* a, void* ptr,
 
 UPB_API_INLINE void* upb_Arena_Realloc(struct upb_Arena* a, void* ptr,
                                        size_t oldsize, size_t size) {
+  uintptr_t tid = enterTid(a);
   UPB_PRIVATE(upb_Xsan_AccessReadWrite)(UPB_XSAN(a));
 
   void* ret;
@@ -145,7 +182,10 @@ UPB_API_INLINE void* upb_Arena_Realloc(struct upb_Arena* a, void* ptr,
   // We want to invalidate pointers to the old region if hwasan is enabled, so
   // we poison and unpoison even if ptr == ret.
   UPB_PRIVATE(upb_Xsan_PoisonRegion)(ptr, oldsize);
-  return UPB_PRIVATE(upb_Xsan_NewUnpoisonedRegion)(UPB_XSAN(a), ret, size);
+  void* ret_final =
+      UPB_PRIVATE(upb_Xsan_NewUnpoisonedRegion)(UPB_XSAN(a), ret, size);
+  exitTid(a, tid);
+  return ret_final;
 }
 
 #ifdef __cplusplus