// Program that loops for ever doing lots of recursions and system calls,
// intended to be used as part of a stress test for GCS context switching.
//
// Copyright 2015-2023 Arm Ltd

#include <asm/unistd.h>

#define sa_sz 32
#define sa_flags 8
#define sa_handler 0
#define sa_mask_sz 8

#define si_code 8

#define SIGINT 2
#define SIGABRT 6
#define SIGUSR1 10
#define SIGSEGV 11
#define SIGUSR2 12
#define SIGTERM 15
#define SEGV_CPERR 10

#define SA_NODEFER 1073741824
#define SA_SIGINFO 4
#define ucontext_regs 184

#define PR_SET_SHADOW_STACK_STATUS      75
# define PR_SHADOW_STACK_ENABLE         (1UL << 0)

#define	GCSPR_EL0 S3_3_C2_C5_1

.macro function name
	.macro endfunction
		.type \name, @function
		.purgem endfunction
	.endm
\name:
.endm

// Print a single character x0 to stdout
// Clobbers x0-x2,x8
function putc
	str	x0, [sp, #-16]!

	mov	x0, #1			// STDOUT_FILENO
	mov	x1, sp
	mov	x2, #1
	mov	x8, #__NR_write
	svc	#0

	add	sp, sp, #16
	ret
endfunction
.globl	putc

// Print a NUL-terminated string starting at address x0 to stdout
// Clobbers x0-x3,x8
function puts
	mov	x1, x0

	mov	x2, #0
0:	ldrb	w3, [x0], #1
	cbz	w3, 1f
	add	x2, x2, #1
	b	0b

1:	mov	w0, #1			// STDOUT_FILENO
	mov	x8, #__NR_write
	svc	#0

	ret
endfunction
.globl	puts

// Utility macro to print a literal string
// Clobbers x0-x4,x8
.macro puts string
	.pushsection .rodata.str1.1, "aMS", @progbits, 1
.L__puts_literal\@: .string "\string"
	.popsection

	ldr	x0, =.L__puts_literal\@
	bl	puts
.endm

// Print an unsigned decimal number x0 to stdout
// Clobbers x0-x4,x8
function putdec
	mov	x1, sp
	str	x30, [sp, #-32]!	// Result can't be > 20 digits

	mov	x2, #0
	strb	w2, [x1, #-1]!		// Write the NUL terminator

	mov	x2, #10
0:	udiv	x3, x0, x2		// div-mod loop to generate the digits
	msub	x0, x3, x2, x0
	add	w0, w0, #'0'
	strb	w0, [x1, #-1]!
	mov	x0, x3
	cbnz	x3, 0b

	ldrb	w0, [x1]
	cbnz	w0, 1f
	mov	w0, #'0'		// Print "0" for 0, not ""
	strb	w0, [x1, #-1]!

1:	mov	x0, x1
	bl	puts

	ldr	x30, [sp], #32
	ret
endfunction
.globl	putdec

// Print an unsigned decimal number x0 to stdout, followed by a newline
// Clobbers x0-x5,x8
function putdecn
	mov	x5, x30

	bl	putdec
	mov	x0, #'\n'
	bl	putc

	ret	x5
endfunction
.globl	putdecn

// Fill x1 bytes starting at x0 with 0.
// Clobbers x1, x2.
function memclr
	mov	w2, #0
endfunction
.globl	memclr
	// fall through to memfill

// Trivial memory fill: fill x1 bytes starting at address x0 with byte w2
// Clobbers x1
function memfill
	cmp	x1, #0
	b.eq	1f

0:	strb	w2, [x0], #1
	subs	x1, x1, #1
	b.ne	0b

1:	ret
endfunction
.globl	memfill

// w0: signal number
// x1: sa_action
// w2: sa_flags
// Clobbers x0-x6,x8
function setsignal
	str	x30, [sp, #-((sa_sz + 15) / 16 * 16 + 16)]!

	mov	w4, w0
	mov	x5, x1
	mov	w6, w2

	add	x0, sp, #16
	mov	x1, #sa_sz
	bl	memclr

	mov	w0, w4
	add	x1, sp, #16
	str	w6, [x1, #sa_flags]
	str	x5, [x1, #sa_handler]
	mov	x2, #0
	mov	x3, #sa_mask_sz
	mov	x8, #__NR_rt_sigaction
	svc	#0

	cbz	w0, 1f

	puts	"sigaction failure\n"
	b	abort

1:	ldr	x30, [sp], #((sa_sz + 15) / 16 * 16 + 16)
	ret
endfunction


function tickle_handler
	// Perhaps collect GCSPR_EL0 here in future?
	ret
endfunction

function terminate_handler
	mov	w21, w0
	mov	x20, x2

	puts	"Terminated by signal "
	mov	w0, w21
	bl	putdec
	puts	", no error\n"

	mov	x0, #0
	mov	x8, #__NR_exit
	svc	#0
endfunction

function segv_handler
	// stash the siginfo_t *
	mov	x20, x1

	// Disable GCS, we don't want additional faults logging things
	mov	x0, PR_SET_SHADOW_STACK_STATUS
	mov	x1, xzr
	mov	x2, xzr
	mov	x3, xzr
	mov	x4, xzr
	mov	x5, xzr
	mov	x8, #__NR_prctl
	svc	#0

	puts	"Got SIGSEGV code "

	ldr	x21, [x20, #si_code]
	mov	x0, x21
	bl	putdec

	// GCS faults should have si_code SEGV_CPERR
	cmp	x21, #SEGV_CPERR
	bne	1f

	puts	" (GCS violation)"
1:
	mov	x0, '\n'
	bl	putc
	b	abort
endfunction

// Recurse x20 times
.macro recurse id
function recurse\id
	stp	x29, x30, [sp, #-16]!
	mov	x29, sp

	cmp	x20, 0
	beq	1f
	sub	x20, x20, 1
	bl	recurse\id

1:
	ldp	x29, x30, [sp], #16

	// Do a syscall immediately prior to returning to try to provoke
	// scheduling and migration at a point where coherency issues
	// might trigger.
	mov	x8, #__NR_getpid
	svc	#0

	ret
endfunction
.endm

// Generate and use two copies so we're changing the GCS contents
recurse 1
recurse 2

.globl _start
function _start
	// Run with GCS
	mov	x0, PR_SET_SHADOW_STACK_STATUS
	mov	x1, PR_SHADOW_STACK_ENABLE
	mov	x2, xzr
	mov	x3, xzr
	mov	x4, xzr
	mov	x5, xzr
	mov	x8, #__NR_prctl
	svc	#0
	cbz	x0, 1f
	puts	"Failed to enable GCS\n"
	b	abort
1:

	mov	w0, #SIGTERM
	adr	x1, terminate_handler
	mov	w2, #SA_SIGINFO
	bl	setsignal

	mov	w0, #SIGUSR1
	adr	x1, tickle_handler
	mov	w2, #SA_SIGINFO
	orr	w2, w2, #SA_NODEFER
	bl	setsignal

	mov	w0, #SIGSEGV
	adr	x1, segv_handler
	mov	w2, #SA_SIGINFO
	orr	w2, w2, #SA_NODEFER
	bl	setsignal

	puts	"Running\n"

loop:
	// Small recursion depth so we're frequently flipping between
	// the two recursors and changing what's on the stack
	mov	x20, #5
	bl	recurse1
	mov	x20, #5
	bl	recurse2
	b	loop
endfunction

abort:
	mov	x0, #255
	mov	x8, #__NR_exit
	svc	#0