arm64: move preemption disablement to prctl handlers

In the next patch, we will start reading sctlr_user from
mte_update_sctlr_user and subsequently writing a new value based on the
task's TCF setting and potentially the per-CPU TCF preference. This
means that we need to be careful to disable preemption around any
code sequences that read from sctlr_user and subsequently write to
sctlr_user and/or SCTLR_EL1, so that we don't end up writing a stale
value (based on the previous CPU's TCF preference) to either of them.

We currently have four such sequences, in the prctl handlers for
PR_SET_TAGGED_ADDR_CTRL and PR_PAC_SET_ENABLED_KEYS, as well as in
the task initialization code that resets the prctl settings. Change
the prctl handlers to disable preemption in the handlers themselves
rather than the functions that they call, and change the task
initialization code to call the respective prctl handlers instead of
setting sctlr_user directly.

As a result of this change, we no longer need the helper function
set_task_sctlr_el1, nor does its behavior make sense any more, so
remove it.

Signed-off-by: Peter Collingbourne <[e-mail]>
Link: https://linux-review.googlesource.com/id/Ic0e8a0c00bb47d786c1e8011df0b7fe99bee4bb5
Link: https://lore.kernel.org/r/[e-mail]
Acked-by: Will Deacon <[e-mail]>
Signed-off-by: Catalin Marinas <[e-mail]>

arch/arm64/include/asm/pointer_auth.h

···
 #include <asm/memory.h>
 #include <asm/sysreg.h>
 
+#define PR_PAC_ENABLED_KEYS_MASK                                               \
+	(PR_PAC_APIAKEY | PR_PAC_APIBKEY | PR_PAC_APDAKEY | PR_PAC_APDBKEY)
+
 #ifdef CONFIG_ARM64_PTR_AUTH
 /*
  * Each key is a 128-bit quantity which is split across a pair of 64-bit
···
 									       \
 		/* enable all keys */                                          \
 		if (system_supports_address_auth())                            \
-			set_task_sctlr_el1(current->thread.sctlr_user |        \
-					   SCTLR_ELx_ENIA | SCTLR_ELx_ENIB |   \
-					   SCTLR_ELx_ENDA | SCTLR_ELx_ENDB);   \
+			ptrauth_set_enabled_keys(current,                      \
+						 PR_PAC_ENABLED_KEYS_MASK,     \
+						 PR_PAC_ENABLED_KEYS_MASK);    \
 	} while (0)
 
 #define ptrauth_thread_switch_user(tsk)                                        \
···
 #define ptrauth_thread_init_kernel(tsk)
 #define ptrauth_thread_switch_kernel(tsk)
 #endif /* CONFIG_ARM64_PTR_AUTH_KERNEL */
-
-#define PR_PAC_ENABLED_KEYS_MASK                                               \
-	(PR_PAC_APIAKEY | PR_PAC_APIBKEY | PR_PAC_APDAKEY | PR_PAC_APDBKEY)
 
 #endif /* __ASM_POINTER_AUTH_H */

arch/arm64/include/asm/processor.h

···
 
 unsigned long get_wchan(struct task_struct *p);
 
-void set_task_sctlr_el1(u64 sctlr);
+void update_sctlr_el1(u64 sctlr);
 
 /* Thread switching */
 extern struct task_struct *cpu_switch_to(struct task_struct *prev,

arch/arm64/kernel/mte.c

···
 	write_sysreg_s(0, SYS_TFSRE0_EL1);
 	clear_thread_flag(TIF_MTE_ASYNC_FAULT);
 	/* disable tag checking and reset tag generation mask */
-	current->thread.mte_ctrl = MTE_CTRL_GCR_USER_EXCL_MASK;
-	mte_update_sctlr_user(current);
-	set_task_sctlr_el1(current->thread.sctlr_user);
+	set_mte_ctrl(current, 0);
 }
 
 void mte_thread_switch(struct task_struct *next)
···
 
 	task->thread.mte_ctrl = mte_ctrl;
 	if (task == current) {
+		preempt_disable();
 		mte_update_sctlr_user(task);
-		set_task_sctlr_el1(task->thread.sctlr_user);
+		update_sctlr_el1(task->thread.sctlr_user);
+		preempt_enable();
 	}
 
 	return 0;

arch/arm64/kernel/pointer_auth.c

···
 int ptrauth_set_enabled_keys(struct task_struct *tsk, unsigned long keys,
 			     unsigned long enabled)
 {
-	u64 sctlr = tsk->thread.sctlr_user;
+	u64 sctlr;
 
 	if (!system_supports_address_auth())
 		return -EINVAL;
···
 	if ((keys & ~PR_PAC_ENABLED_KEYS_MASK) || (enabled & ~keys))
 		return -EINVAL;
 
+	preempt_disable();
+	sctlr = tsk->thread.sctlr_user;
 	sctlr &= ~arg_to_enxx_mask(keys);
 	sctlr |= arg_to_enxx_mask(enabled);
+	tsk->thread.sctlr_user = sctlr;
 	if (tsk == current)
-		set_task_sctlr_el1(sctlr);
-	else
-		tsk->thread.sctlr_user = sctlr;
+		update_sctlr_el1(sctlr);
+	preempt_enable();
 
 	return 0;
 }

arch/arm64/kernel/process.c

···
 		set_tsk_thread_flag(next, TIF_NOTIFY_RESUME);
 }
 
-static void update_sctlr_el1(u64 sctlr)
+/*
+ * __switch_to() checks current->thread.sctlr_user as an optimisation. Therefore
+ * this function must be called with preemption disabled and the update to
+ * sctlr_user must be made in the same preemption disabled block so that
+ * __switch_to() does not see the variable update before the SCTLR_EL1 one.
+ */
+void update_sctlr_el1(u64 sctlr)
 {
 	/*
 	 * EnIA must not be cleared while in the kernel as this is necessary for
···
 
 	/* ISB required for the kernel uaccess routines when setting TCF0. */
 	isb();
-}
-
-void set_task_sctlr_el1(u64 sctlr)
-{
-	/*
-	 * __switch_to() checks current->thread.sctlr as an
-	 * optimisation. Disable preemption so that it does not see
-	 * the variable update before the SCTLR_EL1 one.
-	 */
-	preempt_disable();
-	current->thread.sctlr_user = sctlr;
-	update_sctlr_el1(sctlr);
-	preempt_enable();
 }
 
 /*