[libc-commits] [libc] [libc] fortify jmp buffer for x86-64 (PR #112769)

[Nick Desaulniers via libc-commits]

================
@@ -11,32 +11,82 @@
 #include "src/__support/common.h"
 #include "src/__support/macros/config.h"
 
+#if LIBC_COPT_SETJMP_FORTIFICATION
+#include "src/setjmp/checksum.h"
+#endif
+
 #if !defined(LIBC_TARGET_ARCH_IS_X86_64)
 #error "Invalid file include"
 #endif
 
 namespace LIBC_NAMESPACE_DECL {
 
-[[gnu::naked]]
-LLVM_LIBC_FUNCTION(void, longjmp, (jmp_buf, int)) {
-  asm(R"(
-      cmpl $0x1, %%esi
-      adcl $0x0, %%esi
-      movq %%rsi, %%rax
-
-      movq %c[rbx](%%rdi), %%rbx
-      movq %c[rbp](%%rdi), %%rbp
-      movq %c[r12](%%rdi), %%r12
-      movq %c[r13](%%rdi), %%r13
-      movq %c[r14](%%rdi), %%r14
-      movq %c[r15](%%rdi), %%r15
-      movq %c[rsp](%%rdi), %%rsp
-      jmpq *%c[rip](%%rdi)
-      )" ::[rbx] "i"(offsetof(__jmp_buf, rbx)),
-      [rbp] "i"(offsetof(__jmp_buf, rbp)), [r12] "i"(offsetof(__jmp_buf, r12)),
-      [r13] "i"(offsetof(__jmp_buf, r13)), [r14] "i"(offsetof(__jmp_buf, r14)),
-      [r15] "i"(offsetof(__jmp_buf, r15)), [rsp] "i"(offsetof(__jmp_buf, rsp)),
-      [rip] "i"(offsetof(__jmp_buf, rip)));
+#define CALCULATE_RETURN_VALUE()                                               \
+  "cmpl $0x1, %%esi\n\t"                                                       \
+  "adcl $0x0, %%esi\n\t"                                                       \
+  "movq %%rsi, %%rax\n\t"
+
+#if LIBC_COPT_SETJMP_FORTIFICATION
+#include "src/setjmp/x86_64/checksum.def"
+
+// clang-format off
+#define RESTORE_REG(DST)                                                       \
+  "movq %c[" #DST "](%%rdi), %%rax\n\t"                                        \
+  "movq %%rax, %%" #DST "\n\t"                                                 \
+  "xor %[mask], %%" #DST "\n\t"                                                \
+  ACCUMULATE_CHECKSUM()
+
+#define RESTORE_RIP()                                                          \
+  "movq %c[rip](%%rdi), %%rax\n\t"                                             \
+  "xor %%rax, %[mask]\n\t"                                                     \
+  ACCUMULATE_CHECKSUM()                                                        \
+  "cmp %c[__chksum](%%rdi), %%rdx\n\t"                                         \
+  "jne __libc_jmpbuf_corruption\n\t"                                           \
+  CALCULATE_RETURN_VALUE()                                                     \
+  "jmp *%[mask]\n\t"
+// clang-format on
+#else
+#define LOAD_CHKSUM_STATE_REGISTERS()
+#define RESTORE_REG(DST) "movq %c[" #DST "](%%rdi), %%" #DST "\n\t"
+#define RESTORE_RIP()                                                          \
+  CALCULATE_RETURN_VALUE()                                                     \
+  "jmpq *%c[rip](%%rdi)\n\t"
+#endif
+
+[[gnu::naked]] LLVM_LIBC_FUNCTION(void, longjmp, (jmp_buf, int)) {
+  // use registers to make sure values propagate correctly across the asm blocks
+  [[maybe_unused]] register __UINTPTR_TYPE__ mask asm("rcx");
+  [[maybe_unused]] register __UINT64_TYPE__ checksum asm("rdx");
+
+  LOAD_CHKSUM_STATE_REGISTERS()
+  asm volatile(
+      // clang-format off
+      RESTORE_REG(rbx)
+      RESTORE_REG(rbp)
+      RESTORE_REG(r12)
+      RESTORE_REG(r13)
+      RESTORE_REG(r14)
+      RESTORE_REG(r15)
+      RESTORE_REG(rsp)
+      RESTORE_RIP()
+      // clang-format on
+      : /* outputs */
+#if LIBC_COPT_SETJMP_FORTIFICATION
+      [mask] "+r"(mask), [checksum] "+r"(checksum)
----------------
nickdesaulniers wrote:

Does this work? (using `+c` and `+d` rather than declaring variables with explicit register storage)?
```suggestion
  LOAD_CHKSUM_STATE_REGISTERS()
  asm volatile(
      // clang-format off
      RESTORE_REG(rbx)
      RESTORE_REG(rbp)
      RESTORE_REG(r12)
      RESTORE_REG(r13)
      RESTORE_REG(r14)
      RESTORE_REG(r15)
      RESTORE_REG(rsp)
      RESTORE_RIP()
      // clang-format on
      : /* outputs */
#if LIBC_COPT_SETJMP_FORTIFICATION
      [mask] "+c"(mask), [checksum] "+d"(checksum)
```
Though I'm also curious why you're using rdx and rcx specifically?

https://github.com/llvm/llvm-project/pull/112769