HAOC: Fix conflict between IEE and KPTI while modifying ASID

community inclusion
category: bugfix

--------------------------------

Fix 2 bugs:

Bug 1: Crash upon resume from suspend.
This issue occurred due to missing the proper use of the ISB (Instruction
Synchronization Barrier) instruction after writing to the TTBR0 register via
the IEE SIP interface. The absence of ISB prevented the modification from
immediately taking effect, causing errors in subsequent code accessing
addresses within the TTBR0 range. The problem was fixed by adding an ISB
instruction in the IEE SIP gate.

Bug 2: IEE ASID usage incompatible with KPTI.
When Kernel Page Table Isolation (KPTI) was enabled along with HAOC
initialization, conflicts in the Translation Lookaside Buffer (TLB) due to ASID
handling incompatibility led to hangs during initialization. To resolve this,
the modifications by IEE to the original ARM kernel ASID handling logic were
reverted. Now, the storage of the IEE ASID has been shifted from TTBR1 to
TTBR0, while the kernel continues to use TTBR1 to store user ASIDs.

Fixes: 6d2d4fa ("HAOC: Add support for AArch64 Isolated Execution Environment(IEE).")
Signed-off-by: Lyu Jinglin <lvjl2022@zgclab.edu.cn>
Signed-off-by: Liu Zhehui <liuzhh@zgclab.edu.cn>

Author: Jinglin-lyu

arch/arm64/include/asm/haoc/iee-asm.h

@@ -17,15 +17,9 @@
 
 #define ASID_BIT		(UL(1) << 48)
 /* 
- * We reserves the bigest ASID for IEE and always stores it in TTBR1. As KPTI also reserves
- * odd ASIDs for user-viewed TTBR1, we should use even number for IEE ASID to allow KPTI to
- * switch between them at kernel entry/exit.
+ * We reserves the bigest ASID for IEE and always stores it in TTBR0.
  */
-#ifdef CONFIG_UNMAP_KERNEL_AT_EL0
-#define IEE_ASID			0xfffe
-#else
 #define IEE_ASID			0xffff
-#endif
 #define IEE_ASM_ASID		(UL(IEE_ASID) << 48)
 
 #define TCR_HPD1		(UL(1) << 42)

arch/arm64/include/asm/haoc/iee-si.h

@@ -13,6 +13,7 @@ enum {
 	IEE_SI_SET_VBAR,
 	IEE_SI_SET_TTBR0,
 	IEE_SI_CONTEXT_SWITCH,
+	IEE_SI_CONTEXT_SWITCH_PRE_INIT,
 	IEE_SI_FLAGS
 };
 
@@ -37,9 +38,4 @@ extern u64 __iee_si_text_start[];
 extern u64 __iee_si_text_end[];
 extern u64 iee_si_reserved_pg_dir;
 
-static inline void iee_si_setup_data(void)
-{
-	iee_si_reserved_pg_dir = phys_to_ttbr(__pa_symbol(reserved_pg_dir));
-}
-
 #endif

arch/arm64/include/asm/mmu_context.h

@@ -57,6 +57,10 @@ static inline void cpu_set_reserved_ttbr0_nosync(void)
 #ifdef CONFIG_IEE_SIP
 	iee_rwx_gate(IEE_SI_SET_TTBR0, ttbr);
 #else
+#ifdef CONFIG_IEE
+	if (iee_init_done)
+		ttbr |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
+#endif
 	write_sysreg(ttbr, ttbr0_el1);
 #endif
 }
@@ -166,6 +170,10 @@ static inline void cpu_install_ttbr0(phys_addr_t ttbr0, unsigned long t0sz)
 #ifdef CONFIG_IEE_SIP
 	iee_rwx_gate(IEE_SI_SET_TTBR0, ttbr0);
 #else
+#ifdef CONFIG_IEE
+	if (iee_init_done)
+		ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
+#endif
 	write_sysreg(ttbr0, ttbr0_el1);
 	isb();
 #endif
@@ -196,10 +204,6 @@ static inline void cpu_replace_ttbr1(pgd_t *pgdp, pgd_t *idmap)
 		 */
 		ttbr1 |= TTBR_CNP_BIT;
 	}
-#ifdef CONFIG_IEE
-	if (iee_init_done)
-		ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
-#endif
 
 	replace_phys = (void *)__pa_symbol(idmap_cpu_replace_ttbr1);
 

arch/arm64/kernel/haoc/iee/iee-gate.S

@@ -50,7 +50,7 @@ SYM_FUNC_START(iee_protected_rw_gate)
 	/* entry gate */
 	mrs x12, tcr_el1
 	orr x12, x12, #TCR_HPD1
-	orr x12, x12, #TCR_A1
+	bic x12, x12, #TCR_A1
 	msr tcr_el1, x12
 	isb
 	/* Check TCR */
@@ -79,7 +79,7 @@ SYM_FUNC_START(iee_protected_rw_gate)
 	/* exit gate */
 	mrs x12, tcr_el1
 	bic x12, x12, #TCR_HPD1
-	bic x12, x12, #TCR_A1
+	orr x12, x12, #TCR_A1
 	msr tcr_el1, x12
 	isb
 	/* Check TCR */

arch/arm64/kernel/haoc/iee/iee-init.c

@@ -40,28 +40,12 @@ static void __init iee_stack_alloc(void)
 	flush_tlb_all();
 }
 
-/* Setup TCR for this cpu and move ASID from ttbr1 to ttbr0 */
 void iee_setup_asid(void)
 {
-	unsigned long asid, ttbr0, ttbr1;
-
-	ttbr1 = read_sysreg(ttbr1_el1);
-	asid = FIELD_GET(TTBR_ASID_MASK, ttbr1);
-	ttbr0 = read_sysreg(ttbr0_el1) | FIELD_PREP(TTBR_ASID_MASK, asid);
-	ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
-	write_sysreg(ttbr1, ttbr1_el1);
+	unsigned long ttbr0 = read_sysreg(ttbr0_el1) | FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
+	/* Load IEE ASID into ttbr0 to use it in IEE. */
 	write_sysreg(ttbr0, ttbr0_el1);
-	write_sysreg(read_sysreg(tcr_el1) & ~TCR_A1, tcr_el1);
 	isb();
-
-	/* Flush tlb to enable IEE. */
-	local_flush_tlb_all();
-}
-
-static void iee_setup_init_data(void){
-	for (u64 addr = (u64)iee_init_data_begin; addr < (u64)iee_init_data_end;
-			addr += PAGE_SIZE)
-		iee_set_logical_mem(addr, 0, true);
 }
 
 void __init iee_init_post(void)
@@ -82,7 +66,6 @@ void __init iee_init_post(void)
 	extern void iee_si_init(void);
 	iee_si_init();
 #endif
-	iee_setup_init_data();
 }
 
 void __init iee_stack_init(void)

arch/arm64/kernel/haoc/iee/iee-mmu.c

@@ -60,12 +60,12 @@ void __init iee_init_tcr(void)
 
 	__set_fixmap(FIX_PTE, __pa_symbol(&kernel_tcr), FIXMAP_PAGE_NORMAL);
 	ptr += (unsigned long)(&kernel_tcr) & (PAGE_SIZE - 1);
-	*((u64 *)ptr) = read_sysreg(tcr_el1) & IEE_TCR_MASK & ~(TCR_HPD1 | TCR_A1);
+	*((u64 *)ptr) = read_sysreg(tcr_el1) & IEE_TCR_MASK & ~TCR_HPD1;
 	clear_fixmap(FIX_PTE);
 	ptr = (unsigned long)(fix_to_virt(FIX_PTE));
 	__set_fixmap(FIX_PTE, __pa_symbol(&iee_tcr), FIXMAP_PAGE_NORMAL);
 	ptr += (unsigned long)(&iee_tcr) & (PAGE_SIZE - 1);
-	*((u64 *)ptr) = kernel_tcr | TCR_HPD1 | TCR_A1;
+	*((u64 *)ptr) = (kernel_tcr | TCR_HPD1) & ~TCR_A1;
 	clear_fixmap(FIX_PTE);
 }
 

arch/arm64/kernel/haoc/iee/iee-si-gate.S

@@ -20,6 +20,7 @@ SYM_FUNC_START(iee_rwx_gate)
 	bl iee_si_handler
 	ldp x29, x30, [sp, #16]
 	ldr x13, [sp], #32
+	isb
 	b 2f
 1:
 	stp x29, x30, [sp, #-16]!
@@ -40,7 +41,7 @@ SYM_FUNC_START(iee_rwx_gate_tramp)
 	/* entry gate */
 	mrs x12, tcr_el1
 	orr x12, x12, #TCR_HPD1
-	orr x12, x12, #TCR_A1
+	bic x12, x12, #TCR_A1
 	msr tcr_el1, x12
 	isb
 	/* Check TCR */
@@ -66,7 +67,7 @@ SYM_FUNC_START(iee_rwx_gate_tramp)
 	/* exit gate */
 	mrs x12, tcr_el1
 	bic x12, x12, #TCR_HPD1
-	bic x12, x12, #TCR_A1
+	orr x12, x12, #TCR_A1
 	msr tcr_el1, x12
 	isb
 	/* Check TCR */

arch/arm64/kernel/haoc/iee/iee-si.c

@@ -16,10 +16,30 @@ static inline unsigned long iee_si_mask(unsigned long mask,
 	return (new_val & mask) | (old_val & ~mask);
 }
 
+#ifdef CONFIG_IEE_PTRP
+static inline void iee_si_check_ttbr0(void)
+{
+	u64 old_ttbr0 = read_sysreg(ttbr0_el1);
+	u64 old_phys = (old_ttbr0 & PAGE_MASK) & ~TTBR_ASID_MASK;
+	struct task_token *token = (struct task_token *)__addr_to_iee(current);
+	/* Phys in TTBR0 shall be the same with current->mm->pgd. */
+	if (!(current == &init_task) && token->pgd
+		&& old_phys != iee_si_reserved_pg_dir) {
+		u64 token_phys = __pa(token->pgd);
+
+		if (old_phys != token_phys)
+			pr_err("IEE: Pgd set error. old ttbr0:%llx, token ttbr0:%llx",
+				old_phys, token_phys);
+	}
+}
+#else
+static inline void iee_si_check_ttbr0(void) {}
+#endif
+
 unsigned long __iee_si_code iee_si_handler(int flag, ...)
 {
 	va_list pArgs;
-	unsigned long old_val, new_val;
+	unsigned long old_val, new_val, ttbr1, ttbr0;
 
 	va_start(pArgs, flag);
 	switch (flag) {
@@ -32,33 +52,38 @@ unsigned long __iee_si_code iee_si_handler(int flag, ...)
 		write_sysreg(new_val, sctlr_el1);
 		break;
 	case IEE_SI_SET_TTBR0:
+		ttbr0 = va_arg(pArgs, u64) & ~TTBR_ASID_MASK;
+		/* Skip checking before init IEE. */
+		if (iee_init_done)
+			iee_si_check_ttbr0();
+		/* Load the reserved IEE ASID into TTBR0.*/
+		ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
+		write_sysreg(ttbr0, ttbr0_el1);
+		break;
 	case IEE_SI_CONTEXT_SWITCH:
-		old_val = read_sysreg(ttbr0_el1);
-		new_val = va_arg(pArgs, u64);
+		ttbr1 = va_arg(pArgs, u64);
+		ttbr0 = va_arg(pArgs, u64) & ~TTBR_ASID_MASK;
+		ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
 		/* Skip checking before init IEE. */
-		if (iee_tcr != 0) {
-			u64 new_asid;
+		if (iee_init_done) {
+			u64 new_asid = ttbr1 >> 48;
 			/* ASID shall not be the one reserved for IEE. */
-			new_asid = new_val >> 48;
 			if (new_asid == IEE_ASID)
 				panic("IEE SI: Using reserved IEE ASID in TTRB0.");
 
-			#ifdef CONFIG_IEE_PTRP
-			u64 old_phys;
-			struct task_token *token = (struct task_token *)__addr_to_iee(current);
-			/* Phys in TTBR0 shall be the same with current->mm->pgd. */
-			old_phys = (old_val & PAGE_MASK) & ~TTBR_ASID_MASK;
-			if (!(current == &init_task) && token->pgd
-				&& old_phys != iee_si_reserved_pg_dir) {
-				u64 token_phys = __pa(token->pgd);
-
-				if (old_phys != token_phys)
-					panic("IEE: Pgd set error. old ttbr0:%llx, token ttbr0:%llx",
-						old_phys, token_phys);
-			}
-			#endif
+			iee_si_check_ttbr0();
 		}
-		write_sysreg(new_val, ttbr0_el1);
+		write_sysreg(ttbr1, ttbr1_el1);
+		write_sysreg(ttbr0, ttbr0_el1);
+		break;
+	case IEE_SI_CONTEXT_SWITCH_PRE_INIT:
+		ttbr1 = va_arg(pArgs, u64);
+		ttbr0 = va_arg(pArgs, u64);
+
+		if (iee_init_done)
+			panic("IEE: Using legacy context switch after IEE init.");
+		write_sysreg(ttbr1, ttbr1_el1);
+		write_sysreg(ttbr0, ttbr0_el1);
 		break;
 	case IEE_SI_SET_VBAR:
 		new_val = va_arg(pArgs, u64);
@@ -77,6 +102,11 @@ unsigned long __iee_si_code iee_si_handler(int flag, ...)
 	return 0;
 }
 
+static inline void iee_si_setup_data(void)
+{
+	iee_si_reserved_pg_dir = phys_to_ttbr(__pa_symbol(reserved_pg_dir));
+}
+
 /* Map iee si code PXNTable=1 on pmd page tables. */
 static int __init iee_si_init_code(void)
 {

arch/arm64/mm/context.c

@@ -368,30 +368,34 @@ void cpu_do_switch_mm(phys_addr_t pgd_phys, struct mm_struct *mm)
 		ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, asid);
 
 #ifdef CONFIG_IEE
-	if (iee_init_done) {
-		/*
-		 * IEE requires the reserved ASID stored in TTBR1 and User ASID stored
-		 * in TTBR0 to support ASID switch by changing TCR.A1.
-		 */
-		ttbr0 &= ~TTBR_ASID_MASK;
-		ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, asid);
+	/* Set ASID in TTBR1 since TCR.A1 is set */
+	ttbr1 &= ~TTBR_ASID_MASK;
+	ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, asid);
 
-		cpu_set_reserved_ttbr0_nosync();
+	cpu_set_reserved_ttbr0_nosync();
+	/*
+	 * IEE requires the reserved IEE ASID stored in TTBR0_EL1 so that IEE can
+	 * switch the ASID by setting TCR_EL1.A1 from 1 to 0 in IEE gate to avoid
+	 * tlb entry disclosure when calling IEE functions.
+	 */
+	if (iee_init_done) {
 		#ifdef CONFIG_IEE_SIP
-		iee_rwx_gate(IEE_SI_CONTEXT_SWITCH, ttbr0);
+		iee_rwx_gate(IEE_SI_CONTEXT_SWITCH, ttbr1, ttbr0);
 		#else
+		ttbr0 &= ~TTBR_ASID_MASK;
+		ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, IEE_ASID);
+		write_sysreg(ttbr1, ttbr1_el1);
 		write_sysreg(ttbr0, ttbr0_el1);
 		isb();
 		#endif
 	} else {
-		/* Set ASID in TTBR1 since TCR.A1 is set */
-		ttbr1 &= ~TTBR_ASID_MASK;
-		ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, asid);
-
-		cpu_set_reserved_ttbr0_nosync();
+		#ifdef CONFIG_IEE_SIP
+		iee_rwx_gate(IEE_SI_CONTEXT_SWITCH_PRE_INIT, ttbr1, ttbr0);
+		#else
 		write_sysreg(ttbr1, ttbr1_el1);
 		write_sysreg(ttbr0, ttbr0_el1);
 		isb();
+		#endif
 	}
 #else
 	/* Set ASID in TTBR1 since TCR.A1 is set */
@@ -417,7 +421,7 @@ static int asids_update_limit(void)
 #ifdef CONFIG_IEE
 		if (haoc_enabled && pinned_asid_map) {
 			__set_bit(ctxid2asid(IEE_ASID), pinned_asid_map);
-			__set_bit(ctxid2asid(IEE_ASID | ASID_BIT), pinned_asid_map);
+			__set_bit(ctxid2asid(IEE_ASID & ~ASID_BIT), pinned_asid_map);
 		}
 #endif
 	}
@@ -461,8 +465,8 @@ static int asids_init(void)
 #ifdef CONFIG_IEE
 	if (haoc_enabled) {
 		#ifdef CONFIG_UNMAP_KERNEL_AT_EL0
-		__set_bit(ctxid2asid(IEE_ASID | ASID_BIT), asid_map);
-		__set_bit(ctxid2asid(IEE_ASID | ASID_BIT), pinned_asid_map);
+		__set_bit(ctxid2asid(IEE_ASID & ~ASID_BIT), asid_map);
+		__set_bit(ctxid2asid(IEE_ASID & ~ASID_BIT), pinned_asid_map);
 		#endif
 		__set_bit(ctxid2asid(IEE_ASID), asid_map);
 		__set_bit(ctxid2asid(IEE_ASID), pinned_asid_map);