Merge "Reland^2 "[MTE] remap stacks with PROT_MTE when requested by dlopened library"" into main

16 files changed, 572 insertions, 57 deletions

diff --git a/libc/bionic/heap_tagging.cpp b/libc/bionic/heap_tagging.cpp
index 48ec955bc..0c1e5060d 100644
--- a/libc/bionic/heap_tagging.cpp
+++ b/libc/bionic/heap_tagging.cpp
@@ -57,6 +57,7 @@ void SetDefaultHeapTaggingLevel() {
         break;
       case M_HEAP_TAGGING_LEVEL_SYNC:
       case M_HEAP_TAGGING_LEVEL_ASYNC:
+        atomic_store(&globals->memtag, true);
         atomic_store(&globals->memtag_stack, __libc_shared_globals()->initial_memtag_stack);
         break;
       default:
@@ -113,6 +114,7 @@ bool SetHeapTaggingLevel(HeapTaggingLevel tag_level) {
           globals->heap_pointer_tag = static_cast<uintptr_t>(0xffull << UNTAG_SHIFT);
         }
         atomic_store(&globals->memtag_stack, false);
+        atomic_store(&globals->memtag, false);
       });
 
       if (heap_tagging_level != M_HEAP_TAGGING_LEVEL_TBI) {
diff --git a/libc/bionic/libc_init_dynamic.cpp b/libc/bionic/libc_init_dynamic.cpp
index c61810e34..1180a513e 100644
--- a/libc/bionic/libc_init_dynamic.cpp
+++ b/libc/bionic/libc_init_dynamic.cpp
@@ -39,11 +39,12 @@
  *   all dynamic linking has been performed.
  */
 
+#include <elf.h>
 #include <stddef.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <stdint.h>
-#include <elf.h>
+#include "bionic/pthread_internal.h"
 #include "libc_init_common.h"
 
 #include "private/bionic_defs.h"
@@ -59,6 +60,11 @@ extern "C" {
   extern int __cxa_atexit(void (*)(void *), void *, void *);
 };
 
+void memtag_stack_dlopen_callback() {
+  async_safe_format_log(ANDROID_LOG_INFO, "libc", "remapping stacks as PROT_MTE");
+  __pthread_internal_remap_stack_with_mte();
+}
+
 // Use an initializer so __libc_sysinfo will have a fallback implementation
 // while .preinit_array constructors run.
 #if defined(__i386__)
@@ -156,6 +162,10 @@ __noreturn void __libc_init(void* raw_args,
 
   __libc_init_mte_late();
 
+  // This roundabout way is needed so we don't use the static libc linked into the linker, which
+  // will not affect the process.
+  __libc_shared_globals()->memtag_stack_dlopen_callback = memtag_stack_dlopen_callback;
+
   exit(slingshot(args.argc - __libc_shared_globals()->initial_linker_arg_count,
                  args.argv + __libc_shared_globals()->initial_linker_arg_count,
                  args.envp));
diff --git a/libc/bionic/libc_init_static.cpp b/libc/bionic/libc_init_static.cpp
index 00faa5ba7..f091ff876 100644
--- a/libc/bionic/libc_init_static.cpp
+++ b/libc/bionic/libc_init_static.cpp
@@ -305,6 +305,14 @@ __attribute__((no_sanitize("hwaddress", "memtag"))) void __libc_init_mte(
   bool memtag_stack = false;
   HeapTaggingLevel level =
       __get_tagging_level(memtag_dynamic_entries, phdr_start, phdr_ct, load_bias, &memtag_stack);
+  // This is used by the linker (in linker.cpp) to communicate than any library linked by this
+  // executable enables memtag-stack.
+  if (__libc_shared_globals()->initial_memtag_stack) {
+    if (!memtag_stack) {
+      async_safe_format_log(ANDROID_LOG_INFO, "libc", "enabling PROT_MTE as requested by linker");
+    }
+    memtag_stack = true;
+  }
   char* env = getenv("BIONIC_MEMTAG_UPGRADE_SECS");
   static const char kAppProcessName[] = "app_process64";
   const char* progname = __libc_shared_globals()->init_progname;
@@ -373,6 +381,8 @@ __attribute__((no_sanitize("hwaddress", "memtag"))) void __libc_init_mte(
   }
   // We did not enable MTE, so we do not need to arm the upgrade timer.
   __libc_shared_globals()->heap_tagging_upgrade_timer_sec = 0;
+  // We also didn't enable memtag_stack.
+  __libc_shared_globals()->initial_memtag_stack = false;
 }
 #else   // __aarch64__
 void __libc_init_mte(const memtag_dynamic_entries_t*, const void*, size_t, uintptr_t, void*) {}
diff --git a/libc/bionic/pthread_attr.cpp b/libc/bionic/pthread_attr.cpp
index de4cc9e8d..f6c0401bb 100644
--- a/libc/bionic/pthread_attr.cpp
+++ b/libc/bionic/pthread_attr.cpp
@@ -155,36 +155,6 @@ int pthread_attr_setstack(pthread_attr_t* attr, void* stack_base, size_t stack_s
   return 0;
 }
 
-static uintptr_t __get_main_stack_startstack() {
-  FILE* fp = fopen("/proc/self/stat", "re");
-  if (fp == nullptr) {
-    async_safe_fatal("couldn't open /proc/self/stat: %m");
-  }
-
-  char line[BUFSIZ];
-  if (fgets(line, sizeof(line), fp) == nullptr) {
-    async_safe_fatal("couldn't read /proc/self/stat: %m");
-  }
-
-  fclose(fp);
-
-  // See man 5 proc. There's no reason comm can't contain ' ' or ')',
-  // so we search backwards for the end of it. We're looking for this field:
-  //
-  //  startstack %lu (28) The address of the start (i.e., bottom) of the stack.
-  uintptr_t startstack = 0;
-  const char* end_of_comm = strrchr(line, ')');
-  if (sscanf(end_of_comm + 1, " %*c "
-             "%*d %*d %*d %*d %*d "
-             "%*u %*u %*u %*u %*u %*u %*u "
-             "%*d %*d %*d %*d %*d %*d "
-             "%*u %*u %*d %*u %*u %*u %" SCNuPTR, &startstack) != 1) {
-    async_safe_fatal("couldn't parse /proc/self/stat");
-  }
-
-  return startstack;
-}
-
 static int __pthread_attr_getstack_main_thread(void** stack_base, size_t* stack_size) {
   ErrnoRestorer errno_restorer;
 
@@ -198,28 +168,11 @@ static int __pthread_attr_getstack_main_thread(void** stack_base, size_t* stack_
   if (stack_limit.rlim_cur == RLIM_INFINITY) {
     stack_limit.rlim_cur = 8 * 1024 * 1024;
   }
-
-  // Ask the kernel where our main thread's stack started.
-  uintptr_t startstack = __get_main_stack_startstack();
-
-  // Hunt for the region that contains that address.
-  FILE* fp = fopen("/proc/self/maps", "re");
-  if (fp == nullptr) {
-    async_safe_fatal("couldn't open /proc/self/maps: %m");
-  }
-  char line[BUFSIZ];
-  while (fgets(line, sizeof(line), fp) != nullptr) {
-    uintptr_t lo, hi;
-    if (sscanf(line, "%" SCNxPTR "-%" SCNxPTR, &lo, &hi) == 2) {
-      if (lo <= startstack && startstack <= hi) {
-        *stack_size = stack_limit.rlim_cur;
-        *stack_base = reinterpret_cast<void*>(hi - *stack_size);
-        fclose(fp);
-        return 0;
-      }
-    }
-  }
-  async_safe_fatal("stack not found in /proc/self/maps");
+  uintptr_t lo, hi;
+  __find_main_stack_limits(&lo, &hi);
+  *stack_size = stack_limit.rlim_cur;
+  *stack_base = reinterpret_cast<void*>(hi - *stack_size);
+  return 0;
 }
 
 __BIONIC_WEAK_FOR_NATIVE_BRIDGE
diff --git a/libc/bionic/pthread_internal.cpp b/libc/bionic/pthread_internal.cpp
index 6a7ee2f5b..bfe2f982e 100644
--- a/libc/bionic/pthread_internal.cpp
+++ b/libc/bionic/pthread_internal.cpp
@@ -40,6 +40,7 @@
 #include "private/ErrnoRestorer.h"
 #include "private/ScopedRWLock.h"
 #include "private/bionic_futex.h"
+#include "private/bionic_globals.h"
 #include "private/bionic_tls.h"
 
 static pthread_internal_t* g_thread_list = nullptr;
@@ -119,6 +120,89 @@ pthread_internal_t* __pthread_internal_find(pthread_t thread_id, const char* cal
   return nullptr;
 }
 
+static uintptr_t __get_main_stack_startstack() {
+  FILE* fp = fopen("/proc/self/stat", "re");
+  if (fp == nullptr) {
+    async_safe_fatal("couldn't open /proc/self/stat: %m");
+  }
+
+  char line[BUFSIZ];
+  if (fgets(line, sizeof(line), fp) == nullptr) {
+    async_safe_fatal("couldn't read /proc/self/stat: %m");
+  }
+
+  fclose(fp);
+
+  // See man 5 proc. There's no reason comm can't contain ' ' or ')',
+  // so we search backwards for the end of it. We're looking for this field:
+  //
+  //  startstack %lu (28) The address of the start (i.e., bottom) of the stack.
+  uintptr_t startstack = 0;
+  const char* end_of_comm = strrchr(line, ')');
+  if (sscanf(end_of_comm + 1,
+             " %*c "
+             "%*d %*d %*d %*d %*d "
+             "%*u %*u %*u %*u %*u %*u %*u "
+             "%*d %*d %*d %*d %*d %*d "
+             "%*u %*u %*d %*u %*u %*u %" SCNuPTR,
+             &startstack) != 1) {
+    async_safe_fatal("couldn't parse /proc/self/stat");
+  }
+
+  return startstack;
+}
+
+void __find_main_stack_limits(uintptr_t* low, uintptr_t* high) {
+  // Ask the kernel where our main thread's stack started.
+  uintptr_t startstack = __get_main_stack_startstack();
+
+  // Hunt for the region that contains that address.
+  FILE* fp = fopen("/proc/self/maps", "re");
+  if (fp == nullptr) {
+    async_safe_fatal("couldn't open /proc/self/maps: %m");
+  }
+  char line[BUFSIZ];
+  while (fgets(line, sizeof(line), fp) != nullptr) {
+    uintptr_t lo, hi;
+    if (sscanf(line, "%" SCNxPTR "-%" SCNxPTR, &lo, &hi) == 2) {
+      if (lo <= startstack && startstack <= hi) {
+        *low = lo;
+        *high = hi;
+        fclose(fp);
+        return;
+      }
+    }
+  }
+  async_safe_fatal("stack not found in /proc/self/maps");
+}
+
+void __pthread_internal_remap_stack_with_mte() {
+#if defined(__aarch64__)
+  // If process doesn't have MTE enabled, we don't need to do anything.
+  if (!__libc_globals->memtag) return;
+  bool prev = true;
+  __libc_globals.mutate(
+      [&prev](libc_globals* globals) { prev = atomic_exchange(&globals->memtag_stack, true); });
+  if (prev) return;
+  uintptr_t lo, hi;
+  __find_main_stack_limits(&lo, &hi);
+
+  if (mprotect(reinterpret_cast<void*>(lo), hi - lo,
+               PROT_READ | PROT_WRITE | PROT_MTE | PROT_GROWSDOWN)) {
+    async_safe_fatal("error: failed to set PROT_MTE on main thread");
+  }
+  ScopedWriteLock creation_locker(&g_thread_creation_lock);
+  ScopedReadLock list_locker(&g_thread_list_lock);
+  for (pthread_internal_t* t = g_thread_list; t != nullptr; t = t->next) {
+    if (t->terminating || t->is_main()) continue;
+    if (mprotect(t->mmap_base_unguarded, t->mmap_size_unguarded,
+                 PROT_READ | PROT_WRITE | PROT_MTE)) {
+      async_safe_fatal("error: failed to set PROT_MTE on thread: %d", t->tid);
+    }
+  }
+#endif
+}
+
 bool android_run_on_all_threads(bool (*func)(void*), void* arg) {
   // Take the locks in this order to avoid inversion (pthread_create ->
   // __pthread_internal_add).
diff --git a/libc/bionic/pthread_internal.h b/libc/bionic/pthread_internal.h
index 3b9e6a481..091f711eb 100644
--- a/libc/bionic/pthread_internal.h
+++ b/libc/bionic/pthread_internal.h
@@ -178,6 +178,7 @@ class pthread_internal_t {
   bionic_tls* bionic_tls;
 
   int errno_value;
+  bool is_main() { return start_routine == nullptr; }
 };
 
 struct ThreadMapping {
@@ -207,6 +208,7 @@ __LIBC_HIDDEN__ pthread_internal_t* __pthread_internal_find(pthread_t pthread_id
 __LIBC_HIDDEN__ pid_t __pthread_internal_gettid(pthread_t pthread_id, const char* caller);
 __LIBC_HIDDEN__ void __pthread_internal_remove(pthread_internal_t* thread);
 __LIBC_HIDDEN__ void __pthread_internal_remove_and_free(pthread_internal_t* thread);
+__LIBC_HIDDEN__ void __find_main_stack_limits(uintptr_t* low, uintptr_t* high);
 
 static inline __always_inline bionic_tcb* __get_bionic_tcb() {
   return reinterpret_cast<bionic_tcb*>(&__get_tls()[MIN_TLS_SLOT]);
@@ -266,6 +268,9 @@ __LIBC_HIDDEN__ extern void __bionic_atfork_run_prepare();
 __LIBC_HIDDEN__ extern void __bionic_atfork_run_child();
 __LIBC_HIDDEN__ extern void __bionic_atfork_run_parent();
 
+// Re-map all threads and successively launched threads with PROT_MTE.
+__LIBC_HIDDEN__ void __pthread_internal_remap_stack_with_mte();
+
 extern "C" bool android_run_on_all_threads(bool (*func)(void*), void* arg);
 
 extern pthread_rwlock_t g_thread_creation_lock;
diff --git a/libc/private/bionic_globals.h b/libc/private/bionic_globals.h
index 23f2953f5..6f1e3895e 100644
--- a/libc/private/bionic_globals.h
+++ b/libc/private/bionic_globals.h
@@ -50,6 +50,7 @@ struct libc_globals {
   uintptr_t heap_pointer_tag;
   _Atomic(bool) memtag_stack;
   _Atomic(bool) decay_time_enabled;
+  _Atomic(bool) memtag;
 
   // In order to allow a complete switch between dispatch tables without
   // the need for copying each function by function in the structure,
@@ -137,6 +138,7 @@ struct libc_shared_globals {
   bool initial_memtag_stack = false;
   int64_t heap_tagging_upgrade_timer_sec = 0;
 
+  void (*memtag_stack_dlopen_callback)() = nullptr;
   pthread_mutex_t crash_detail_page_lock = PTHREAD_MUTEX_INITIALIZER;
   crash_detail_page_t* crash_detail_page = nullptr;
 };
diff --git a/linker/linker.cpp b/linker/linker.cpp
index a12388c38..b0caeddc1 100644
--- a/linker/linker.cpp
+++ b/linker/linker.cpp
@@ -1695,13 +1695,31 @@ bool find_libraries(android_namespace_t* ns,
   }
 
   // Step 3: pre-link all DT_NEEDED libraries in breadth first order.
+  bool any_memtag_stack = false;
   for (auto&& task : load_tasks) {
     soinfo* si = task->get_soinfo();
     if (!si->is_linked() && !si->prelink_image()) {
       return false;
     }
+    // si->memtag_stack() needs to be called after si->prelink_image() which populates
+    // the dynamic section.
+    if (si->has_min_version(7) && si->memtag_stack()) {
+      any_memtag_stack = true;
+      LD_LOG(kLogDlopen,
+             "... load_library requesting stack MTE for: realpath=\"%s\", soname=\"%s\"",
+             si->get_realpath(), si->get_soname());
+    }
     register_soinfo_tls(si);
   }
+  if (any_memtag_stack) {
+    if (auto* cb = __libc_shared_globals()->memtag_stack_dlopen_callback) {
+      cb();
+    } else {
+      // find_library is used by the initial linking step, so we communicate that we
+      // want memtag_stack enabled to __libc_init_mte.
+      __libc_shared_globals()->initial_memtag_stack = true;
+    }
+  }
 
   // Step 4: Construct the global group. DF_1_GLOBAL bit is force set for LD_PRELOADed libs because
   // they must be added to the global group. Note: The DF_1_GLOBAL bit for a library is normally set
diff --git a/linker/linker_main.cpp b/linker/linker_main.cpp
index 5f5eba468..d6592af14 100644
--- a/linker/linker_main.cpp
+++ b/linker/linker_main.cpp
@@ -404,9 +404,6 @@ static ElfW(Addr) linker_main(KernelArgumentBlock& args, const char* exe_to_load
                      strerror(errno));
     }
   }
-
-  __libc_init_mte(somain->memtag_dynamic_entries(), somain->phdr, somain->phnum, somain->load_bias,
-                  args.argv);
 #endif
 
   // Register the main executable and the linker upfront to have
@@ -496,6 +493,12 @@ static ElfW(Addr) linker_main(KernelArgumentBlock& args, const char* exe_to_load
     }
     si->increment_ref_count();
   }
+#if defined(__aarch64__)
+  // This has to happen after the find_libraries, which will have collected any possible
+  // libraries that request memtag_stack in the dynamic section.
+  __libc_init_mte(somain->memtag_dynamic_entries(), somain->phdr, somain->phnum, somain->load_bias,
+                  args.argv);
+#endif
 
   linker_finalize_static_tls();
   __libc_init_main_thread_final();
diff --git a/tests/Android.bp b/tests/Android.bp
index a62ababfa..78c2c1083 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -1108,6 +1108,36 @@ cc_test {
 }
 
 cc_test {
+    name: "memtag_stack_dlopen_test",
+    enabled: false,
+    // This does not use bionic_tests_defaults because it is not supported on
+    // host.
+    arch: {
+        arm64: {
+            enabled: true,
+        },
+    },
+    sanitize: {
+        memtag_heap: true,
+        memtag_stack: false,
+    },
+    srcs: [
+        "memtag_stack_dlopen_test.cpp",
+    ],
+    shared_libs: [
+        "libbase",
+    ],
+    data_libs: ["libtest_simple_memtag_stack", "libtest_depends_on_simple_memtag_stack"],
+    data_bins: [
+        "testbinary_depends_on_simple_memtag_stack",
+        "testbinary_depends_on_depends_on_simple_memtag_stack",
+        "testbinary_is_stack_mte_after_dlopen"
+    ],
+    header_libs: ["bionic_libc_platform_headers"],
+    test_suites: ["device-tests"],
+}
+
+cc_test {
     name: "bionic-stress-tests",
     defaults: [
         "bionic_tests_defaults",
diff --git a/tests/libs/Android.bp b/tests/libs/Android.bp
index 039d1e101..06ee132d4 100644
--- a/tests/libs/Android.bp
+++ b/tests/libs/Android.bp
@@ -234,6 +234,61 @@ cc_test_library {
 }
 
 // -----------------------------------------------------------------------------
+// Libraries and binaries used by memtag_stack_dlopen_test tests
+// -----------------------------------------------------------------------------
+cc_test_library {
+    name: "libtest_simple_memtag_stack",
+    sanitize: {
+        memtag_stack: true,
+    },
+    srcs: ["dlopen_testlib_simple.cpp"],
+}
+
+cc_test_library {
+    name: "libtest_depends_on_simple_memtag_stack",
+    sanitize: {
+        memtag_stack: false,
+    },
+    shared_libs: [
+        "libtest_simple_memtag_stack",
+    ],
+    srcs: ["dlopen_testlib_depends_on_simple.cpp"],
+}
+
+cc_binary {
+    name: "testbinary_is_stack_mte_after_dlopen",
+    sanitize: {
+        memtag_stack: false,
+        memtag_heap: true,
+    },
+    srcs: ["testbinary_is_stack_mte_after_dlopen.cpp"],
+}
+
+cc_binary {
+    name: "testbinary_depends_on_simple_memtag_stack",
+    sanitize: {
+        memtag_stack: false,
+        memtag_heap: true,
+    },
+    shared_libs: [
+        "libtest_simple_memtag_stack",
+    ],
+    srcs: ["testbinary_is_stack_mte.cpp"],
+}
+
+cc_binary {
+    name: "testbinary_depends_on_depends_on_simple_memtag_stack",
+    sanitize: {
+        memtag_stack: false,
+        memtag_heap: true,
+    },
+    shared_libs: [
+        "libtest_depends_on_simple_memtag_stack",
+    ],
+    srcs: ["testbinary_is_stack_mte.cpp"],
+}
+
+// -----------------------------------------------------------------------------
 // Libraries used by hwasan_test
 // -----------------------------------------------------------------------------
 cc_test_library {
diff --git a/tests/libs/dlopen_testlib_depends_on_simple.cpp b/tests/libs/dlopen_testlib_depends_on_simple.cpp
new file mode 100644
index 000000000..9e130d426
--- /dev/null
+++ b/tests/libs/dlopen_testlib_depends_on_simple.cpp
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <stdint.h>
+#include <stdlib.h>
+
+extern "C" bool dlopen_testlib_simple_func();
+
+extern "C" bool dlopen_testlib_call_simple_func() {
+  return dlopen_testlib_simple_func();
+}
diff --git a/tests/libs/testbinary_is_stack_mte.cpp b/tests/libs/testbinary_is_stack_mte.cpp
new file mode 100644
index 000000000..8dde83c34
--- /dev/null
+++ b/tests/libs/testbinary_is_stack_mte.cpp
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "../mte_utils.h"
+#include "CHECK.h"
+
+#if defined(__BIONIC__) && defined(__aarch64__)
+
+extern "C" int main(int, char**) {
+  int ret = is_stack_mte_on() ? 0 : 1;
+  printf("RAN\n");
+  return ret;
+}
+
+#else
+
+extern "C" int main(int, char**) {
+  printf("RAN\n");
+  return 1;
+}
+#endif
diff --git a/tests/libs/testbinary_is_stack_mte_after_dlopen.cpp b/tests/libs/testbinary_is_stack_mte_after_dlopen.cpp
new file mode 100644
index 000000000..c5e686808
--- /dev/null
+++ b/tests/libs/testbinary_is_stack_mte_after_dlopen.cpp
@@ -0,0 +1,133 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include <condition_variable>
+#include <thread>
+
+#include <dlfcn.h>
+#include <inttypes.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+#include "../mte_utils.h"
+#include "CHECK.h"
+
+#if defined(__BIONIC__) && defined(__aarch64__)
+
+enum State { kInit, kThreadStarted, kStackRemapped };
+
+// We can't use pthread_getattr_np because that uses the rlimit rather than the actual mapping
+// bounds.
+static void find_main_stack_limits(uintptr_t* low, uintptr_t* high) {
+  uintptr_t startstack = reinterpret_cast<uintptr_t>(__builtin_frame_address(0));
+
+  // Hunt for the region that contains that address.
+  FILE* fp = fopen("/proc/self/maps", "re");
+  if (fp == nullptr) {
+    abort();
+  }
+  char line[BUFSIZ];
+  while (fgets(line, sizeof(line), fp) != nullptr) {
+    uintptr_t lo, hi;
+    if (sscanf(line, "%" SCNxPTR "-%" SCNxPTR, &lo, &hi) == 2) {
+      if (lo <= startstack && startstack <= hi) {
+        *low = lo;
+        *high = hi;
+        fclose(fp);
+        return;
+      }
+    }
+  }
+  abort();
+}
+
+template <typename Fn>
+unsigned int fault_new_stack_page(uintptr_t low, Fn f) {
+  uintptr_t new_low;
+  uintptr_t new_high;
+  volatile char buf[4096];
+  buf[4095] = 1;
+  find_main_stack_limits(&new_low, &new_high);
+  if (new_low < low) {
+    f();
+    return new_high;
+  }
+  // Useless, but should defeat TCO.
+  return new_low + fault_new_stack_page(low, f);
+}
+extern "C" int main(int argc, char** argv) {
+  if (argc < 2) {
+    return 1;
+  }
+  const char* path = argv[1];
+  CHECK(access(path, F_OK) == 0);  // Verify test setup.
+  CHECK(!is_stack_mte_on());
+  std::mutex m;
+  std::condition_variable cv;
+  State state = kInit;
+
+  bool is_early_thread_mte_on = false;
+  std::thread early_th([&] {
+    {
+      std::lock_guard lk(m);
+      state = kThreadStarted;
+    }
+    cv.notify_one();
+    {
+      std::unique_lock lk(m);
+      cv.wait(lk, [&] { return state == kStackRemapped; });
+    }
+    is_early_thread_mte_on = is_stack_mte_on();
+  });
+  {
+    std::unique_lock lk(m);
+    cv.wait(lk, [&] { return state == kThreadStarted; });
+  }
+  void* handle = dlopen(path, RTLD_NOW);
+  {
+    std::lock_guard lk(m);
+    state = kStackRemapped;
+  }
+  cv.notify_one();
+  CHECK(handle != nullptr);
+  CHECK(is_stack_mte_on());
+
+  bool new_stack_page_mte_on = false;
+  uintptr_t low;
+  uintptr_t high;
+  find_main_stack_limits(&low, &high);
+  fault_new_stack_page(low, [&] { new_stack_page_mte_on = is_stack_mte_on(); });
+  CHECK(new_stack_page_mte_on);
+
+  bool is_late_thread_mte_on = false;
+  std::thread late_th([&] { is_late_thread_mte_on = is_stack_mte_on(); });
+  late_th.join();
+  early_th.join();
+  CHECK(is_late_thread_mte_on);
+  CHECK(is_early_thread_mte_on);
+  printf("RAN\n");
+  return 0;
+}
+
+#else
+extern "C" int main(int, char**) {
+  return 1;
+}
+#endif
diff --git a/tests/memtag_stack_dlopen_test.cpp b/tests/memtag_stack_dlopen_test.cpp
new file mode 100644
index 000000000..68ddb81bd
--- /dev/null
+++ b/tests/memtag_stack_dlopen_test.cpp
@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in
+ *    the documentation and/or other materials provided with the
+ *    distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+ * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+ * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+ * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+ * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+ * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
+
+#include <thread>
+
+#include <dlfcn.h>
+#include <stdlib.h>
+
+#include <gtest/gtest.h>
+
+#include <android-base/silent_death_test.h>
+#include <android-base/test_utils.h>
+#include "mte_utils.h"
+#include "utils.h"
+
+TEST(MemtagStackDlopenTest, DependentBinaryGetsMemtagStack) {
+#if defined(__BIONIC__) && defined(__aarch64__)
+  if (!running_with_mte()) GTEST_SKIP() << "Test requires MTE.";
+  if (is_stack_mte_on())
+    GTEST_SKIP() << "Stack MTE needs to be off for this test. Are you running fullmte?";
+
+  std::string path =
+      android::base::GetExecutableDirectory() + "/testbinary_depends_on_simple_memtag_stack";
+  ExecTestHelper eth;
+  std::string ld_library_path = "LD_LIBRARY_PATH=" + android::base::GetExecutableDirectory();
+  eth.SetArgs({path.c_str(), nullptr});
+  eth.SetEnv({ld_library_path.c_str(), nullptr});
+  eth.Run([&]() { execve(path.c_str(), eth.GetArgs(), eth.GetEnv()); }, 0, "RAN");
+#else
+  GTEST_SKIP() << "requires bionic arm64";
+#endif
+}
+
+TEST(MemtagStackDlopenTest, DependentBinaryGetsMemtagStack2) {
+#if defined(__BIONIC__) && defined(__aarch64__)
+  if (!running_with_mte()) GTEST_SKIP() << "Test requires MTE.";
+  if (is_stack_mte_on())
+    GTEST_SKIP() << "Stack MTE needs to be off for this test. Are you running fullmte?";
+
+  std::string path = android::base::GetExecutableDirectory() +
+                     "/testbinary_depends_on_depends_on_simple_memtag_stack";
+  ExecTestHelper eth;
+  std::string ld_library_path = "LD_LIBRARY_PATH=" + android::base::GetExecutableDirectory();
+  eth.SetArgs({path.c_str(), nullptr});
+  eth.SetEnv({ld_library_path.c_str(), nullptr});
+  eth.Run([&]() { execve(path.c_str(), eth.GetArgs(), eth.GetEnv()); }, 0, "RAN");
+#else
+  GTEST_SKIP() << "requires bionic arm64";
+#endif
+}
+
+TEST(MemtagStackDlopenTest, DlopenRemapsStack) {
+#if defined(__BIONIC__) && defined(__aarch64__)
+  // If this test is failing, look at crash logcat for why the test binary died.
+  if (!running_with_mte()) GTEST_SKIP() << "Test requires MTE.";
+  if (is_stack_mte_on())
+    GTEST_SKIP() << "Stack MTE needs to be off for this test. Are you running fullmte?";
+
+  std::string path =
+      android::base::GetExecutableDirectory() + "/testbinary_is_stack_mte_after_dlopen";
+  std::string lib_path =
+      android::base::GetExecutableDirectory() + "/libtest_simple_memtag_stack.so";
+  ExecTestHelper eth;
+  std::string ld_library_path = "LD_LIBRARY_PATH=" + android::base::GetExecutableDirectory();
+  eth.SetArgs({path.c_str(), lib_path.c_str(), nullptr});
+  eth.SetEnv({ld_library_path.c_str(), nullptr});
+  eth.Run([&]() { execve(path.c_str(), eth.GetArgs(), eth.GetEnv()); }, 0, "RAN");
+#else
+  GTEST_SKIP() << "requires bionic arm64";
+#endif
+}
+
+TEST(MemtagStackDlopenTest, DlopenRemapsStack2) {
+#if defined(__BIONIC__) && defined(__aarch64__)
+  // If this test is failing, look at crash logcat for why the test binary died.
+  if (!running_with_mte()) GTEST_SKIP() << "Test requires MTE.";
+  if (is_stack_mte_on())
+    GTEST_SKIP() << "Stack MTE needs to be off for this test. Are you running fullmte?";
+
+  std::string path =
+      android::base::GetExecutableDirectory() + "/testbinary_is_stack_mte_after_dlopen";
+  std::string lib_path =
+      android::base::GetExecutableDirectory() + "/libtest_depends_on_simple_memtag_stack.so";
+  ExecTestHelper eth;
+  std::string ld_library_path = "LD_LIBRARY_PATH=" + android::base::GetExecutableDirectory();
+  eth.SetArgs({path.c_str(), lib_path.c_str(), nullptr});
+  eth.SetEnv({ld_library_path.c_str(), nullptr});
+  eth.Run([&]() { execve(path.c_str(), eth.GetArgs(), eth.GetEnv()); }, 0, "RAN");
+#else
+  GTEST_SKIP() << "requires bionic arm64";
+#endif
+}
diff --git a/tests/mte_utils.h b/tests/mte_utils.h
new file mode 100644
index 000000000..0f18442fa
--- /dev/null
+++ b/tests/mte_utils.h
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#if defined(__BIONIC__) && defined(__aarch64__)
+
+__attribute__((target("mte"))) static bool is_stack_mte_on() {
+  alignas(16) int x = 0;
+  void* p = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(&x) + (1UL << 57));
+  void* p_cpy = p;
+  __builtin_arm_stg(p);
+  p = __builtin_arm_ldg(p);
+  __builtin_arm_stg(&x);
+  return p == p_cpy;
+}
+
+#endif