From 59e71f2347ac5cc22a1f283026924b3525c20db2 Mon Sep 17 00:00:00 2001
From: Augustin Cavalier <waddlesplash@gmail.com>
Date: Mon, 16 Dec 2024 15:39:14 -0500
Subject: [PATCH] kernel: Pre-fault the stack and IP page addresses on fork().

Since we don't pre-map anything in a forked process, without this
we'd just fault immediately upon restoring the fork frame multiple
times in a row. This way, we can do some of the faulting up front,
before even returning to userland.

2	0	headers/private/kernel/arch/thread.h
10	0	src/system/kernel/arch/arm/arch_thread.cpp
7	0	src/system/kernel/arch/arm64/arch_thread.cpp
10	2	src/system/kernel/arch/m68k/arch_thread.cpp
10	2	src/system/kernel/arch/ppc/arch_thread.cpp
10	2	src/system/kernel/arch/riscv64/arch_thread.cpp
10	3	src/system/kernel/arch/sparc/arch_thread.cpp
13	0	src/system/kernel/arch/x86/arch_thread.cpp
29	13	src/system/kernel/thread.cpp

diff --git a/headers/private/kernel/arch/thread.h b/headers/private/kernel/arch/thread.h
index b19066a520..9bea968a7e 100644
--- a/headers/private/kernel/arch/thread.h
+++ b/headers/private/kernel/arch/thread.h
@@ -33,6 +33,8 @@ status_t arch_setup_signal_frame(Thread *thread, struct sigaction *action,
 int64 arch_restore_signal_frame(struct signal_frame_data* signalFrameData);
 
 void arch_store_fork_frame(struct arch_fork_arg *arg);
+void arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount);
 void arch_restore_fork_frame(struct arch_fork_arg *arg);
 
 #define arch_syscall_64_bit_return_value()
diff --git a/src/system/kernel/arch/arm/arch_thread.cpp b/src/system/kernel/arch/arm/arch_thread.cpp
index 7dae8a1fa0..37d96d4af9 100644
--- a/src/system/kernel/arch/arm/arch_thread.cpp
+++ b/src/system/kernel/arch/arm/arch_thread.cpp
@@ -360,6 +360,16 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
 }
 
 
+/*!	Gets the pages that are expected to be touched immediately
+	upon restoring the fork frame.
+*/
+void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+}
+
+
 /** Restores the frame from a forked team as specified by the provided
  *	arch_fork_arg structure.
  *	Needs to be called from within the child team, ie. instead of
diff --git a/src/system/kernel/arch/arm64/arch_thread.cpp b/src/system/kernel/arch/arm64/arch_thread.cpp
index a7f152e7ae..af4924d05b 100644
--- a/src/system/kernel/arch/arm64/arch_thread.cpp
+++ b/src/system/kernel/arch/arm64/arch_thread.cpp
@@ -174,6 +174,13 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
 
 
 void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+}
+
+
+void
 arch_restore_fork_frame(struct arch_fork_arg *arg)
 {
 }
diff --git a/src/system/kernel/arch/m68k/arch_thread.cpp b/src/system/kernel/arch/m68k/arch_thread.cpp
index 0e7bcc3cde..88477f904f 100644
--- a/src/system/kernel/arch/m68k/arch_thread.cpp
+++ b/src/system/kernel/arch/m68k/arch_thread.cpp
@@ -258,13 +258,22 @@ arch_restore_signal_frame(struct signal_frame_data* signalFrameData)
  *	arch_fork_arg structure to be passed to arch_restore_fork_frame().
  *	Also makes sure to return the right value.
  */
-
 void
 arch_store_fork_frame(struct arch_fork_arg *arg)
 {
 }
 
 
+/*!	Gets the pages that are expected to be touched immediately
+	upon restoring the fork frame.
+*/
+void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+}
+
+
 /** Restores the frame from a forked team as specified by the provided
  *	arch_fork_arg structure.
  *	Needs to be called from within the child team, ie. instead of
@@ -272,7 +281,6 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
  *	This function does not return to the caller, but will enter userland
  *	in the child team at the same position where the parent team left of.
  */
-
 void
 arch_restore_fork_frame(struct arch_fork_arg *arg)
 {
diff --git a/src/system/kernel/arch/ppc/arch_thread.cpp b/src/system/kernel/arch/ppc/arch_thread.cpp
index c9856cb009..1f7e649813 100644
--- a/src/system/kernel/arch/ppc/arch_thread.cpp
+++ b/src/system/kernel/arch/ppc/arch_thread.cpp
@@ -242,6 +242,16 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
 }
 
 
+/*!	Gets the pages that are expected to be touched immediately
+	upon restoring the fork frame.
+*/
+void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+}
+
+
 /** Restores the frame from a forked team as specified by the provided
  *	arch_fork_arg structure.
  *	Needs to be called from within the child team, ie. instead of
@@ -249,9 +259,7 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
  *	This function does not return to the caller, but will enter userland
  *	in the child team at the same position where the parent team left of.
  */
-
 void
 arch_restore_fork_frame(struct arch_fork_arg *arg)
 {
 }
-
diff --git a/src/system/kernel/arch/riscv64/arch_thread.cpp b/src/system/kernel/arch/riscv64/arch_thread.cpp
index 3236c5d092..d504145d4f 100644
--- a/src/system/kernel/arch/riscv64/arch_thread.cpp
+++ b/src/system/kernel/arch/riscv64/arch_thread.cpp
@@ -313,7 +313,6 @@ arch_restore_signal_frame(struct signal_frame_data* signalFrameData)
  *	arch_fork_arg structure to be passed to arch_restore_fork_frame().
  *	Also makes sure to return the right value.
  */
-
 void
 arch_store_fork_frame(struct arch_fork_arg *arg)
 {
@@ -329,6 +328,16 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
 }
 
 
+/*!	Gets the pages that are expected to be touched immediately
+	upon restoring the fork frame.
+*/
+void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+}
+
+
 /** Restores the frame from a forked team as specified by the provided
  *	arch_fork_arg structure.
  *	Needs to be called from within the child team, ie. instead of
@@ -336,7 +345,6 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
  *	This function does not return to the caller, but will enter userland
  *	in the child team at the same position where the parent team left of.
  */
-
 void
 arch_restore_fork_frame(struct arch_fork_arg *arg)
 {
diff --git a/src/system/kernel/arch/sparc/arch_thread.cpp b/src/system/kernel/arch/sparc/arch_thread.cpp
index fd4e840926..274e0b5c55 100644
--- a/src/system/kernel/arch/sparc/arch_thread.cpp
+++ b/src/system/kernel/arch/sparc/arch_thread.cpp
@@ -142,13 +142,22 @@ arch_restore_signal_frame(struct signal_frame_data* signalFrameData)
  *	arch_fork_arg structure to be passed to arch_restore_fork_frame().
  *	Also makes sure to return the right value.
  */
-
 void
 arch_store_fork_frame(struct arch_fork_arg *arg)
 {
 }
 
 
+/*!	Gets the pages that are expected to be touched immediately
+	upon restoring the fork frame.
+*/
+void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+}
+
+
 /** Restores the frame from a forked team as specified by the provided
  *	arch_fork_arg structure.
  *	Needs to be called from within the child team, ie. instead of
@@ -156,9 +165,7 @@ arch_store_fork_frame(struct arch_fork_arg *arg)
  *	This function does not return to the caller, but will enter userland
  *	in the child team at the same position where the parent team left of.
  */
-
 void
 arch_restore_fork_frame(struct arch_fork_arg *arg)
 {
 }
-
diff --git a/src/system/kernel/arch/x86/arch_thread.cpp b/src/system/kernel/arch/x86/arch_thread.cpp
index 4af28e33a0..5b7a8acaba 100644
--- a/src/system/kernel/arch/x86/arch_thread.cpp
+++ b/src/system/kernel/arch/x86/arch_thread.cpp
@@ -267,6 +267,19 @@ arch_store_fork_frame(struct arch_fork_arg* arg)
 }
 
 
+/*!	Gets the pages that are expected to be touched immediately
+	upon restoring the fork frame.
+*/
+void
+arch_get_fork_prefault_pages(struct arch_fork_arg *arg,
+	addr_t pages[], uint32 maxCount)
+{
+	ASSERT(maxCount >= 2);
+	pages[0] = ROUNDDOWN(arg->iframe.ip, B_PAGE_SIZE) | B_EXECUTE_AREA;
+	pages[1] = ROUNDDOWN(arg->iframe.user_sp, B_PAGE_SIZE) | (B_READ_AREA | B_WRITE_AREA);
+}
+
+
 /*!	Restores the frame from a forked team as specified by the provided
 	arch_fork_arg structure.
 	Needs to be called from within the child team, i.e. instead of
diff --git a/src/system/kernel/thread.cpp b/src/system/kernel/thread.cpp
index e53755322a..b7dc6ee5dc 100644
--- a/src/system/kernel/thread.cpp
+++ b/src/system/kernel/thread.cpp
@@ -660,16 +660,20 @@ enter_userspace(Thread* thread, UserThreadEntryArguments* args)
 	userThread->pending_signals = 0;
 	arch_cpu_disable_user_access();
 
-	// initialize default TLS fields
-	addr_t tls[TLS_FIRST_FREE_SLOT];
-	memset(tls, 0, sizeof(tls));
-	tls[TLS_BASE_ADDRESS_SLOT] = thread->user_local_storage;
-	tls[TLS_THREAD_ID_SLOT] = thread->id;
-	tls[TLS_USER_THREAD_SLOT] = (addr_t)thread->user_thread;
-
 	if (args->forkArgs == NULL) {
+		// initialize default TLS fields
+		addr_t tls[TLS_FIRST_FREE_SLOT];
+		memset(tls, 0, sizeof(tls));
+		tls[TLS_BASE_ADDRESS_SLOT] = thread->user_local_storage;
+		tls[TLS_THREAD_ID_SLOT] = thread->id;
+		tls[TLS_USER_THREAD_SLOT] = (addr_t)thread->user_thread;
+
 		if (user_memcpy((void*)thread->user_local_storage, tls, sizeof(tls)) != B_OK)
 			return B_BAD_ADDRESS;
+
+		// Jump to the entry point in user space. Only returns, if something fails.
+		return arch_thread_enter_userspace(thread, args->userlandEntry,
+			args->userlandArgument1, args->userlandArgument2);
 	} else {
 		// This is a fork()ed thread.
 
@@ -677,10 +681,26 @@ enter_userspace(Thread* thread, UserThreadEntryArguments* args)
 		arch_cpu_enable_user_access();
 		addr_t* userTls = (addr_t*)thread->user_local_storage;
 		ASSERT(userTls[TLS_BASE_ADDRESS_SLOT] == thread->user_local_storage);
-		userTls[TLS_THREAD_ID_SLOT] = tls[TLS_THREAD_ID_SLOT];
-		userTls[TLS_USER_THREAD_SLOT] = tls[TLS_USER_THREAD_SLOT];
+		userTls[TLS_THREAD_ID_SLOT] = thread->id;
+		userTls[TLS_USER_THREAD_SLOT] = (addr_t)thread->user_thread;
 		arch_cpu_disable_user_access();
 
+		// Pre-fault initial pages.
+		addr_t pages[4] = {};
+		arch_get_fork_prefault_pages(args->forkArgs, pages, B_COUNT_OF(pages));
+		for (size_t i = 0; i < B_COUNT_OF(pages); i++) {
+			if (pages[i] == 0)
+				break;
+
+			addr_t pageAddress = pages[i] & ~(B_PAGE_SIZE - 1);
+			const uint32 flags = pages[i] & (B_PAGE_SIZE - 1);
+			error = vm_page_fault(pageAddress, pageAddress,
+				(flags & B_WRITE_AREA) != 0, (flags & B_EXECUTE_AREA) != 0,
+				true, &pageAddress);
+			if (error < 0)
+				return error;
+		}
+
 		// Copy the fork args onto the stack and free them.
 		arch_fork_arg archArgs = *args->forkArgs;
 		free(args->forkArgs);
@@ -689,10 +709,6 @@ enter_userspace(Thread* thread, UserThreadEntryArguments* args)
 			// this one won't return here
 		return B_ERROR;
 	}
-
-	// Jump to the entry point in user space. Only returns, if something fails.
-	return arch_thread_enter_userspace(thread, args->userlandEntry,
-		args->userlandArgument1, args->userlandArgument2);
 }
 
 
-- 
2.11.0

