[libc-commits] [libc] [libc] add checksum for jmpbuf (PR #101110)

Schrodinger ZHU Yifan via libc-commits libc-commits at lists.llvm.org
Tue Jul 30 14:07:40 PDT 2024


https://github.com/SchrodingerZhu updated https://github.com/llvm/llvm-project/pull/101110

>From 09c276696b2df9f9fdc5e350f2d8aa185ae7f749 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yifzhu at nvidia.com>
Date: Mon, 29 Jul 2024 16:01:45 -0700
Subject: [PATCH 1/4] [libc] add checksum for jmpbuf

Co-authored-by: Nick Desaulniers <ndesaulniers at google.com>
---
 libc/include/llvm-libc-types/jmp_buf.h |  4 ++
 libc/src/setjmp/CMakeLists.txt         | 10 ++++
 libc/src/setjmp/checksum.h             | 67 +++++++++++++++++++++++++
 libc/src/setjmp/x86_64/CMakeLists.txt  |  8 ++-
 libc/src/setjmp/x86_64/longjmp.cpp     | 69 +++++++++++++++++---------
 libc/src/setjmp/x86_64/setjmp.cpp      | 69 +++++++++++++-------------
 6 files changed, 164 insertions(+), 63 deletions(-)
 create mode 100644 libc/src/setjmp/checksum.h

diff --git a/libc/include/llvm-libc-types/jmp_buf.h b/libc/include/llvm-libc-types/jmp_buf.h
index 8949be9fa0ab7..fb48cac5c8f17 100644
--- a/libc/include/llvm-libc-types/jmp_buf.h
+++ b/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];
diff --git a/libc/src/setjmp/CMakeLists.txt b/libc/src/setjmp/CMakeLists.txt
index d85c532e8636c..9120ba459f0a4 100644
--- a/libc/src/setjmp/CMakeLists.txt
+++ b/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
diff --git a/libc/src/setjmp/checksum.h b/libc/src/setjmp/checksum.h
new file mode 100644
index 0000000000000..4ce1777d2bd87
--- /dev/null
+++ b/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
diff --git a/libc/src/setjmp/x86_64/CMakeLists.txt b/libc/src/setjmp/x86_64/CMakeLists.txt
index ae84322a65401..3f5797c607df7 100644
--- a/libc/src/setjmp/x86_64/CMakeLists.txt
+++ b/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
diff --git a/libc/src/setjmp/x86_64/longjmp.cpp b/libc/src/setjmp/x86_64/longjmp.cpp
index f479c7bc96c97..2f17de7b693d1 100644
--- a/libc/src/setjmp/x86_64/longjmp.cpp
+++ b/libc/src/setjmp/x86_64/longjmp.cpp
@@ -7,8 +7,10 @@
 //===----------------------------------------------------------------------===//
 
 #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"
@@ -16,30 +18,51 @@
 
 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)                                                           \
+  asm("mov %c[" #REG "](%%rdi), %%rdx\n\t"                                     \
+      "xor %%rdx, %%rcx\n\t"                                                   \
+      "mov %%rdx, %%" #REG "\n\t" ::[REG] "i"(offsetof(__jmp_buf, REG))        \
+      : "rdx");
+
+  RECOVER(rbx);
+  RECOVER(rbp);
+  RECOVER(r12);
+  RECOVER(r13);
+  RECOVER(r14);
+  RECOVER(r15);
+  RECOVER(rsp);
+
+  asm(R"(
+      xor %%eax,%%eax
+	    cmp $1,%%esi             
+	    adc %%esi,%%eax
+      mov %c[rip](%%rdi),%%rdx
+      xor %%rdx, %%rcx
+	    jmp *%%rdx
+  )" ::[rip] "i"(offsetof(__jmp_buf, rip))
+      : "rdx");
 }
 
 } // namespace LIBC_NAMESPACE_DECL
diff --git a/libc/src/setjmp/x86_64/setjmp.cpp b/libc/src/setjmp/x86_64/setjmp.cpp
index 6a1cc7a83936a..b3c92ded08839 100644
--- a/libc/src/setjmp/x86_64/setjmp.cpp
+++ b/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,40 @@
 
 namespace LIBC_NAMESPACE_DECL {
 
+namespace jmpbuf {} // namespace jmpbuf
+[[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

>From bc4926964c1b95f7cf24559af371081d74e970e7 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yifzhu at nvidia.com>
Date: Mon, 29 Jul 2024 17:33:50 -0700
Subject: [PATCH 2/4] adjust spacing

---
 libc/src/setjmp/x86_64/longjmp.cpp | 39 +++++++++++++++---------------
 libc/src/setjmp/x86_64/setjmp.cpp  | 14 +++++------
 2 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/libc/src/setjmp/x86_64/longjmp.cpp b/libc/src/setjmp/x86_64/longjmp.cpp
index 2f17de7b693d1..fca928e5a0df3 100644
--- a/libc/src/setjmp/x86_64/longjmp.cpp
+++ b/libc/src/setjmp/x86_64/longjmp.cpp
@@ -21,18 +21,19 @@ namespace LIBC_NAMESPACE_DECL {
 [[gnu::naked]]
 LLVM_LIBC_FUNCTION(void, longjmp, (__jmp_buf * buf, int val)) {
   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");
+   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
@@ -55,13 +56,13 @@ LLVM_LIBC_FUNCTION(void, longjmp, (__jmp_buf * buf, int val)) {
   RECOVER(rsp);
 
   asm(R"(
-      xor %%eax,%%eax
-	    cmp $1,%%esi             
-	    adc %%esi,%%eax
-      mov %c[rip](%%rdi),%%rdx
-      xor %%rdx, %%rcx
-	    jmp *%%rdx
-  )" ::[rip] "i"(offsetof(__jmp_buf, rip))
+   xor %%eax,%%eax
+   cmp $1,%%esi       
+   adc %%esi,%%eax
+   mov %c[rip](%%rdi),%%rdx
+   xor %%rdx, %%rcx
+   jmp *%%rdx
+ )" ::[rip] "i"(offsetof(__jmp_buf, rip))
       : "rdx");
 }
 
diff --git a/libc/src/setjmp/x86_64/setjmp.cpp b/libc/src/setjmp/x86_64/setjmp.cpp
index b3c92ded08839..e06997eff106f 100644
--- a/libc/src/setjmp/x86_64/setjmp.cpp
+++ b/libc/src/setjmp/x86_64/setjmp.cpp
@@ -39,13 +39,13 @@ LLVM_LIBC_FUNCTION(int, setjmp, (__jmp_buf * buf)) {
   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)),
+   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");
 

>From 2547493d5e8d4adf8d024b572a211660b848a71b Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yifzhu at nvidia.com>
Date: Tue, 30 Jul 2024 11:16:05 -0700
Subject: [PATCH 3/4] address CR

---
 libc/src/setjmp/x86_64/longjmp.cpp | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/libc/src/setjmp/x86_64/longjmp.cpp b/libc/src/setjmp/x86_64/longjmp.cpp
index fca928e5a0df3..7dfcabc92a7e1 100644
--- a/libc/src/setjmp/x86_64/longjmp.cpp
+++ b/libc/src/setjmp/x86_64/longjmp.cpp
@@ -42,10 +42,13 @@ LLVM_LIBC_FUNCTION(void, longjmp, (__jmp_buf * buf, int val)) {
   // load registers from buffer
   // do not pass any invalid values into registers
 #define RECOVER(REG)                                                           \
-  asm("mov %c[" #REG "](%%rdi), %%rdx\n\t"                                     \
-      "xor %%rdx, %%rcx\n\t"                                                   \
-      "mov %%rdx, %%" #REG "\n\t" ::[REG] "i"(offsetof(__jmp_buf, REG))        \
-      : "rdx");
+  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);
@@ -55,15 +58,18 @@ LLVM_LIBC_FUNCTION(void, longjmp, (__jmp_buf * buf, int val)) {
   RECOVER(r15);
   RECOVER(rsp);
 
-  asm(R"(
-   xor %%eax,%%eax
+  register int eax __asm__("eax");
+  asm volatile(R"(
+   xor %0,%0
    cmp $1,%%esi       
-   adc %%esi,%%eax
+   adc %%esi,%0
    mov %c[rip](%%rdi),%%rdx
    xor %%rdx, %%rcx
    jmp *%%rdx
- )" ::[rip] "i"(offsetof(__jmp_buf, rip))
-      : "rdx");
+ )"
+               : "=r"(eax)
+               : [rip] "i"(offsetof(__jmp_buf, rip))
+               : "rdx");
 }
 
 } // namespace LIBC_NAMESPACE_DECL

>From 21b738853449ca08cdcd22a58bca03624101e045 Mon Sep 17 00:00:00 2001
From: Yifan Zhu <yifzhu at nvidia.com>
Date: Tue, 30 Jul 2024 14:07:25 -0700
Subject: [PATCH 4/4] remove junk

---
 libc/src/setjmp/x86_64/setjmp.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/libc/src/setjmp/x86_64/setjmp.cpp b/libc/src/setjmp/x86_64/setjmp.cpp
index e06997eff106f..94a237e3bfd6e 100644
--- a/libc/src/setjmp/x86_64/setjmp.cpp
+++ b/libc/src/setjmp/x86_64/setjmp.cpp
@@ -17,7 +17,6 @@
 
 namespace LIBC_NAMESPACE_DECL {
 
-namespace jmpbuf {} // namespace jmpbuf
 [[gnu::naked]]
 LLVM_LIBC_FUNCTION(int, setjmp, (__jmp_buf * buf)) {
   register __UINT64_TYPE__ rcx __asm__("rcx");



More information about the libc-commits mailing list