commit 5d844e1b72513cf59b5e7c14295644efdcc66e44
Author: H.J. Lu <hjl.tools@gmail.com>
Date:   Fri Feb 14 14:45:34 2020 -0800
 
    i386: Enable CET support in ucontext functions
    
    1. getcontext and swapcontext are updated to save the caller's shadow
    stack pointer and return address.
    2. setcontext and swapcontext are updated to restore shadow stack and
    jump to new context directly.
    3. makecontext is updated to allocate a new shadow stack and set the
    caller's return address to the helper code, L(exitcode).
    4. Since we no longer save and restore EAX, ECX and EDX in getcontext,
    setcontext and swapcontext, we can use them as scratch register slots
    to enable CET in ucontext functions.
    
    Since makecontext allocates a new shadow stack when making a new
    context and kernel allocates a new shadow stack for clone/fork/vfork
    syscalls, we track the current shadow stack base.  In setcontext and
    swapcontext, if the target shadow stack base is the same as the current
    shadow stack base, we unwind the shadow stack.  Otherwise it is a stack
    switch and we look for a restore token.
    
    We enable shadow stack at run-time only if program and all used shared
    objects, including dlopened ones, are shadow stack enabled, which means
    that they must be compiled with GCC 8 or above and glibc 2.28 or above.
    We need to save and restore shadow stack only if shadow stack is enabled.
    When caller of getcontext, setcontext, swapcontext and makecontext is
    compiled with smaller ucontext_t, shadow stack won't be enabled at
    run-time.  We check if shadow stack is enabled before accessing the
    extended field in ucontext_t.
    
    Tested on i386 CET/non-CET machines.
    
    Reviewed-by: Carlos O'Donell <carlos@redhat.com>
---
 
diff --git a/sysdeps/unix/sysv/linux/i386/getcontext.S b/sysdeps/unix/sysv/linux/i386/getcontext.S
index 6637596..4ed9d03 100644
--- a/sysdeps/unix/sysv/linux/i386/getcontext.S
+++ b/sysdeps/unix/sysv/linux/i386/getcontext.S
@@ -18,6 +18,7 @@
    <http://www.gnu.org/licenses/>.  */
 
 #include <sysdep.h>
+#include <asm/prctl.h>
 
 #include "ucontext_i.h"
 
@@ -42,6 +43,61 @@ ENTRY(__getcontext)
 	movw	%fs, %dx
 	movl	%edx, oFS(%eax)
 
+#if SHSTK_ENABLED
+	/* Check if shadow stack is enabled.  */
+	testl	$X86_FEATURE_1_SHSTK, %gs:FEATURE_1_OFFSET
+	jz	L(no_shstk)
+
+	/* Save EAX in EDX.  */
+	movl	%eax, %edx
+
+	xorl	%eax, %eax
+	cmpl	%gs:SSP_BASE_OFFSET, %eax
+	jnz	L(shadow_stack_bound_recorded)
+
+	/* Save EBX in the first scratch register slot.  */
+	movl	%ebx, oSCRATCH1(%edx)
+
+	/* Get the base address and size of the default shadow stack
+	   which must be the current shadow stack since nothing has
+	   been recorded yet.  */
+	sub	$24, %esp
+	mov	%esp, %ecx
+	movl	$ARCH_CET_STATUS, %ebx
+	movl	$__NR_arch_prctl, %eax
+	ENTER_KERNEL
+	testl	%eax, %eax
+	jz	L(continue_no_err)
+
+	/* This should never happen.  */
+	hlt
+
+L(continue_no_err):
+	/* Restore EBX from the first scratch register slot.  */
+	movl	oSCRATCH1(%edx), %ebx
+
+	/* Record the base of the current shadow stack.  */
+	movl	8(%esp), %eax
+	movl	%eax, %gs:SSP_BASE_OFFSET
+	add	$24, %esp
+
+L(shadow_stack_bound_recorded):
+	/* Load address of the context data structure.  */
+	movl	4(%esp), %eax
+
+	/* Get the current shadow stack pointer.  */
+	rdsspd	%edx
+	/* NB: Save the caller's shadow stack so that we can jump back
+	   to the caller directly.  */
+	addl	$4, %edx
+	movl	%edx, oSSP(%eax)
+
+	/* Save the current shadow stack base in ucontext.  */
+	movl	%gs:SSP_BASE_OFFSET, %edx
+	movl	%edx, (oSSP + 4)(%eax)
+
+L(no_shstk):
+#endif
 	/* We have separate floating-point register content memory on the
 	   stack.  We use the __fpregs_mem block in the context.  Set the
 	   links up correctly.  */
diff --git a/sysdeps/unix/sysv/linux/i386/makecontext.S b/sysdeps/unix/sysv/linux/i386/makecontext.S
index e3ca3dc..2d82ddc 100644
--- a/sysdeps/unix/sysv/linux/i386/makecontext.S
+++ b/sysdeps/unix/sysv/linux/i386/makecontext.S
@@ -18,6 +18,7 @@
    <http://www.gnu.org/licenses/>.  */
 
 #include <sysdep.h>
+#include <asm/prctl.h>
 
 #include "ucontext_i.h"
 
@@ -68,6 +69,127 @@ ENTRY(__makecontext)
 	jnz	1b
 2:
 
+#if SHSTK_ENABLED
+	/* Check if Shadow Stack is enabled.  */
+	testl	$X86_FEATURE_1_SHSTK, %gs:FEATURE_1_OFFSET
+	jz	L(skip_ssp)
+
+	/* Reload the pointer to ucontext.  */
+	movl	4(%esp), %eax
+
+	/* Shadow stack is enabled.  We need to allocate a new shadow
+	   stack.  */
+	subl	oSS_SP(%eax), %edx
+	shrl	$STACK_SIZE_TO_SHADOW_STACK_SIZE_SHIFT, %edx
+
+	/* Align shadow stack size to 8 bytes.  */
+	addl	$7, %edx
+	andl	$-8, %edx
+
+	/* Store shadow stack size in __ssp[2].  */
+	movl	%edx, (oSSP + 8)(%eax)
+
+	/* Save ESI in the second scratch register slot.  */
+	movl	%esi, oSCRATCH2(%eax)
+	/* Save EDI in the third scratch register slot.  */
+	movl	%edi, oSCRATCH3(%eax)
+
+	/* Save the pointer to ucontext.  */
+	movl	%eax, %edi
+
+	/* Get the original shadow stack pointer.  */
+	rdsspd	%esi
+
+	/* Align the saved original shadow stack pointer to the next
+	   8 byte aligned boundary.  */
+	andl	$-8, %esi
+
+	/* Load the top of the new stack into EDX.  */
+	movl	oESP(%eax), %edx
+
+	/* We need to terminate the FDE here because the unwinder looks
+	   at ra-1 for unwind information.  */
+	cfi_endproc
+
+	/* Swap the original stack pointer with the top of the new
+	   stack.  */
+	xchgl	%esp, %edx
+
+	/* Add 4 bytes since CALL will push the 4-byte return address
+	   onto stack.  */
+	addl	$4, %esp
+
+	/* Allocate the new shadow stack.  Save EBX in the first scratch
+	   register slot.  */
+	movl	%ebx, oSCRATCH1(%eax)
+
+	/* CET syscall takes 64-bit sizes.  */
+	subl	$16, %esp
+	movl	(oSSP + 8)(%eax), %ecx
+	movl	%ecx, (%esp)
+	movl	$0, 4(%esp)
+	movl	%ecx, 8(%esp)
+	movl	$0, 12(%esp)
+	movl	%esp, %ecx
+
+	movl	$ARCH_CET_ALLOC_SHSTK, %ebx
+	movl	$__NR_arch_prctl, %eax
+	ENTER_KERNEL
+	testl	%eax, %eax
+	jne	L(hlt)		/* This should never happen.  */
+
+	/* Copy the base address of the new shadow stack to __ssp[1].  */
+	movl	(%esp), %eax
+	movl	%eax, (oSSP + 4)(%edi)
+
+	addl	$16, %esp
+
+	/* Restore EBX from the first scratch register slot.  */
+	movl	oSCRATCH1(%edi), %ebx
+
+	/* Get the size of the new shadow stack.  */
+	movl	(oSSP + 8)(%edi), %ecx
+
+	/* Use the restore stoken to restore the new shadow stack.  */
+	rstorssp -8(%eax, %ecx)
+
+	/* Save the restore token at the next 8 byte aligned boundary
+	   on the original shadow stack.  */
+	saveprevssp
+
+	/* Push the address of "jmp exitcode" onto the new stack as
+	   well as the new shadow stack.  */
+	call	1f
+	jmp	L(exitcode)
+1:
+
+	/* Get the new shadow stack pointer.  */
+	rdsspd	%eax
+
+	/* Use the restore stoken to restore the original shadow stack.  */
+	rstorssp -8(%esi)
+
+	/* Save the restore token on the new shadow stack.  */
+	saveprevssp
+
+	/* Store the new shadow stack pointer in __ssp[0].  */
+	movl	%eax, oSSP(%edi)
+
+	/* Restore the original stack.  */
+	mov	%edx, %esp
+
+	cfi_startproc
+
+	/* Restore ESI from the second scratch register slot.  */
+	movl	oSCRATCH2(%edi), %esi
+	/* Restore EDI from the third scratch register slot.  */
+	movl	oSCRATCH3(%edi), %edi
+
+	ret
+
+L(skip_ssp):
+#endif
+
 	/* If the function we call returns we must continue with the
 	   context which is given in the uc_link element.  To do this
 	   set the return address for the function the user provides
@@ -123,6 +245,7 @@ L(call_exit):
 	call	HIDDEN_JUMPTARGET(exit)
 	/* The 'exit' call should never return.  In case it does cause
 	   the process to terminate.  */
+L(hlt):
 	hlt
 	cfi_startproc
 END(__makecontext)
diff --git a/sysdeps/unix/sysv/linux/i386/setcontext.S b/sysdeps/unix/sysv/linux/i386/setcontext.S
index 7565d7d..7b58918 100644
--- a/sysdeps/unix/sysv/linux/i386/setcontext.S
+++ b/sysdeps/unix/sysv/linux/i386/setcontext.S
@@ -18,6 +18,7 @@
    <http://www.gnu.org/licenses/>.  */
 
 #include <sysdep.h>
+#include <asm/prctl.h>
 
 #include "ucontext_i.h"
 
@@ -56,9 +57,6 @@ ENTRY(__setcontext)
 	movl	oFS(%eax), %ecx
 	movw	%cx, %fs
 
-	/* Fetch the address to return to.  */
-	movl	oEIP(%eax), %ecx
-
 	/* Load the new stack pointer.  */
 	cfi_def_cfa (eax, 0)
 	cfi_offset (edi, oEDI)
@@ -67,6 +65,103 @@ ENTRY(__setcontext)
 	cfi_offset (ebx, oEBX)
 	movl	oESP(%eax), %esp
 
+#if SHSTK_ENABLED
+	/* Check if Shadow Stack is enabled.  */
+	testl	$X86_FEATURE_1_SHSTK, %gs:FEATURE_1_OFFSET
+	jz	L(no_shstk)
+
+	/* If the base of the target shadow stack is the same as the
+	   base of the current shadow stack, we unwind the shadow
+	   stack.  Otherwise it is a stack switch and we look for a
+	   restore token.  */
+	movl	oSSP(%eax), %esi
+	movl	%esi, %edi
+
+	/* Get the base of the target shadow stack.  */
+	movl	(oSSP + 4)(%eax), %ecx
+	cmpl	%gs:SSP_BASE_OFFSET, %ecx
+	je	L(unwind_shadow_stack)
+
+	/* Align the saved original shadow stack pointer to the next
+	   8 byte aligned boundary.  */
+	andl	$-8, %esi
+
+L(find_restore_token_loop):
+	/* Look for a restore token.  */
+	movl	-8(%esi), %ebx
+	andl	$-8, %ebx
+	cmpl	%esi, %ebx
+	je	L(restore_shadow_stack)
+
+	/* Try the next slot.  */
+	subl	$8, %esi
+	jmp	L(find_restore_token_loop)
+
+L(restore_shadow_stack):
+	/* Pop return address from the shadow stack since setcontext
+	   will not return.  */
+	movl	$1, %ebx
+	incsspd	%ebx
+
+	/* Use the restore stoken to restore the target shadow stack.  */
+	rstorssp -8(%esi)
+
+	/* Save the restore token on the old shadow stack.  NB: This
+	   restore token may be checked by setcontext or swapcontext
+	   later.  */
+	saveprevssp
+
+	/* Record the new shadow stack base that was switched to.  */
+	movl	(oSSP + 4)(%eax), %ebx
+	movl	%ebx, %gs:SSP_BASE_OFFSET
+
+L(unwind_shadow_stack):
+	rdsspd	%ebx
+	subl	%edi, %ebx
+	je	L(skip_unwind_shadow_stack)
+	negl	%ebx
+	shrl	$2, %ebx
+	movl	$255, %esi
+L(loop):
+	cmpl	%esi, %ebx
+	cmovb	%ebx, %esi
+	incsspd	%esi
+	subl	%esi, %ebx
+	ja	L(loop)
+
+L(skip_unwind_shadow_stack):
+
+	/* Load the values of all the preserved registers (except ESP).  */
+	movl	oEDI(%eax), %edi
+	movl	oESI(%eax), %esi
+	movl	oEBP(%eax), %ebp
+	movl	oEBX(%eax), %ebx
+
+	/* Get the return address set with getcontext.  */
+	movl	oEIP(%eax), %ecx
+
+	/* Check if return address is valid for the case when setcontext
+	   is invoked from L(exitcode) with linked context.  */
+	rdsspd	%eax
+	cmpl	(%eax), %ecx
+	/* Clear EAX to indicate success.  NB: Don't use xorl to keep
+	   EFLAGS for jne.  */
+	movl	$0, %eax
+	jne	L(jmp)
+	/* Return to the new context if return address valid.  */
+	pushl	%ecx
+	ret
+
+L(jmp):
+	/* Jump to the new context directly.  */
+	jmp	*%ecx
+
+L(no_shstk):
+#endif
+
+	/* Fetch the address to return to.  */
+	movl	oEIP(%eax), %ecx
+
 	/* Push the return address on the new stack so we can return there.  */
 	pushl	%ecx
 
diff --git a/sysdeps/unix/sysv/linux/i386/swapcontext.S b/sysdeps/unix/sysv/linux/i386/swapcontext.S
index ce27d51..d1b648c 100644
--- a/sysdeps/unix/sysv/linux/i386/swapcontext.S
+++ b/sysdeps/unix/sysv/linux/i386/swapcontext.S
@@ -18,6 +18,7 @@
    <http://www.gnu.org/licenses/>.  */
 
 #include <sysdep.h>
+#include <asm/prctl.h>
 
 #include "ucontext_i.h"
 
@@ -76,6 +77,144 @@ ENTRY(__swapcontext)
 	movl	oFS(%eax), %edx
 	movw	%dx, %fs
 
+#if SHSTK_ENABLED
+	/* Check if Shadow Stack is enabled.  */
+	testl	$X86_FEATURE_1_SHSTK, %gs:FEATURE_1_OFFSET
+	jz	L(no_shstk)
+
+	xorl	%eax, %eax
+	cmpl	%gs:SSP_BASE_OFFSET, %eax
+	jnz	L(shadow_stack_bound_recorded)
+
+	/* Get the base address and size of the default shadow stack
+	   which must be the current shadow stack since nothing has
+	   been recorded yet.  */
+	sub	$24, %esp
+	mov	%esp, %ecx
+	movl	$ARCH_CET_STATUS, %ebx
+	movl	$__NR_arch_prctl, %eax
+	ENTER_KERNEL
+	testl	%eax, %eax
+	jz	L(continue_no_err)
+
+	/* This should never happen.  */
+	hlt
+
+L(continue_no_err):
+	/* Record the base of the current shadow stack.  */
+	movl	8(%esp), %eax
+	movl	%eax, %gs:SSP_BASE_OFFSET
+	add	$24, %esp
+
+L(shadow_stack_bound_recorded):
+	/* Load address of the context data structure we save in.  */
+	movl	4(%esp), %eax
+
+	/* Load address of the context data structure we swap in  */
+	movl	8(%esp), %edx
+
+       /* If we unwind the stack, we can't undo stack unwinding.  Just
+	   save the target shadow stack pointer as the current shadow
+	   stack pointer.   */
+	movl	oSSP(%edx), %ecx
+	movl	%ecx, oSSP(%eax)
+
+	/* Save the current shadow stack base in ucontext.  */
+	movl	%gs:SSP_BASE_OFFSET, %ecx
+	movl	%ecx, (oSSP + 4)(%eax)
+
+	/* If the base of the target shadow stack is the same as the
+	   base of the current shadow stack, we unwind the shadow
+	   stack.  Otherwise it is a stack switch and we look for a
+	   restore token.  */
+	movl	oSSP(%edx), %esi
+	movl	%esi, %edi
+
+	/* Get the base of the target shadow stack.  */
+	movl	(oSSP + 4)(%edx), %ecx
+	cmpl	%gs:SSP_BASE_OFFSET, %ecx
+	je	L(unwind_shadow_stack)
+
+	/* Align the saved original shadow stack pointer to the next
+	   8 byte aligned boundary.  */
+	andl	$-8, %esi
+
+L(find_restore_token_loop):
+	/* Look for a restore token.  */
+	movl	-8(%esi), %ebx
+	andl	$-8, %ebx
+	cmpl	%esi, %ebx
+	je	L(restore_shadow_stack)
+
+	/* Try the next slot.  */
+	subl	$8, %esi
+	jmp	L(find_restore_token_loop)
+
+L(restore_shadow_stack):
+	/* The target shadow stack will be restored.  Save the current
+	   shadow stack pointer.  */
+	rdsspd	%ecx
+	movl	%ecx, oSSP(%eax)
+
+	/* Use the restore stoken to restore the target shadow stack.  */
+	rstorssp -8(%esi)
+
+	/* Save the restore token on the old shadow stack.  NB: This
+	   restore token may be checked by setcontext or swapcontext
+	   later.  */
+	saveprevssp
+
+	/* Record the new shadow stack base that was switched to.  */
+	movl	(oSSP + 4)(%edx), %ebx
+	movl	%ebx, %gs:SSP_BASE_OFFSET
+
+L(unwind_shadow_stack):
+	rdsspd	%ebx
+	subl	%edi, %ebx
+	je	L(skip_unwind_shadow_stack)
+	negl	%ebx
+	shrl	$2, %ebx
+	movl	$255, %esi
+L(loop):
+	cmpl	%esi, %ebx
+	cmovb	%ebx, %esi
+	incsspd	%esi
+	subl	%esi, %ebx
+	ja	L(loop)
+
+L(skip_unwind_shadow_stack):
+
+	/* Load the new stack pointer.  */
+	movl	oESP(%edx), %esp
+
+	/* Load the values of all the preserved registers (except ESP).  */
+	movl	oEDI(%edx), %edi
+	movl	oESI(%edx), %esi
+	movl	oEBP(%edx), %ebp
+	movl	oEBX(%edx), %ebx
+
+	/* Get the return address set with getcontext.  */
+	movl	oEIP(%edx), %ecx
+
+	/* Check if return address is valid for the case when setcontext
+	   is invoked from L(exitcode) with linked context.  */
+	rdsspd	%eax
+	cmpl	(%eax), %ecx
+	/* Clear EAX to indicate success.  NB: Don't use xorl to keep
+	   EFLAGS for jne.  */
+	movl	$0, %eax
+	jne	L(jmp)
+	/* Return to the new context if return address valid.  */
+	pushl	%ecx
+	ret
+
+L(jmp):
+	/* Jump to the new context directly.  */
+	jmp	*%ecx
+
+L(no_shstk):
+#endif
+
 	/* Fetch the address to return to.  */
 	movl	oEIP(%eax), %ecx
 
diff --git a/sysdeps/unix/sysv/linux/i386/sysdep.h b/sysdeps/unix/sysv/linux/i386/sysdep.h
index 3255cc7..9344ac7 100644
--- a/sysdeps/unix/sysv/linux/i386/sysdep.h
+++ b/sysdeps/unix/sysv/linux/i386/sysdep.h
@@ -656,4 +656,9 @@ struct libc_do_syscall_args
 # endif
 #endif
 
+/* Each shadow stack slot takes 4 bytes.  Assuming that each stack
+   frame takes 128 bytes, this is used to compute shadow stack size
+   from stack size.  */
+#define STACK_SIZE_TO_SHADOW_STACK_SIZE_SHIFT 5
+
 #endif /* linux/i386/sysdep.h */
diff --git a/sysdeps/unix/sysv/linux/i386/ucontext_i.sym b/sysdeps/unix/sysv/linux/i386/ucontext_i.sym
index 1dfe03d..1d8608e 100644
--- a/sysdeps/unix/sysv/linux/i386/ucontext_i.sym
+++ b/sysdeps/unix/sysv/linux/i386/ucontext_i.sym
@@ -22,6 +22,10 @@ oEBP		mreg (EBP)
 oESP		mreg (ESP)
 oEBX		mreg (EBX)
 oEIP		mreg (EIP)
+oSCRATCH1	mreg (EAX)
+oSCRATCH2	mreg (ECX)
+oSCRATCH3	mreg (EDX)
 oFPREGS		mcontext (fpregs)
 oSIGMASK	ucontext (uc_sigmask)
 oFPREGSMEM	ucontext (__fpregs_mem)
+oSSP		ucontext (__ssp)