Revert^2 "memtag_stack tests."

8162b05ccdc90722499eb2f55f0bbb0d65d7f044

Change-Id: I68c1988b0d76dddfaf69189cfd439192cabda00d
diff --git a/tests/Android.bp b/tests/Android.bp
index 8f4ba2f..71be058 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -393,6 +393,7 @@
         "math_test.cpp",
         "math_force_long_double_test.cpp",
         "membarrier_test.cpp",
+        "memtag_stack_test.cpp",
         "mntent_test.cpp",
         "mte_test.cpp",
         "netdb_test.cpp",
@@ -913,6 +914,8 @@
         "heap_tagging_static_disabled_helper",
         "heap_tagging_static_sync_helper",
         "heap_tagging_sync_helper",
+        "stack_tagging_helper",
+        "stack_tagging_static_helper",
         "ld_config_test_helper",
         "ld_config_test_helper_lib1",
         "ld_config_test_helper_lib2",
@@ -1161,6 +1164,8 @@
         "heap_tagging_static_disabled_helper",
         "heap_tagging_static_sync_helper",
         "heap_tagging_sync_helper",
+        "stack_tagging_helper",
+        "stack_tagging_static_helper",
     ],
 }
 
diff --git a/tests/libs/Android.bp b/tests/libs/Android.bp
index 0046ef6..5edf97c 100644
--- a/tests/libs/Android.bp
+++ b/tests/libs/Android.bp
@@ -1652,6 +1652,33 @@
    },
 }
 
+cc_test {
+   name: "stack_tagging_helper",
+   defaults: ["bionic_testlib_defaults", "bionic_targets_only"],
+   srcs: ["stack_tagging_helper.cpp"],
+   sanitize: {
+     memtag_heap: true,
+     memtag_stack: true,
+     diag: {
+       memtag_heap: true,
+     },
+   },
+}
+
+cc_test {
+   name: "stack_tagging_static_helper",
+   defaults: ["bionic_testlib_defaults", "bionic_targets_only"],
+   srcs: ["stack_tagging_helper.cpp"],
+   static_executable: true,
+   sanitize: {
+     memtag_heap: true,
+     memtag_stack: true,
+     diag: {
+       memtag_heap: true,
+     },
+   },
+}
+
 cc_genrule {
     name: "libdlext_test_zip_zipaligned",
      out: ["bionic-loader-test-libs/libdlext_test_zip/libdlext_test_zip_zipaligned.zip"],
diff --git a/tests/libs/stack_tagging_helper.cpp b/tests/libs/stack_tagging_helper.cpp
new file mode 100644
index 0000000..f6739b6
--- /dev/null
+++ b/tests/libs/stack_tagging_helper.cpp
@@ -0,0 +1,294 @@
+/*
+ * Copyright (C) 2022 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 <errno.h>
+#include <setjmp.h>
+#include <signal.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+#include <thread>
+
+#include "libs_utils.h"
+
+#if defined(__aarch64__)
+
+template <typename T>
+static inline void mte_set_tag(T* p) {
+  __asm__ __volatile__(
+      ".arch_extension memtag\n"
+      "stg %[Ptr], [%[Ptr]]\n"
+      :
+      : [Ptr] "r"(p)
+      : "memory");
+}
+
+template <typename T>
+static inline T* mte_get_tag(T* p) {
+  __asm__ __volatile__(
+      ".arch_extension memtag\n"
+      "ldg %[Ptr], [%[Ptr]]\n"
+      : [Ptr] "+r"(p)
+      :
+      : "memory");
+  return p;
+}
+
+template <typename T>
+static inline T* mte_increment_tag(T* p) {
+  T* res;
+  __asm__ __volatile__(
+      ".arch_extension memtag\n"
+      "addg %[Res], %[Ptr], #0, #1\n"
+      : [Res] "=r"(res)
+      : [Ptr] "r"(p)
+      : "memory");
+  return res;
+}
+
+constexpr size_t kStackAllocationSize = 128 * 1024;
+
+// Prevent optimizations.
+volatile void* sink;
+
+enum struct ChildAction { Exit, Execve, Execl };
+
+// Either execve or _exit, transferring control back to parent.
+__attribute__((no_sanitize("memtag"), optnone, noinline)) void vfork_child2(ChildAction action,
+                                                                            void* fp_parent) {
+  // Make sure that the buffer in the caller has not been optimized out.
+  void* fp = __builtin_frame_address(0);
+  CHECK(reinterpret_cast<uintptr_t>(fp_parent) - reinterpret_cast<uintptr_t>(fp) >=
+        kStackAllocationSize);
+  if (action == ChildAction::Execve) {
+    const char* argv[] = {"/system/bin/true", nullptr};
+    const char* envp[] = {nullptr};
+    execve("/system/bin/true", const_cast<char**>(argv), const_cast<char**>(envp));
+    fprintf(stderr, "execve failed: %m\n");
+    _exit(1);
+  } else if (action == ChildAction::Execl) {
+    execl("/system/bin/true", "/system/bin/true", "unusedA", "unusedB", nullptr);
+    fprintf(stderr, "execl failed: %m\n");
+    _exit(1);
+  } else if (action == ChildAction::Exit) {
+    _exit(0);
+  }
+  CHECK(0);
+}
+
+// Place a tagged buffer on the stack. Do not tag the top half so that the parent does not crash too
+// early even if things go wrong.
+__attribute__((no_sanitize("memtag"), optnone, noinline)) void vfork_child(ChildAction action) {
+  alignas(16) char buf[kStackAllocationSize] __attribute__((uninitialized));
+  sink = &buf;
+
+  for (char* p = buf; p < buf + sizeof(buf) / 2; p += 16) {
+    char* q = mte_increment_tag(p);
+    mte_set_tag(q);
+    CHECK(mte_get_tag(p) == q);
+  }
+  vfork_child2(action, __builtin_frame_address(0));
+}
+
+// Parent. Check that the stack has correct allocation tags.
+__attribute__((no_sanitize("memtag"), optnone, noinline)) void vfork_parent(pid_t pid) {
+  alignas(16) char buf[kStackAllocationSize] __attribute__((uninitialized));
+  fprintf(stderr, "vfork_parent %p\n", &buf);
+  bool success = true;
+  for (char* p = buf; p < buf + sizeof(buf); p += 16) {
+    char* q = mte_get_tag(p);
+    if (p != q) {
+      fprintf(stderr, "tag mismatch at offset %zx: %p != %p\n", p - buf, p, q);
+      success = false;
+      break;
+    }
+  }
+
+  int wstatus;
+  do {
+    int res = waitpid(pid, &wstatus, 0);
+    CHECK(res == pid);
+  } while (!WIFEXITED(wstatus) && !WIFSIGNALED(wstatus));
+
+  CHECK(WIFEXITED(wstatus));
+  CHECK(WEXITSTATUS(wstatus) == 0);
+
+  if (!success) exit(1);
+}
+
+void test_vfork(ChildAction action) {
+  pid_t pid = vfork();
+  if (pid == 0) {
+    vfork_child(action);
+  } else {
+    vfork_parent(pid);
+  }
+}
+
+__attribute__((no_sanitize("memtag"), optnone, noinline)) static void settag_and_longjmp(
+    jmp_buf cont) {
+  alignas(16) char buf[kStackAllocationSize] __attribute__((uninitialized));
+  sink = &buf;
+
+  for (char* p = buf; p < buf + sizeof(buf) / 2; p += 16) {
+    char* q = mte_increment_tag(p);
+    mte_set_tag(q);
+    if (mte_get_tag(p) != q) {
+      fprintf(stderr, "failed to set allocation tags on stack: %p != %p\n", mte_get_tag(p), q);
+      exit(1);
+    }
+  }
+  longjmp(cont, 42);
+}
+
+// Check that the stack has correct allocation tags.
+__attribute__((no_sanitize("memtag"), optnone, noinline)) static void check_stack_tags() {
+  alignas(16) char buf[kStackAllocationSize] __attribute__((uninitialized));
+  for (char* p = buf; p < buf + sizeof(buf); p += 16) {
+    void* q = mte_get_tag(p);
+    if (p != q) {
+      fprintf(stderr, "stack tags mismatch: expected %p, got %p", p, q);
+      exit(1);
+    }
+  }
+}
+
+void check_longjmp_restores_tags() {
+  int value;
+  jmp_buf jb;
+  if ((value = setjmp(jb)) == 0) {
+    settag_and_longjmp(jb);
+    exit(2);  // Unreachable.
+  } else {
+    CHECK(value == 42);
+    check_stack_tags();
+  }
+}
+
+class SigAltStackScoped {
+  stack_t old_ss;
+  void* altstack_start;
+  size_t altstack_size;
+
+ public:
+  SigAltStackScoped(size_t sz) : altstack_size(sz) {
+    altstack_start = mmap(nullptr, altstack_size, PROT_READ | PROT_WRITE | PROT_MTE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
+    if (altstack_start == MAP_FAILED) {
+      fprintf(stderr, "sigaltstack mmap failed: %m\n");
+      exit(1);
+    }
+    stack_t ss = {};
+    ss.ss_sp = altstack_start;
+    ss.ss_size = altstack_size;
+    int res = sigaltstack(&ss, &old_ss);
+    CHECK(res == 0);
+  }
+
+  ~SigAltStackScoped() {
+    int res = sigaltstack(&old_ss, nullptr);
+    CHECK(res == 0);
+    munmap(altstack_start, altstack_size);
+  }
+};
+
+class SigActionScoped {
+  int signo;
+  struct sigaction oldsa;
+
+ public:
+  using handler_t = void (*)(int, siginfo_t* siginfo, void*);
+
+  SigActionScoped(int signo, handler_t handler) : signo(signo) {
+    struct sigaction sa = {};
+    sa.sa_sigaction = handler;
+    sa.sa_flags = SA_SIGINFO | SA_ONSTACK;
+    int res = sigaction(signo, &sa, &oldsa);
+    CHECK(res == 0);
+  }
+
+  ~SigActionScoped() {
+    int res = sigaction(signo, &oldsa, nullptr);
+    CHECK(res == 0);
+  }
+};
+
+void test_longjmp() {
+  check_longjmp_restores_tags();
+
+  std::thread t([]() { check_longjmp_restores_tags(); });
+  t.join();
+}
+
+void test_longjmp_sigaltstack() {
+  constexpr size_t kAltStackSize = kStackAllocationSize + PAGE_SIZE * 16;
+  SigAltStackScoped sigAltStackScoped(kAltStackSize);
+  SigActionScoped sigActionScoped(
+      SIGUSR1, [](int, siginfo_t*, void*) { check_longjmp_restores_tags(); });
+  raise(SIGUSR1);
+
+  // same for a secondary thread
+  std::thread t([]() {
+    SigAltStackScoped sigAltStackScoped(kAltStackSize);
+    raise(SIGUSR1);
+  });
+  t.join();
+}
+
+int main(int argc, char** argv) {
+  if (argc < 2) {
+    printf("nothing to do\n");
+    return 1;
+  }
+
+  if (strcmp(argv[1], "vfork_execve") == 0) {
+    test_vfork(ChildAction::Execve);
+    return 0;
+  }
+
+  if (strcmp(argv[1], "vfork_execl") == 0) {
+    test_vfork(ChildAction::Execl);
+    return 0;
+  }
+
+  if (strcmp(argv[1], "vfork_exit") == 0) {
+    test_vfork(ChildAction::Exit);
+    return 0;
+  }
+
+  if (strcmp(argv[1], "longjmp") == 0) {
+    test_longjmp();
+    return 0;
+  }
+
+  if (strcmp(argv[1], "longjmp_sigaltstack") == 0) {
+    test_longjmp_sigaltstack();
+    return 0;
+  }
+
+  printf("unrecognized command: %s\n", argv[1]);
+  return 1;
+}
+#else
+int main(int, char**) {
+  printf("aarch64 only\n");
+  return 1;
+}
+#endif  // defined(__aarch64__)
diff --git a/tests/memtag_stack_test.cpp b/tests/memtag_stack_test.cpp
new file mode 100644
index 0000000..9ff8899
--- /dev/null
+++ b/tests/memtag_stack_test.cpp
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2022 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 <tuple>
+
+#include <gtest/gtest.h>
+
+#if defined(__BIONIC__)
+#include "gtest_globals.h"
+#include "platform/bionic/mte.h"
+#include "utils.h"
+#endif
+
+class MemtagStackTest : public testing::TestWithParam<std::tuple<const char*, bool>> {};
+
+TEST_P(MemtagStackTest, test) {
+#if defined(__BIONIC__) && defined(__aarch64__)
+  if (!mte_supported()) {
+    return;
+  }
+  bool is_static = std::get<1>(GetParam());
+  std::string helper =
+      GetTestlibRoot() + (is_static ? "/stack_tagging_static_helper" : "/stack_tagging_helper");
+  const char* arg = std::get<0>(GetParam());
+  chmod(helper.c_str(), 0755);
+  ExecTestHelper eth;
+  eth.SetArgs({helper.c_str(), arg, nullptr});
+  eth.Run([&]() { execve(helper.c_str(), eth.GetArgs(), eth.GetEnv()); }, 0, "");
+#else
+  GTEST_SKIP() << "bionic/arm64 only";
+#endif
+}
+
+INSTANTIATE_TEST_SUITE_P(, MemtagStackTest,
+                         testing::Combine(testing::Values("vfork_execve", "vfork_execl",
+                                                          "vfork_exit", "longjmp",
+                                                          "longjmp_sigaltstack"),
+                                          testing::Bool()),
+                         [](const ::testing::TestParamInfo<MemtagStackTest::ParamType>& info) {
+                           std::string s = std::get<0>(info.param);
+                           if (std::get<1>(info.param)) s += "_static";
+                           return s;
+                         });