Merge cherrypicks of ['googleplex-android-review.googlesource.com/33294988', 'googleplex-android-review.googlesource.com/33313547'] into 25Q2-release. Change-Id: I4cbf70423928c5722ce5a2d8738bff3eb27d3c67

diff --git a/libc/arch-arm64/bionic/setjmp.S b/libc/arch-arm64/bionic/setjmp.S
index c408998..e94e5f4 100644
--- a/libc/arch-arm64/bionic/setjmp.S
+++ b/libc/arch-arm64/bionic/setjmp.S

@@ -114,6 +114,9 @@
   .cfi_rel_offset x0, 0
   .cfi_rel_offset x30, 8
 
+  // Commit SME's ZA lazy save. Note that the call preserves x1.
+  bl __arm_za_disable
+
   // Get the cookie and store it along with the signal flag.
   mov x0, x1
   bl __bionic_setjmp_cookie_get
@@ -183,6 +186,17 @@
 
 // void siglongjmp(sigjmp_buf env, int value);
 ENTRY_WEAK_FOR_NATIVE_BRIDGE(siglongjmp)
+  // First of all, disable SME's ZA, so that it does not interfere
+  // with anything else. Note that __arm_za_disable is guaranteed to
+  // preserve x0 and x1.
+  str x30, [sp, #-16]!
+  .cfi_adjust_cfa_offset 16
+  .cfi_rel_offset x30, 0
+  bl __arm_za_disable
+  ldr x30, [sp], #16
+  .cfi_adjust_cfa_offset -16
+  .cfi_restore x30
+
   // Check the checksum before doing anything.
   m_calculate_checksum x12, x0, x2
   ldr x2, [x0, #(_JB_CHECKSUM * 8)]

diff --git a/tests/setjmp_test.cpp b/tests/setjmp_test.cpp
index 9469285..836aadc 100644
--- a/tests/setjmp_test.cpp
+++ b/tests/setjmp_test.cpp

@@ -18,6 +18,7 @@
 
 #include <setjmp.h>
 #include <stdlib.h>
+#include <sys/auxv.h>
 #include <sys/syscall.h>
 #include <unistd.h>
 
@@ -364,3 +365,42 @@
   GTEST_SKIP() << "tests uses functions not in glibc";
 #endif
 }
+
+#if defined(__aarch64__)
+TEST(setjmp, sigsetjmp_sme) {
+  if (!(getauxval(AT_HWCAP2) & HWCAP2_SME)) {
+    GTEST_SKIP() << "SME is not enabled on device.";
+  }
+
+  uint64_t svcr, za_state;
+  sigjmp_buf jb;
+  __asm__ __volatile__(".arch_extension sme; smstart za");
+  sigsetjmp(jb, 0);
+  __asm__ __volatile__(".arch_extension sme; mrs %0, SVCR" : "=r"(svcr));
+  __asm__ __volatile__(".arch_extension sme; smstop za");  // Turn ZA off anyway.
+  za_state = svcr & 0x2UL;
+  ASSERT_EQ(0UL, za_state);
+}
+
+TEST(setjmp, siglongjmp_sme) {
+  if (!(getauxval(AT_HWCAP2) & HWCAP2_SME)) {
+    GTEST_SKIP() << "SME is not enabled on device.";
+  }
+
+  uint64_t svcr, za_state;
+  int value;
+  sigjmp_buf jb;
+  if ((value = sigsetjmp(jb, 0)) == 0) {
+    __asm__ __volatile__(".arch_extension sme; smstart za");
+    siglongjmp(jb, 789);
+    __asm__ __volatile__(".arch_extension sme; smstop za");
+    FAIL();  // Unreachable.
+  } else {
+    __asm__ __volatile__(".arch_extension sme; mrs %0, SVCR" : "=r"(svcr));
+    __asm__ __volatile__(".arch_extension sme; smstop za");  // Turn ZA off anyway.
+    za_state = svcr & 0x2UL;
+    ASSERT_EQ(789, value);
+    ASSERT_EQ(0UL, za_state);
+  }
+}
+#endif

diff --git a/tests/signal_test.cpp b/tests/signal_test.cpp
index f8bfcef..27f5c6c 100644
--- a/tests/signal_test.cpp
+++ b/tests/signal_test.cpp

@@ -16,6 +16,7 @@
 
 #include <errno.h>
 #include <signal.h>
+#include <sys/auxv.h>
 #include <sys/cdefs.h>
 #include <sys/syscall.h>
 #include <sys/types.h>
@@ -1070,3 +1071,35 @@
   GTEST_SKIP() << "our old glibc doesn't have str2sig";
 #endif
 }
+
+#if defined(__aarch64__)
+__attribute__((target("arch=armv9+sme"))) __arm_new("za") static void FunctionUsingZA() {
+  raise(SIGUSR1);
+}
+
+TEST(signal, sme_tpidr2_clear) {
+  // When using SME, on entering a signal handler the kernel should clear TPIDR2_EL0, but this was
+  // not always correctly done. This tests checks if the kernel correctly clears it or not.
+  if (!(getauxval(AT_HWCAP2) & HWCAP2_SME)) {
+    GTEST_SKIP() << "SME is not enabled on device.";
+  }
+
+  static uint64_t tpidr2 = 0;
+  struct sigaction handler = {};
+  handler.sa_sigaction = [](int, siginfo_t*, void*) {
+    uint64_t zero = 0;
+    __asm__ __volatile__(".arch_extension sme; mrs %0, TPIDR2_EL0" : "=r"(tpidr2));
+    __asm__ __volatile__(".arch_extension sme; msr TPIDR2_EL0, %0" : : "r"(zero));  // Clear TPIDR2.
+  };
+  handler.sa_flags = SA_SIGINFO;
+
+  ASSERT_EQ(0, sigaction(SIGUSR1, &handler, nullptr));
+
+  FunctionUsingZA();
+
+  ASSERT_EQ(0x0UL, tpidr2)
+      << "Broken kernel! TPIDR2_EL0 was not null in the signal handler! "
+      << "Please make sure the following patch has been applied to the kernel: "
+      << "https://lore.kernel.org/linux-arm-kernel/20250417190113.3778111-1-mark.rutland@arm.com/";
+}
+#endif