Properly protect the startup stack with a mutex.

Because we start the APs sequentially, it is unlikely they will coincide
for the brief period that they use the temporary startup stack, but we
should guard against it. This allows us to remove the mutex around the
restart of each AP when relocating, which should improve test times.

app/main.c

@@ -126,9 +126,10 @@ static void run_at(uintptr_t addr, int my_pcpu)
     }
     BARRIER(true);
 
-    // We use a lock to ensure that only one CPU at a time jumps to
-    // the new code. Some of the startup stuff is not thread safe!
-    spin_lock(start_mutex);
+#ifndef __x86_64__
+    // The 32-bit startup code needs to know where it is located.
+    __asm__ __volatile__("movl %0, %%edi" : : "r" (new_start_addr));
+#endif
 
     goto *new_start_addr;
 }
@@ -384,9 +385,6 @@ void main(void)
         } else {
             pcpu_state[my_pcpu] = CPU_STATE_RUNNING;
         }
-    } else {
-        // Release the lock taken in run_at().
-        spin_unlock(start_mutex);
     }
     BARRIER(true);
     init_state = 2;

boot/boot.h

@@ -49,6 +49,8 @@
 
 extern uint8_t	_start[];
 
+extern uint8_t	startup32[];
+
 extern uint8_t	startup[];
 
 extern uint64_t	pml4[];

boot/efisetup.c

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: GPL-2.0-only
-// Copyright (C) 2020-2021 Martin Whitaker.
+// Copyright (C) 2020-2022 Martin Whitaker.
 //
 // Derived from Linux 5.6 arch/x86/boot/compressed/eboot.c and extracts
 // from drivers/firmware/efi/libstub:
@@ -541,6 +541,8 @@ boot_params_t *efi_setup(efi_handle_t handle, efi_system_table_t *sys_table_arg,
         memset(boot_params, 0, sizeof(boot_params_t));
     }
 
+    boot_params->code32_start = (uintptr_t)startup32;
+
     status = set_screen_info(boot_params);
     if (status != EFI_SUCCESS) {
         print_string("set_screen_info() failed\n");

boot/startup32.S

@@ -31,7 +31,7 @@
 startup32:
 	cld
 	cli
-	jmp	startup
+	jmp	first_start
 
 # The Linux 32-bit EFI handover point.
 
@@ -62,18 +62,28 @@ efi_handover:
 	call	efi_setup
 
 	# Fall though to the shared 32-bit entry point with the boot
-	# params pointer in %esi.
+	# params pointer in %esi...
 
 	movl	%eax, %esi
 
+	# ...and with the startup address in %edi.
+
+first_start:
+	movl	0x214(%esi), %ebx	# bootparams.code32_start
+	leal	(startup - startup32)(%ebx), %edi
+
 # The 32-bit entry point for AP boot and for restart after relocation.
 
 	.globl	startup
 startup:
-	# Use a temporary stack until we pick the correct one. We can
-	# safely use the high address, even if we are loaded low.
+	# Use the startup stack until we pick the correct one. We
+	# need to take the mutex to protect our use of the stack.
 
-	movl	$(HIGH_LOAD_ADDR + startup_stack_top - startup32), %esp
+	leal	(startup_stack_mutex - startup)(%edi), %eax
+0: lock bts	$0, (%eax)
+	jc	0b
+
+	leal	(startup_stack_top - startup)(%edi), %esp
 
 	# Load the GOT pointer.
 
@@ -98,6 +108,10 @@ startup:
 	leal	_end@GOTOFF(%ebx), %esp
 	addl	%eax, %esp
 
+	# Release the mutex that protects the startup stack.
+
+	movl	$0, startup_stack_mutex@GOTOFF(%ebx)
+
 	# Initialise the GDT descriptor.
 
 	leal	gdt@GOTOFF(%ebx), %eax
@@ -495,15 +509,15 @@ ap_trampoline:
 	movw	%cs, %ax
 	movw	%ax, %ds
 
-	# Patch the jump address.
+	# Load the startup address and use it to patch the jump address.
 
-	movl	(ap_startup_addr - ap_trampoline), %ebx
-	movl	%ebx, (ap_jump - ap_trampoline + 2)
+	movl	(ap_startup_addr - ap_trampoline), %edi
+	movl	%edi, (ap_jump - ap_trampoline + 2)
 
 	# Patch and load the GDT descriptor. It should point to the main
 	# GDT descriptor, which has already been initialised by the BSP.
 
-	movl	%ebx, %eax
+	movl	%edi, %eax
 	addl	$(gdt - startup), %eax
 	movl	%eax, (ap_gdt_descr - ap_trampoline + 2)
 	lgdt	ap_gdt_descr - ap_trampoline
@@ -544,11 +558,15 @@ ap_trampoline_end:
 # Variables.
 
 	.data
+	.align	4
 
 	.globl	boot_params_addr
 boot_params_addr:
 	.long	0
 
+startup_stack_mutex:
+	.long	0
+
 first_boot:
 	.long	1
 

boot/startup64.S

@@ -32,19 +32,18 @@ startup32:
 	cld
 	cli
 
-	# Use a temporary stack until we pick the correct one. We can
-	# safely use the high address, even if we are loaded low.
-
-	movl	$(HIGH_LOAD_ADDR + startup_stack_top - startup32), %esp
-
 	# Get the load address.
 
-	movl	0x214(%esi), %ebx
+	movl	0x214(%esi), %ebx	# bootparams.code32_start
 
 	# Save the boot params pointer.
 
 	movl	%esi, (boot_params_addr - startup32)(%ebx)
 
+	# Use the startup stack until we pick the correct one.
+
+	leal	(startup_stack_top - startup32)(%ebx), %esp
+
 	# Initialise the pml4 and pdp tables.
 
 	leal	(pml4 - startup32)(%ebx), %ecx
@@ -143,8 +142,11 @@ efi_handover:
 
 	.globl	startup
 startup:
-	# Use a temporary stack until we pick the correct one.
+	# Use the startup stack until we pick the correct one. We
+	# need to take a mutex to protect our use of the stack.
 
+0: lock bts	$0, startup_stack_mutex(%rip)
+	jc	0b
 	leaq	startup_stack_top(%rip), %rsp
 
 	# Pick the correct stack. The stacks are allocated immediately
@@ -158,6 +160,9 @@ startup:
 	leaq	_end(%rip), %rsp
 	addq	%rax, %rsp
 
+	# Release the mutex that protects the startup stack.
+	movl	$0, startup_stack_mutex(%rip)
+
 	# Initialise the pml4 and pdp tables.
 
 	leaq	pml4(%rip), %rcx
@@ -589,11 +594,15 @@ ap_trampoline_end:
 # Variables.
 
 	.data
+	.align	4
 
 	.globl	boot_params_addr
 boot_params_addr:
 	.quad	0
 
+startup_stack_mutex:
+	.long	0
+
 first_boot:
 	.long	1