From 82f4975f455523b168733711e1877e9cc33c120e Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Sun, 4 May 2025 13:54:35 +0300
Subject: [PATCH] LibC: Save callee saved registers on setjmp

---
 userspace/libraries/LibC/arch/i686/setjmp.S   | 15 +++++++++++++--
 userspace/libraries/LibC/arch/x86_64/setjmp.S | 19 +++++++++++++++++--
 userspace/libraries/LibC/include/setjmp.h     | 10 ++++++++--
 userspace/libraries/LibC/setjmp.cpp           | 10 +++++-----
 4 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/userspace/libraries/LibC/arch/i686/setjmp.S b/userspace/libraries/LibC/arch/i686/setjmp.S
index 3a33e21c..cbb318e5 100644
--- a/userspace/libraries/LibC/arch/i686/setjmp.S
+++ b/userspace/libraries/LibC/arch/i686/setjmp.S
@@ -1,3 +1,5 @@
+// jmp_buf: esp, eip, ebx, ebp, edi, esi
+
 // int setjmp(jmp_buf env)
 .global setjmp
 setjmp:
@@ -9,6 +11,11 @@ setjmp:
 	movl (%esp), %eax
 	movl %eax, 4(%edx)
 
+	movl %ebx,  8(%edx)
+	movl %ebp, 12(%edx)
+	movl %edi, 16(%edx)
+	movl %esi, 20(%edx)
+
 	xorl %eax, %eax
 
 	ret
@@ -25,8 +32,12 @@ longjmp:
 	testl %ecx, %ecx
 	cmovnzl %ecx, %eax
 
-	movl 0(%edx), %esp
-	movl 4(%edx), %ecx
+	movl  0(%edx), %esp
+	movl  4(%edx), %ecx
+	movl  8(%edx), %ebx
+	movl 12(%edx), %ebp
+	movl 16(%edx), %edi
+	movl 20(%edx), %esi
 	jmp *%ecx
 
  .size longjmp, . - longjmp
diff --git a/userspace/libraries/LibC/arch/x86_64/setjmp.S b/userspace/libraries/LibC/arch/x86_64/setjmp.S
index fd1634e8..da77a1ff 100644
--- a/userspace/libraries/LibC/arch/x86_64/setjmp.S
+++ b/userspace/libraries/LibC/arch/x86_64/setjmp.S
@@ -1,3 +1,5 @@
+// jmp_buf: rsp, rip, rbx, rbp, r12-r15
+
 // int setjmp(jmp_buf env)
 .global setjmp
 setjmp:
@@ -7,6 +9,13 @@ setjmp:
 	movq (%rsp), %rax
 	movq %rax, 8(%rdi)
 
+	movq %rbx, 16(%rdi)
+	movq %rbp, 24(%rdi)
+	movq %r12, 32(%rdi)
+	movq %r13, 40(%rdi)
+	movq %r14, 48(%rdi)
+	movq %r15, 56(%rdi)
+
 	xorq %rax, %rax
 
 	ret
@@ -19,7 +28,13 @@ longjmp:
 	testq %rsi, %rsi
 	cmovnzq %rsi, %rax
 
-	movq 0(%rdi), %rsp
-	movq 8(%rdi), %rcx
+	movq  0(%rdi), %rsp
+	movq  8(%rdi), %rcx
+	movq 16(%rdi), %rbx
+	movq 24(%rdi), %rbp
+	movq 32(%rdi), %r12
+	movq 40(%rdi), %r13
+	movq 48(%rdi), %r14
+	movq 56(%rdi), %r15
 	jmp *%rcx
  .size longjmp, . - longjmp
diff --git a/userspace/libraries/LibC/include/setjmp.h b/userspace/libraries/LibC/include/setjmp.h
index ada6e58e..b3214313 100644
--- a/userspace/libraries/LibC/include/setjmp.h
+++ b/userspace/libraries/LibC/include/setjmp.h
@@ -7,8 +7,14 @@
 
 __BEGIN_DECLS
 
-typedef long jmp_buf[2];
-typedef long sigjmp_buf[2 + 1 + (sizeof(long long) / sizeof(long))];
+#if defined(__x86_64__)
+#define _JMP_BUF_REGS 8 // rsp, rip, rbx, rbp, r12-r15
+#elif defined(__i686__)
+#define _JMP_BUF_REGS 6 // esp, eip, ebx, ebp, edi, esi
+#endif
+
+typedef long jmp_buf[_JMP_BUF_REGS];
+typedef long sigjmp_buf[_JMP_BUF_REGS + 1 + (sizeof(long long) / sizeof(long))];
 
 #define _longjmp longjmp
 void longjmp(jmp_buf env, int val);
diff --git a/userspace/libraries/LibC/setjmp.cpp b/userspace/libraries/LibC/setjmp.cpp
index 9fb94cc1..0774b1b4 100644
--- a/userspace/libraries/LibC/setjmp.cpp
+++ b/userspace/libraries/LibC/setjmp.cpp
@@ -1,19 +1,19 @@
 #include <setjmp.h>
 #include <signal.h>
 
-static_assert(sizeof(sigjmp_buf) == sizeof(jmp_buf) + sizeof(long) + sizeof(sigset_t));
+static_assert(sizeof(sigjmp_buf) == (_JMP_BUF_REGS + 1) * sizeof(long) + sizeof(sigset_t));
 
 void siglongjmp(sigjmp_buf env, int val)
 {
-	if (env[2])
-		pthread_sigmask(SIG_SETMASK, reinterpret_cast<sigset_t*>(&env[3]), nullptr);
+	if (env[_JMP_BUF_REGS])
+		pthread_sigmask(SIG_SETMASK, reinterpret_cast<sigset_t*>(&env[_JMP_BUF_REGS + 1]), nullptr);
 	return longjmp(env, val);
 }
 
 int sigsetjmp(sigjmp_buf env, int savemask)
 {
-	env[2] = savemask;
+	env[_JMP_BUF_REGS] = savemask;
 	if (savemask)
-		pthread_sigmask(0, nullptr, reinterpret_cast<sigset_t*>(&env[3]));
+		pthread_sigmask(0, nullptr, reinterpret_cast<sigset_t*>(&env[_JMP_BUF_REGS + 1]));
 	return setjmp(env);
 }