[libc] add checksum for jmpbuf

libc/include/llvm-libc-types/jmp_buf.h

@@ -38,6 +38,10 @@ typedef struct {
 #else
 #error "__jmp_buf not available for your target architecture."
 #endif
+  __UINT64_TYPE__ __sigmask;
+  __UINT64_TYPE__ __has_sigmask : 1;
+  __UINT64_TYPE__ __unused : 63;
+  __UINT64_TYPE__ __chksum;
 } __jmp_buf;
 
 typedef __jmp_buf jmp_buf[1];

libc/src/setjmp/CMakeLists.txt

@@ -2,6 +2,16 @@ if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_ARCHITECTURE})
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_ARCHITECTURE})
 endif()
 
+add_header_library(
+  checksum
+  HDRS
+    checksum.h
+  DEPENDS
+    libc.src.__support.hash
+    libc.src.stdlib.abort
+    libc.src.unistd.write
+)
+
 add_entrypoint_object(
   setjmp
   ALIAS

libc/src/setjmp/checksum.h

@@ -0,0 +1,67 @@
+//===-- Implementation header for jmpbuf checksum ---------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_SETJMP_CHECKSUM_H
+#define LLVM_LIBC_SRC_SETJMP_CHECKSUM_H
+
+#include "src/__support/hash.h"
+#include "src/__support/macros/attributes.h"
+#include "src/__support/macros/config.h"
+#include "src/setjmp/setjmp_impl.h"
+#include "src/stdlib/abort.h"
+#include "src/unistd/write.h"
+
+namespace LIBC_NAMESPACE_DECL {
+
+namespace jmpbuf {
+using HashState = internal::HashState;
+// Initial values generated by
+// https://www.random.org/cgi-bin/randbyte?nbytes=48&format=h
+// These values are only used for overlay targets.
+LIBC_INLINE_VAR uint64_t register_mangle_cookie = 0xdf8a883867040cbc;
+LIBC_INLINE_VAR uint64_t checksum_mangle_cookie = 0x9ed4fe406ebe9cf9;
+LIBC_INLINE_VAR uint64_t randomness[4] = {
+    0x83b9df7dddf5ab3d,
+    0x06c931cca75e15c6,
+    0x08280ec9e9a778bf,
+    0x111f67f4aafc9276,
+};
+
+LIBC_INLINE int update_checksum(__jmp_buf *buf) {
+  HashState state{
+      randomness[0],
+      randomness[1],
+      randomness[2],
+      randomness[3],
+  };
+  state.update(buf, offsetof(__jmp_buf, __chksum));
+  buf->__chksum = state.finish() ^ checksum_mangle_cookie;
+  return 0;
+}
+
+LIBC_INLINE void verify(const __jmp_buf *buf) {
+  HashState state{
+      randomness[0],
+      randomness[1],
+      randomness[2],
+      randomness[3],
+  };
+  state.update(buf, offsetof(__jmp_buf, __chksum));
+  auto chksum = state.finish() ^ checksum_mangle_cookie;
+  if (chksum != buf->__chksum) {
+    constexpr char MSG[] = "jump buffer corrupted\n";
+    LIBC_NAMESPACE::write(2, MSG, sizeof(MSG) - 1);
+    LIBC_NAMESPACE::abort();
+  }
+}
+
+} // namespace jmpbuf
+
+} // namespace LIBC_NAMESPACE_DECL
+
+#endif // LLVM_LIBC_SRC_SETJMP_CHECKSUM_H

libc/src/setjmp/x86_64/CMakeLists.txt

@@ -6,14 +6,11 @@ add_entrypoint_object(
     ../setjmp_impl.h
   DEPENDS
     libc.include.setjmp
+    libc.src.setjmp.checksum
   COMPILE_OPTIONS
     -O3
     -fno-omit-frame-pointer
-    # TODO: Remove once one of these lands:
-    # https://github.com/llvm/llvm-project/pull/87837
-    # https://github.com/llvm/llvm-project/pull/88054
-    # https://github.com/llvm/llvm-project/pull/88157
-    -ftrivial-auto-var-init=uninitialized
+    -momit-leaf-frame-pointer
 )
 
 add_entrypoint_object(
@@ -24,6 +21,7 @@ add_entrypoint_object(
     ../longjmp.h
   DEPENDS
     libc.include.setjmp
+    libc.src.setjmp.checksum
   COMPILE_OPTIONS
     -O3
     -fomit-frame-pointer

libc/src/setjmp/x86_64/longjmp.cpp

@@ -7,39 +7,69 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/setjmp/longjmp.h"
+#include "include/llvm-libc-types/jmp_buf.h"
 #include "src/__support/common.h"
 #include "src/__support/macros/config.h"
+#include "src/setjmp/checksum.h"
 
 #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 * buf, int val)) {
-  register __UINT64_TYPE__ rbx __asm__("rbx");
-  register __UINT64_TYPE__ rbp __asm__("rbp");
-  register __UINT64_TYPE__ r12 __asm__("r12");
-  register __UINT64_TYPE__ r13 __asm__("r13");
-  register __UINT64_TYPE__ r14 __asm__("r14");
-  register __UINT64_TYPE__ r15 __asm__("r15");
-  register __UINT64_TYPE__ rsp __asm__("rsp");
-  register __UINT64_TYPE__ rax __asm__("rax");
-
-  // ABI requires that the return value should be stored in rax. So, we store
-  // |val| in rax. Note that this has to happen before we restore the registers
-  // from values in |buf|. Otherwise, once rsp and rbp are updated, we cannot
-  // read |val|.
-  val = val == 0 ? 1 : val;
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(rax) : "m"(val) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(rbx) : "m"(buf->rbx) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(rbp) : "m"(buf->rbp) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(r12) : "m"(buf->r12) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(r13) : "m"(buf->r13) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(r14) : "m"(buf->r14) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(r15) : "m"(buf->r15) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=r"(rsp) : "m"(buf->rsp) :);
-  LIBC_INLINE_ASM("jmp *%0\n\t" : : "m"(buf->rip));
+  asm(R"(
+   pushq %%rbp
+   pushq %%rbx
+   mov  %%rdi, %%rbp
+   mov  %%esi, %%ebx
+   subq $8, %%rsp
+   call %P0
+   addq $8, %%rsp
+   mov  %%ebx, %%esi
+   mov  %%rbp, %%rdi
+   popq %%rbx
+   popq %%rbp
+ )" ::"i"(jmpbuf::verify)
+      : "rax", "rcx", "rdx", "r8", "r9", "r10", "r11");
+
+  register __UINT64_TYPE__ rcx __asm__("rcx");
+  // Load cookie
+  asm("mov %1, %0\n\t" : "=r"(rcx) : "m"(jmpbuf::register_mangle_cookie));
+
+  // load registers from buffer
+  // do not pass any invalid values into registers
+#define RECOVER(REG)                                                           \
+  register __UINT64_TYPE__ REG __asm__(#REG);                                  \
+  asm volatile("mov %c[" #REG "](%%rdi), %%rdx\n\t"                            \
+               "xor %%rdx, %1\n\t"                                             \
+               "mov %%rdx, %0\n\t"                                             \
+               : "=r"(REG)                                                     \
+               : "r"(rcx), [REG] "i"(offsetof(__jmp_buf, REG))                 \
+               : "rdx");
+
+  RECOVER(rbx);
+  RECOVER(rbp);
+  RECOVER(r12);
+  RECOVER(r13);
+  RECOVER(r14);
+  RECOVER(r15);
+  RECOVER(rsp);
+
+  register int eax __asm__("eax");
+  asm volatile(R"(
+   xor %0,%0
+   cmp $1,%%esi       
+   adc %%esi,%0
+   mov %c[rip](%%rdi),%%rdx
+   xor %%rdx, %%rcx
+   jmp *%%rdx
+ )"
+               : "=r"(eax)
+               : [rip] "i"(offsetof(__jmp_buf, rip))
+               : "rdx");
 }
 
 } // namespace LIBC_NAMESPACE_DECL

libc/src/setjmp/x86_64/setjmp.cpp

@@ -8,6 +8,7 @@
 
 #include "src/__support/common.h"
 #include "src/__support/macros/config.h"
+#include "src/setjmp/checksum.h"
 #include "src/setjmp/setjmp_impl.h"
 
 #if !defined(LIBC_TARGET_ARCH_IS_X86_64)
@@ -16,42 +17,39 @@
 
 namespace LIBC_NAMESPACE_DECL {
 
+[[gnu::naked]]
 LLVM_LIBC_FUNCTION(int, setjmp, (__jmp_buf * buf)) {
-  register __UINT64_TYPE__ rbx __asm__("rbx");
-  register __UINT64_TYPE__ r12 __asm__("r12");
-  register __UINT64_TYPE__ r13 __asm__("r13");
-  register __UINT64_TYPE__ r14 __asm__("r14");
-  register __UINT64_TYPE__ r15 __asm__("r15");
-
-  // We want to store the register values as is. So, we will suppress the
-  // compiler warnings about the uninitialized variables declared above.
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wuninitialized"
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=m"(buf->rbx) : "r"(rbx) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=m"(buf->r12) : "r"(r12) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=m"(buf->r13) : "r"(r13) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=m"(buf->r14) : "r"(r14) :);
-  LIBC_INLINE_ASM("mov %1, %0\n\t" : "=m"(buf->r15) : "r"(r15) :);
-#pragma GCC diagnostic pop
-
-  // We want the rbp of the caller, which is what __builtin_frame_address(1)
-  // should return. But, compilers generate a warning that calling
-  // __builtin_frame_address with non-zero argument is unsafe. So, we use
-  // the knowledge of the x86_64 ABI to fetch the callers rbp. As per the ABI,
-  // the rbp of the caller is pushed on to the stack and then new top is saved
-  // in this function's rbp. So, we fetch it from location at which this
-  // functions's rbp is pointing.
-  buf->rbp = *reinterpret_cast<__UINTPTR_TYPE__ *>(__builtin_frame_address(0));
-
-  // The callers stack address is exactly 2 pointer widths ahead of the current
-  // frame pointer - between the current frame pointer and the rsp of the caller
-  // are the return address (pushed by the x86_64 call instruction) and the
-  // previous stack pointer as required by the x86_64 ABI.
-  // The stack pointer is ahead because the stack grows down on x86_64.
-  buf->rsp = reinterpret_cast<__UINTPTR_TYPE__>(__builtin_frame_address(0)) +
-             sizeof(__UINTPTR_TYPE__) * 2;
-  buf->rip = reinterpret_cast<__UINTPTR_TYPE__>(__builtin_return_address(0));
-  return 0;
+  register __UINT64_TYPE__ rcx __asm__("rcx");
+  // Load cookie
+  asm("mov %1, %0\n\t" : "=r"(rcx) : "m"(jmpbuf::register_mangle_cookie));
+  // store registers to buffer
+  // do not pass any invalid values into registers
+#define STORE(REG)                                                             \
+  asm("mov %%" #REG ", %%rdx\n\t"                                              \
+      "xor %%rdx, %%rcx\n\t"                                                   \
+      "mov %%rdx, %c[" #REG                                                    \
+      "](%%rdi)\n\t" ::[REG] "i"(offsetof(__jmp_buf, REG))                     \
+      : "rdx");
+
+  STORE(rbx);
+  STORE(rbp);
+  STORE(r12);
+  STORE(r13);
+  STORE(r14);
+  STORE(r15);
+  asm(R"(
+   lea 8(%%rsp),%%rdx
+   xor %%rdx, %%rcx
+   mov %%rdx,%c[rsp](%%rdi)
+   mov (%%rsp),%%rdx
+   xor %%rdx, %%rcx     
+   mov %%rdx,%c[rip](%%rdi)
+ )" ::[rsp] "i"(offsetof(__jmp_buf, rsp)),
+      [rip] "i"(offsetof(__jmp_buf, rip))
+      : "rdx");
+
+  // tail call to update checksum
+  asm("jmp %P0" : : "i"(jmpbuf::update_checksum));
 }
 
 } // namespace LIBC_NAMESPACE_DECL