diff --git a/arch/arm64/include/asm/haoc/iee.h b/arch/arm64/include/asm/haoc/iee.h
index 7137aa186d29a..9e653dd34c92e 100644
--- a/arch/arm64/include/asm/haoc/iee.h
+++ b/arch/arm64/include/asm/haoc/iee.h
@@ -66,7 +66,7 @@ extern bool haoc_enabled;
 
 #define TCR_HPD1		(UL(1) << 42)
 
-void iee_init_mappings(pgd_t *pgdp);
+void __init iee_init_mappings(pgd_t *pgdp);
 void iee_init_post(void);
 void iee_stack_init(void);
 void iee_init_tcr(void);
diff --git a/arch/arm64/kernel/haoc/Kconfig b/arch/arm64/kernel/haoc/Kconfig
index 7aa900608650f..b2fe32a2f0701 100644
--- a/arch/arm64/kernel/haoc/Kconfig
+++ b/arch/arm64/kernel/haoc/Kconfig
@@ -34,6 +34,7 @@ config IEE_SIP
 config CREDP
 	bool "Struct cred protection(CREDP)"
 	depends on IEE
+	select IEE_PTRP
 	help
 		Protects kernel struct cred. All modifications of cred must be made and
 		verified by IEE APIs, and critical dereferences of cred would be monitored
diff --git a/arch/arm64/kernel/haoc/iee/iee-mmu.c b/arch/arm64/kernel/haoc/iee/iee-mmu.c
index c7b227907c3f0..b3e73381d3461 100644
--- a/arch/arm64/kernel/haoc/iee/iee-mmu.c
+++ b/arch/arm64/kernel/haoc/iee/iee-mmu.c
@@ -487,7 +487,7 @@ static void __init __create_pgd_mapping_for_iee_locked(pgd_t *pgdir, phys_addr_t
 	} while (pgdp++, addr = next, addr != end);
 }
 
-static void __create_pgd_mapping_for_iee(pgd_t *pgdir, phys_addr_t phys,
+static void __init __create_pgd_mapping_for_iee(pgd_t *pgdir, phys_addr_t phys,
 				 unsigned long virt, phys_addr_t size,
 				 pgprot_t prot,
 				 phys_addr_t (*pgtable_alloc)(int),
diff --git a/arch/x86/Kconfig b/arch/x86/Kconfig
index 91685de954096..a8b1123c2b142 100644
--- a/arch/x86/Kconfig
+++ b/arch/x86/Kconfig
@@ -1576,6 +1576,7 @@ config IEE_SIP
 config CREDP
 	bool "Struct cred protection(CREDP)"
 	depends on IEE
+	select IEE_PTRP
 	help
 		Protects kernel struct cred. All modifications of cred must be made and
 		verified by IEE APIs, and critical dereferences of cred would be monitored
diff --git a/arch/x86/boot/compressed/ident_map_64.c b/arch/x86/boot/compressed/ident_map_64.c
index 848a958c639f8..8ecb4d40e20dd 100644
--- a/arch/x86/boot/compressed/ident_map_64.c
+++ b/arch/x86/boot/compressed/ident_map_64.c
@@ -180,14 +180,7 @@ void initialize_identity_maps(void *rmode)
 	sev_prep_identity_maps(top_level_pgt);
 
 	/* Load the new page-table. */
-	#ifdef CONFIG_IEE_SIP
-	if(haoc_enabled)
-		iee_write_cr3_early(top_level_pgt);
-	else
-		write_cr3(top_level_pgt);
-	#else
 	write_cr3(top_level_pgt);
-	#endif
 
 	/*
 	 * Now that the required page table mappings are established and a
@@ -231,14 +224,7 @@ static pte_t *split_large_pmd(struct x86_mapping_info *info,
 	pmd = __pmd((unsigned long)pte | info->kernpg_flag);
 	set_pmd(pmdp, pmd);
 	/* Flush TLB to establish the new PMD */
-	#ifdef CONFIG_IEE_SIP
-	if(haoc_enabled)
-		iee_write_cr3_early(top_level_pgt);
-	else
-		write_cr3(top_level_pgt);
-	#else
 	write_cr3(top_level_pgt);
-	#endif
 
 	return pte + pte_index(__address);
 }
@@ -339,14 +325,7 @@ static int set_clr_page_flags(struct x86_mapping_info *info,
 		snp_set_page_private(__pa(address & PAGE_MASK));
 
 	/* Flush TLB after changing encryption attribute */
-	#ifdef CONFIG_IEE_SIP
-	if(haoc_enabled)
-		iee_write_cr3_early(top_level_pgt);
-	else
-		write_cr3(top_level_pgt);
-	#else
 	write_cr3(top_level_pgt);
-	#endif
 
 	return 0;
 }
diff --git a/arch/x86/boot/compressed/pgtable_64.c b/arch/x86/boot/compressed/pgtable_64.c
index ab495a0b16f22..dc21b0d89d667 100644
--- a/arch/x86/boot/compressed/pgtable_64.c
+++ b/arch/x86/boot/compressed/pgtable_64.c
@@ -200,14 +200,7 @@ asmlinkage void configure_5level_paging(struct boot_params *bp, void *pgtable)
 	 * Move the top level page table out of trampoline memory.
 	 */
 	memcpy(pgtable, trampoline_32bit, PAGE_SIZE);
-	#ifdef CONFIG_IEE_SIP
-	if(haoc_enabled)
-		iee_write_cr3_early((unsigned long)pgtable);
-	else
-		native_write_cr3((unsigned long)pgtable);
-	#else
 	native_write_cr3((unsigned long)pgtable);
-	#endif
 
 	/* Restore trampoline memory */
 	memcpy(trampoline_32bit, trampoline_save, TRAMPOLINE_32BIT_SIZE);
diff --git a/arch/x86/include/asm/desc.h b/arch/x86/include/asm/desc.h
index 4a90cf3c78dc7..94a8674872fd8 100644
--- a/arch/x86/include/asm/desc.h
+++ b/arch/x86/include/asm/desc.h
@@ -13,7 +13,7 @@
 #include <linux/smp.h>
 #include <linux/percpu.h>
 
-#ifdef CONFIG_IEE_SIP
+#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 #include <asm/haoc/iee-si.h>
 extern bool haoc_enabled;
 #endif
@@ -215,7 +215,7 @@ static inline void native_load_gdt(const struct desc_ptr *dtr)
 	asm volatile("lgdt %0"::"m" (*dtr));
 }
 
-#ifdef CONFIG_IEE_SIP
+#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 static __always_inline void iee_load_idt_early(const struct desc_ptr *dtr)
 {
 	asm volatile("lidt %0"::"m" (*dtr));
@@ -224,7 +224,7 @@ static __always_inline void iee_load_idt_early(const struct desc_ptr *dtr)
 
 static __always_inline void native_load_idt(const struct desc_ptr *dtr)
 {
-	#ifdef CONFIG_IEE_SIP
+	#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 	if(haoc_enabled)
 		iee_load_idt((void *)dtr);
 	else
@@ -261,7 +261,7 @@ static inline void native_idt_invalidate(void)
 		.size = 0
 	};
 
-#ifdef CONFIG_IEE_SIP
+#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 	// The native_idt_invalidate() is only called by machine_kexec(). 
 	// In the kdump path, IEE_SIP should not be used, so we directly
 	// call iee_load_idt_early(), which is the original version of
diff --git a/arch/x86/include/asm/haoc/haoc.h b/arch/x86/include/asm/haoc/haoc.h
index ef665161259de..53526217cb5f6 100644
--- a/arch/x86/include/asm/haoc/haoc.h
+++ b/arch/x86/include/asm/haoc/haoc.h
@@ -75,4 +75,8 @@ void _iee_set_cred_security(unsigned long __unused, struct cred *cred,void *secu
 void _iee_set_cred_rcu(unsigned long __unused, struct cred *cred, struct rcu_head *rcu);
 void _iee_set_cred_ucounts(unsigned long __unused, struct cred *cred, struct ucounts *ucounts);
 #endif
+
+/* Called from iee_rw_gate (asm) to dispatch IEE ops to the _iee_*() functions. */
+unsigned long iee_dispatch(int flag, unsigned long arg1,
+			   unsigned long arg2, unsigned long arg3);
 #endif
diff --git a/arch/x86/include/asm/haoc/iee-si.h b/arch/x86/include/asm/haoc/iee-si.h
index 3baf1813f19d3..07b5400039ff6 100644
--- a/arch/x86/include/asm/haoc/iee-si.h
+++ b/arch/x86/include/asm/haoc/iee-si.h
@@ -7,9 +7,10 @@
 #define __iee_si_code   __section(".iee.si_text")
 #define __iee_si_data   __section(".iee.si_data")
 
-extern unsigned long cr4_pinned_mask;
+extern const unsigned long cr4_pinned_mask;
 extern struct static_key_false cr_pinning;
 extern unsigned long cr4_pinned_bits;
+extern unsigned long iee_cr4_set_mask;
 
 extern unsigned long __iee_si_text_start[];
 extern unsigned long __iee_si_text_end[];
diff --git a/arch/x86/include/asm/haoc/iee-token.h b/arch/x86/include/asm/haoc/iee-token.h
index 53e75396577d3..ff2e18822b6fa 100644
--- a/arch/x86/include/asm/haoc/iee-token.h
+++ b/arch/x86/include/asm/haoc/iee-token.h
@@ -10,6 +10,7 @@ extern unsigned long long iee_rw_gate(int flag, ...);
 
 extern struct kmem_cache *task_struct_cachep;
 
+extern void iee_prepare_init_task_token(void);
 extern void iee_set_token_page_valid(unsigned long token, unsigned long new,
 				     unsigned int order);
 extern void iee_set_token_page_invalid(unsigned long token_addr,
diff --git a/arch/x86/include/asm/special_insns.h b/arch/x86/include/asm/special_insns.h
index 1ab08117dea8c..0a3737ba7a7c9 100644
--- a/arch/x86/include/asm/special_insns.h
+++ b/arch/x86/include/asm/special_insns.h
@@ -9,7 +9,7 @@
 #include <asm/processor-flags.h>
 #include <linux/irqflags.h>
 #include <linux/jump_label.h>
-#ifdef CONFIG_IEE_SIP
+#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 #include <asm/haoc/iee-si.h>
 extern bool haoc_enabled;
 #endif
@@ -53,7 +53,7 @@ static inline unsigned long __native_read_cr3(void)
 	return val;
 }
 
-#ifdef CONFIG_IEE_SIP
+#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 static inline void iee_write_cr3_early(unsigned long val)
 {
 	asm volatile("mov %0,%%cr3" : : "r" (val) : "memory");
@@ -62,7 +62,7 @@ static inline void iee_write_cr3_early(unsigned long val)
  
 static inline void native_write_cr3(unsigned long val)
 {
-	#ifdef CONFIG_IEE_SIP
+	#if defined(CONFIG_IEE_SIP) && !defined(__DISABLE_EXPORTS)
 	if(haoc_enabled)
 		iee_write_cr3(val);
 	else
diff --git a/arch/x86/kernel/cpu/common.c b/arch/x86/kernel/cpu/common.c
index d9e596eae225c..778fa180e0ef1 100644
--- a/arch/x86/kernel/cpu/common.c
+++ b/arch/x86/kernel/cpu/common.c
@@ -406,7 +406,7 @@ static __always_inline void setup_umip(struct cpuinfo_x86 *c)
 
 /* These bits should not change their value after CPU init is finished. */
 #ifdef CONFIG_IEE_SIP
-unsigned long cr4_pinned_mask =
+const unsigned long cr4_pinned_mask =
 	X86_CR4_SMEP | X86_CR4_SMAP | X86_CR4_UMIP |
 	X86_CR4_FSGSBASE | X86_CR4_CET;
 DEFINE_STATIC_KEY_FALSE_RO(cr_pinning);
@@ -2088,6 +2088,15 @@ static void identify_cpu(struct cpuinfo_x86 *c)
 	/* Disable the PN if appropriate */
 	squash_the_stupid_serial_number(c);
 
+#ifdef CONFIG_IEE_SIP
+	/*
+	 * The IEE rwx gate re-enables SMEP/SMAP from this mask; keep it in
+	 * sync with what setup_smep()/setup_smap() are about to enable.
+	 */
+	iee_cr4_set_mask = (cpu_has(c, X86_FEATURE_SMEP) ? X86_CR4_SMEP : 0) |
+			   (cpu_has(c, X86_FEATURE_SMAP) ? X86_CR4_SMAP : 0);
+#endif /* CONFIG_IEE_SIP */
+
 	/* Set up SMEP/SMAP/UMIP */
 	setup_smep(c);
 	setup_smap(c);
diff --git a/arch/x86/kernel/haoc/haoc.c b/arch/x86/kernel/haoc/haoc.c
index b9576af3fc913..8fbc3268b4884 100644
--- a/arch/x86/kernel/haoc/haoc.c
+++ b/arch/x86/kernel/haoc/haoc.c
@@ -7,53 +7,186 @@
  *          Hu Bing <hubing2023@zgclab.edu.cn>
  */
 
+#include <linux/ftrace.h>
+#include <linux/panic.h>
 #include <asm/haoc/haoc.h>
+#include <asm/haoc/haoc-def.h>
 
-typedef void (*iee_func)(void);
-iee_func iee_funcs[] = {
-	(iee_func)_iee_memcpy,
-	(iee_func)_iee_memset,
-	(iee_func)_iee_set_freeptr,
-	(iee_func)_iee_test_and_clear_bit,
+/*
+ * IEE function dispatcher.
+ *
+ * Called from iee_rw_gate (asm) on the per-cpu IEE stack with interrupts
+ * disabled and CR0.WP cleared. The first argument of every _iee_*()
+ * function absorbs @flag, keeping the register convention of the old
+ * indirect function table.
+ *
+ * Direct calls only: no indirect branch (retpoline/objtool clean) and no
+ * dispatch table in writable memory.
+ */
+notrace unsigned long iee_dispatch(int flag, unsigned long arg1,
+				   unsigned long arg2, unsigned long arg3)
+{
+	switch (flag) {
+	case IEE_OP_MEMCPY:
+		_iee_memcpy(flag, (void *)arg1, (void *)arg2, arg3);
+		break;
+	case IEE_OP_MEMSET:
+		_iee_memset(flag, (void *)arg1, (int)arg2, arg3);
+		break;
+	case IEE_OP_SET_FREEPTR:
+		_iee_set_freeptr(flag, (void **)arg1, (void *)arg2);
+		break;
+	case IEE_OP_TEST_CLEAR_BIT:
+		return _iee_test_and_clear_bit(flag, (long)arg1,
+					       (unsigned long *)arg2);
 #ifdef CONFIG_IEE_PTRP
-	(iee_func)_iee_set_token_pgd,
-	(iee_func)_iee_invalidate_token,
-	(iee_func)_iee_validate_token,
+	case IEE_OP_SET_TOKEN_PGD:
+		_iee_set_token_pgd(flag, (struct task_struct *)arg1,
+				   (pgd_t *)arg2);
+		break;
+	case IEE_OP_INVALIDATE_TOKEN:
+		_iee_invalidate_token(flag, (struct task_struct *)arg1);
+		break;
+	case IEE_OP_VALIDATE_TOKEN:
+		_iee_validate_token(flag, (struct task_struct *)arg1);
+		break;
 #endif
 #ifdef CONFIG_CREDP
-	(iee_func)_iee_copy_cred,
-	(iee_func)_iee_copy_kernel_cred,
-	(iee_func)_iee_init_copied_cred,
-	(iee_func)_iee_commit_creds,
-	(iee_func)_iee_abort_cred,
-	(iee_func)_iee_set_cred_uid,
-	(iee_func)_iee_set_cred_gid,
-	(iee_func)_iee_set_cred_suid,
-	(iee_func)_iee_set_cred_sgid,
-	(iee_func)_iee_set_cred_euid,
-	(iee_func)_iee_set_cred_egid,
-	(iee_func)_iee_set_cred_fsuid,
-	(iee_func)_iee_set_cred_fsgid,
-	(iee_func)_iee_set_cred_user,
-	(iee_func)_iee_set_cred_user_ns,
-	(iee_func)_iee_set_cred_group_info,
-	(iee_func)_iee_set_cred_securebits,
-	(iee_func)_iee_set_cred_cap_inheritable,
-	(iee_func)_iee_set_cred_cap_permitted,
-	(iee_func)_iee_set_cred_cap_effective,
-	(iee_func)_iee_set_cred_cap_bset,
-	(iee_func)_iee_set_cred_cap_ambient,
-	(iee_func)_iee_set_cred_jit_keyring,
-	(iee_func)_iee_set_cred_session_keyring,
-	(iee_func)_iee_set_cred_process_keyring,
-	(iee_func)_iee_set_cred_thread_keyring,
-	(iee_func)_iee_set_cred_request_key_auth,
-	(iee_func)_iee_set_cred_non_rcu,
-	(iee_func)_iee_set_cred_atomic_set_usage,
-	(iee_func)_iee_set_cred_atomic_op_usage,
-	(iee_func)_iee_set_cred_security,
-	(iee_func)_iee_set_cred_rcu,
-	(iee_func)_iee_set_cred_ucounts,
+	case IEE_OP_COPY_CRED:
+		_iee_copy_cred(flag, (struct cred *)arg1);
+		break;
+	case IEE_OP_COPY_KERNEL_CRED:
+		_iee_copy_kernel_cred(flag, (const struct cred *)arg1,
+				      (struct cred *)arg2);
+		break;
+	case IEE_OP_INIT_COPIED_CRED:
+		_iee_init_copied_cred(flag, (struct task_struct *)arg1,
+				      (struct cred *)arg2);
+		break;
+	case IEE_OP_COMMIT_CRED:
+		_iee_commit_creds(flag, (const struct cred *)arg1);
+		break;
+	case IEE_OP_ABORT_CRED:
+		_iee_abort_cred(flag, (const struct cred *)arg1);
+		break;
+	case IEE_OP_SET_CRED_UID:
+		_iee_set_cred_uid(flag, (struct cred *)arg1,
+				  KUIDT_INIT((uid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_GID:
+		_iee_set_cred_gid(flag, (struct cred *)arg1,
+				  KGIDT_INIT((gid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_SUID:
+		_iee_set_cred_suid(flag, (struct cred *)arg1,
+				   KUIDT_INIT((uid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_SGID:
+		_iee_set_cred_sgid(flag, (struct cred *)arg1,
+				   KGIDT_INIT((gid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_EUID:
+		_iee_set_cred_euid(flag, (struct cred *)arg1,
+				   KUIDT_INIT((uid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_EGID:
+		_iee_set_cred_egid(flag, (struct cred *)arg1,
+				   KGIDT_INIT((gid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_FSUID:
+		_iee_set_cred_fsuid(flag, (struct cred *)arg1,
+				    KUIDT_INIT((uid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_FSGID:
+		_iee_set_cred_fsgid(flag, (struct cred *)arg1,
+				    KGIDT_INIT((gid_t)arg2));
+		break;
+	case IEE_OP_SET_CRED_USER:
+		_iee_set_cred_user(flag, (struct cred *)arg1,
+				   (struct user_struct *)arg2);
+		break;
+	case IEE_OP_SET_CRED_USER_NS:
+		_iee_set_cred_user_ns(flag, (struct cred *)arg1,
+				      (struct user_namespace *)arg2);
+		break;
+	case IEE_OP_SET_CRED_GROUP_INFO:
+		_iee_set_cred_group_info(flag, (struct cred *)arg1,
+					 (struct group_info *)arg2);
+		break;
+	case IEE_OP_SET_CRED_SECUREBITS:
+		_iee_set_cred_securebits(flag, (struct cred *)arg1,
+					 (unsigned int)arg2);
+		break;
+	case IEE_OP_SET_CRED_CAP_INHER:
+		_iee_set_cred_cap_inheritable(flag, (struct cred *)arg1,
+					      (kernel_cap_t){ .val = arg2 });
+		break;
+	case IEE_OP_SET_CRED_CAP_PERM:
+		_iee_set_cred_cap_permitted(flag, (struct cred *)arg1,
+					    (kernel_cap_t){ .val = arg2 });
+		break;
+	case IEE_OP_SET_CRED_CAP_EFFECT:
+		_iee_set_cred_cap_effective(flag, (struct cred *)arg1,
+					    (kernel_cap_t){ .val = arg2 });
+		break;
+	case IEE_OP_SET_CRED_CAP_BSET:
+		_iee_set_cred_cap_bset(flag, (struct cred *)arg1,
+				       (kernel_cap_t){ .val = arg2 });
+		break;
+	case IEE_OP_SET_CRED_CAP_AMBIENT:
+		_iee_set_cred_cap_ambient(flag, (struct cred *)arg1,
+					  (kernel_cap_t){ .val = arg2 });
+		break;
+	case IEE_OP_SET_CRED_JIT_KEYRING:
+		_iee_set_cred_jit_keyring(flag, (struct cred *)arg1,
+					  (unsigned char)arg2);
+		break;
+	case IEE_OP_SET_CRED_SESS_KEYRING:
+		_iee_set_cred_session_keyring(flag, (struct cred *)arg1,
+					      (struct key *)arg2);
+		break;
+	case IEE_OP_SET_CRED_PROC_KEYRING:
+		_iee_set_cred_process_keyring(flag, (struct cred *)arg1,
+					      (struct key *)arg2);
+		break;
+	case IEE_OP_SET_CRED_THREAD_KEYRING:
+		_iee_set_cred_thread_keyring(flag, (struct cred *)arg1,
+					     (struct key *)arg2);
+		break;
+	case IEE_OP_SET_CRED_REQ_KEYRING:
+		_iee_set_cred_request_key_auth(flag, (struct cred *)arg1,
+					       (struct key *)arg2);
+		break;
+	case IEE_OP_SET_CRED_NON_RCU:
+		_iee_set_cred_non_rcu(flag, (struct cred *)arg1, (int)arg2);
+		break;
+	case IEE_OP_SET_CRED_ATSET_USAGE:
+		_iee_set_cred_atomic_set_usage(flag, (struct cred *)arg1,
+					       (int)arg2);
+		break;
+	case IEE_OP_SET_CRED_ATOP_USAGE:
+		return _iee_set_cred_atomic_op_usage(flag, (struct cred *)arg1,
+						     (int)arg2, (int)arg3);
+	case IEE_OP_SET_CRED_SECURITY:
+		_iee_set_cred_security(flag, (struct cred *)arg1, (void *)arg2);
+		break;
+	case IEE_OP_SET_CRED_RCU:
+		_iee_set_cred_rcu(flag, (struct cred *)arg1,
+				  (struct rcu_head *)arg2);
+		break;
+	case IEE_OP_SET_CRED_UCOUNTS:
+		_iee_set_cred_ucounts(flag, (struct cred *)arg1,
+				      (struct ucounts *)arg2);
+		break;
 #endif
-	NULL
-};
+	default:
+		/*
+		 * Unreachable in practice: all flags are compile-time
+		 * constants in the iee_*() wrappers. An unknown flag means
+		 * memory corruption or ROP into the gate; halt the machine
+		 * instead of continuing with CR0.WP cleared on this CPU.
+		 */
+		panic("iee_dispatch: unknown flag %d\n", flag);
+	}
+	return 0;
+}
diff --git a/arch/x86/kernel/haoc/iee/iee-func.c b/arch/x86/kernel/haoc/iee/iee-func.c
index 0054c1b3d6169..8001ce51dd305 100644
--- a/arch/x86/kernel/haoc/iee/iee-func.c
+++ b/arch/x86/kernel/haoc/iee/iee-func.c
@@ -34,11 +34,12 @@ struct iee_free_slab_work {
 void iee_free_slab(struct kmem_cache *s, struct slab *slab,
 		   void (*do_free_slab)(struct work_struct *work))
 {
-	if(haoc_enabled)
-	return;
 	struct iee_free_slab_work *iee_free_slab_work =
 		kmalloc(sizeof(struct iee_free_slab_work), GFP_ATOMIC);
 
+	if (!iee_free_slab_work)
+		return;
+
 	iee_free_slab_work->s = s;
 	iee_free_slab_work->slab = slab;
 	INIT_WORK(&iee_free_slab_work->work, do_free_slab);
diff --git a/arch/x86/kernel/haoc/iee/iee-gate.S b/arch/x86/kernel/haoc/iee/iee-gate.S
index 5497c3788b026..fc81d79f80503 100644
--- a/arch/x86/kernel/haoc/iee/iee-gate.S
+++ b/arch/x86/kernel/haoc/iee/iee-gate.S
@@ -9,12 +9,11 @@
 
 #include <linux/linkage.h>
 #include <asm/asm-offsets.h>
+#include <asm/ibt.h>
 #include <asm/percpu.h>
 #include <asm/processor-flags.h>
 #include <asm/asm.h>
 
-#define X86_CR4_SMEP_SMAP (X86_CR4_SMEP | X86_CR4_SMAP)
-
 /*
  * scratch_reg would be changed,
  * caller should dertimine if scratch_reg should be saved and restored.
@@ -27,21 +26,34 @@
 .endm
 
 .macro ENABLE_WP scratch_reg:req
-	/* Enable write protection */
+	/*
+	 * Enable write protection. A CR0 write either takes effect or
+	 * faults, so there is nothing to retry; keep this a single write.
+	 * (Pinning enforcement stays in native_write_cr0()/_iee_si_handler.)
+	 */
 	movq %cr0, %\scratch_reg
-1:
 	orq $X86_CR0_WP, %\scratch_reg
 	movq %\scratch_reg, %cr0
-	testq $X86_CR0_WP, %\scratch_reg
-	je 1b
 .endm
 
 /*
  * IEE memory access gate.
  * Kernel calls the gate to modify IEE-protected memory.
+ *
+ * The gates switch to the per-cpu IEE stack in the middle of the
+ * function. objtool cannot model that: %rsp is untraceable after the
+ * per-cpu stack load, so no provable frame pointer state exists at the
+ * call sites (the only modelled stack-switch idiom yields CFI_SP_INDIRECT,
+ * which also fails has_valid_stack_frame()). Declare the gates as code
+ * (not function) symbols to exempt them from stack validation, the same
+ * way x86 entry code handles mid-function stack switches, and maintain a
+ * real %rbp frame chain manually so the frame pointer unwinder still
+ * works at runtime.
  */
 
-SYM_FUNC_START(iee_rw_gate)
+SYM_CODE_START(iee_rw_gate)
+	ENDBR
+
 	/* save Interrupt flag */
 	pushfq
 	/* close irq*/
@@ -51,17 +63,22 @@ SYM_FUNC_START(iee_rw_gate)
 
 	DISABLE_WP r12
 
+	/* set up a real frame for the frame pointer unwinder */
+	pushq %rbp
+	movq %rsp, %rbp
+
 	/* switch to iee stack */
 	movq %rsp, %r12
 	movq PER_CPU_VAR(iee_stacks) + IEE_STACK, %rsp
 
-	/* call iee func */
-	leaq iee_funcs(%rip), %rax
-	call *(%rax, %rdi, 8)
+	/* direct dispatch: no indirect call, no writable func table */
+	call iee_dispatch
 
 	/* switch to kernel stack */
 	movq %r12, %rsp
 
+	popq %rbp
+
 	ENABLE_WP r12
 
 	popq %r12
@@ -70,11 +87,13 @@ SYM_FUNC_START(iee_rw_gate)
 	popfq
 
 	jmp __x86_return_thunk  /* ret */
-SYM_FUNC_END(iee_rw_gate)
+SYM_CODE_END(iee_rw_gate)
 EXPORT_SYMBOL(iee_rw_gate)
 
 #ifdef CONFIG_IEE_SIP
-SYM_FUNC_START(iee_rwx_gate)
+SYM_CODE_START(iee_rwx_gate)
+	ENDBR
+
 	pushq %r12
  
 	/* save Interrupt flag*/
@@ -88,33 +107,43 @@ SYM_FUNC_START(iee_rwx_gate)
 	movq %rax, %cr4
  
 	DISABLE_WP r12
- 
+
+	/* set up a real frame for the frame pointer unwinder */
+	pushq %rbp
+	movq %rsp, %rbp
+
 	movq %rsp, %r12
 	/* If iee hasn't been initialized, skip stack switch. */
 	cmpb $0, iee_init_done(%rip)
 	jz 2f
- 
+
 	/* switch to iee stack */
 	movq PER_CPU_VAR(iee_stacks) + IEE_STACK, %rsp
- 
+
 2:  call _iee_si_handler
+
 	/* switch to kernel stack. If iee hasn't been initialized, skip switch*/
 	movq %r12, %rsp
- 
+
+	popq %rbp
+
 	ENABLE_WP r12
  
-	/* set SMEP=1 to disable supervisor-mode exec user-mode insn */
-	movq %cr4, %rax  /* rax -> cr4 */
-1:  orq $X86_CR4_SMEP_SMAP, %rax
+	/*
+	 * Re-enable SMEP/SMAP after running the handler. Only the bits the
+	 * CPU actually supports may be set (setting an unsupported CR4 bit
+	 * #GPs); iee_cr4_set_mask is initialized in identify_cpu() before
+	 * setup_smep()/setup_smap(). No read-back retry: a CR4 write either
+	 * takes effect or faults, there is nothing to retry.
+	 */
+	movq %cr4, %rax		/* value written by the handler */
+	orq iee_cr4_set_mask(%rip), %rax
 	movq %rax, %cr4
-	andq $X86_CR4_SMEP_SMAP, %rax
-	cmpq $X86_CR4_SMEP_SMAP, %rax
-	jnz 1
  
 	/* restore irq*/
 	popfq
  
 	popq %r12
 	jmp __x86_return_thunk  /* ret */
-SYM_FUNC_END(iee_rwx_gate)
+SYM_CODE_END(iee_rwx_gate)
 #endif
\ No newline at end of file
diff --git a/arch/x86/kernel/haoc/iee/iee-init.c b/arch/x86/kernel/haoc/iee/iee-init.c
index 7ac4704769f4d..9aff3d6b9083c 100644
--- a/arch/x86/kernel/haoc/iee/iee-init.c
+++ b/arch/x86/kernel/haoc/iee/iee-init.c
@@ -17,6 +17,9 @@
 #ifdef CONFIG_IEE_SIP
 #include <asm/haoc/iee-si.h>
 #endif
+#ifdef CONFIG_IEE_PTRP
+#include <asm/haoc/iee-token.h>
+#endif
 
 /* IEE_OFFSET = pgtable_l5_enabled() ? 0x40000000000000 : 0x200000000000; */
 unsigned long IEE_OFFSET = 0x200000000000;
@@ -138,16 +141,18 @@ static void __init _iee_stack_init(void)
 {
 	int cpu;
 	struct iee_stack *iee_stack;
-	void *stack_base;
 	struct page *page;
 
 	for_each_possible_cpu(cpu) {
-		stack_base = (void *)page_address(alloc_pages(GFP_KERNEL, IEE_STACK_ORDER));
 		iee_stack = per_cpu_ptr(&iee_stacks, cpu);
 		page = alloc_pages(GFP_KERNEL, IEE_STACK_ORDER);
-		iee_stack->stack = (void *)page_address(page) + PAGE_SIZE * (1 << IEE_STACK_ORDER);
-		pr_info("IEE: cpu %d, iee_stack 0x%lx", cpu, (unsigned long)iee_stack->stack);
-		set_memory_ro((unsigned long)stack_base, (1 << IEE_STACK_ORDER));
+		if (!page)
+			panic("IEE: failed to allocate iee stack for cpu %d\n",
+			      cpu);
+		iee_stack->stack = (void *)page_address(page) +
+				   PAGE_SIZE * (1 << IEE_STACK_ORDER);
+		pr_info("IEE: cpu %d, iee_stack 0x%lx", cpu,
+			(unsigned long)iee_stack->stack);
 	}
 }
 
@@ -162,24 +167,17 @@ void __init iee_init(void)
 	_iee_offset_init();
 	_iee_mapping_init();
 	_iee_stack_init();
+#ifdef CONFIG_IEE_PTRP
+	iee_prepare_init_task_token();
+#endif
 }
 
 bool __ro_after_init haoc_enabled;
 EXPORT_SYMBOL(haoc_enabled);
-#ifdef CONFIG_IEE_SIP
-extern unsigned long cr4_pinned_mask;
-#endif
 static int __init parse_haoc_enabled(char *str)
 {
 	int ret = kstrtobool(str, &haoc_enabled);
-	#ifdef CONFIG_IEE_SIP
-	if(haoc_enabled)
-	{
-		cr4_pinned_mask =
-		X86_CR4_SMEP | X86_CR4_SMAP | X86_CR4_UMIP |
-		X86_CR4_FSGSBASE | X86_CR4_CET;
-	}
-	#endif
+
 	return ret;
 }
 early_param("haoc", parse_haoc_enabled);
diff --git a/arch/x86/kernel/haoc/iee/iee-si.c b/arch/x86/kernel/haoc/iee/iee-si.c
index 214c7617cfed0..1ed1cf772994b 100644
--- a/arch/x86/kernel/haoc/iee/iee-si.c
+++ b/arch/x86/kernel/haoc/iee/iee-si.c
@@ -1,10 +1,20 @@
 // SPDX-License-Identifier: GPL-2.0
+#include <linux/cache.h>
 #include <linux/init.h>
+
 #include <asm/haoc/iee-si.h>
 #include <asm/haoc/haoc-def.h>
 #include <asm/set_memory.h>
 #include <asm/haoc/iee.h>
 
+/*
+ * SMEP/SMAP bits the boot CPU actually enabled. The rwx gate exit
+ * re-enables these after running the handler; it must never set bits
+ * the CPU does not support (writing them to CR4 would #GP).
+ * Initialized in identify_cpu() before setup_smep()/setup_smap().
+ */
+unsigned long iee_cr4_set_mask __read_mostly;
+
 unsigned long __iee_si_code notrace _iee_si_handler(int flag, ...)
 {
 	va_list pArgs;
@@ -36,9 +46,30 @@ unsigned long __iee_si_code notrace _iee_si_handler(int flag, ...)
 			break;
 		}
 		case IEE_WRITE_CR4: {
+			unsigned long bits_changed = 0;
+			/*
+			 * The handler runs from user-mapped .iee.si_text with SMEP
+			 * already cleared by the gate: writing CR4 with SMEP=1
+			 * here would fault on the next instruction fetch. SMEP is
+			 * restored by the gate exit per iee_cr4_set_mask, so it
+			 * takes no part in the pinning enforcement below.
+			 */
+			const unsigned long check_mask = cr4_pinned_mask & ~X86_CR4_SMEP;
+			const unsigned long check_bits = cr4_pinned_bits & ~X86_CR4_SMEP;
+
 			val = va_arg(pArgs, u64);
-			val &= ~(X86_CR4_SMEP);
+			val &= ~X86_CR4_SMEP;
+			if (static_branch_likely(&cr_pinning)) {
+				if (unlikely((val & check_mask) != check_bits)) {
+					bits_changed = (val & check_mask) ^ check_bits;
+					val = (val & ~check_mask) | check_bits;
+				}
+			}
 			asm volatile("mov %0,%%cr4" : "+r" (val) : : "memory");
+			/* Warn after we've corrected the changed bits. */
+			if (static_branch_likely(&cr_pinning))
+				WARN_ONCE(bits_changed, "pinned CR4 bits changed: 0x%lx!?\n",
+					  bits_changed);
 			break;
 		}
 		case IEE_LOAD_IDT: {
diff --git a/arch/x86/kernel/haoc/iee/iee-token.c b/arch/x86/kernel/haoc/iee/iee-token.c
index b39ba03d9c565..8b665ea836afb 100644
--- a/arch/x86/kernel/haoc/iee/iee-token.c
+++ b/arch/x86/kernel/haoc/iee/iee-token.c
@@ -1,8 +1,41 @@
 // SPDX-License-Identifier: GPL-2.0
+#include <linux/init.h>
+#include <linux/gfp.h>
+#include <linux/mm.h>
 #include <asm/set_memory.h>
 #include <asm/haoc/iee-token.h>
 #include "slab.h"
 
+/*
+ * The token of a slab-allocated task_struct lives on a dedicated page
+ * (see iee_alloc_task_token_slab()).  init_task is a kernel image
+ * symbol, so its IEE alias still points at its own physical page and
+ * token writes (e.g. token->new_cred in CREDP) would land on
+ * init_task itself.  Give it a dedicated token page too.
+ */
+void __init iee_prepare_init_task_token(void)
+{
+	unsigned long token = (unsigned long)__kimg_to_iee(&init_task);
+	unsigned long token_page;
+	unsigned int order = 0;
+
+	/* Allocate one more page if the token crosses a page boundary. */
+	if (ALIGN(token + sizeof(struct task_token), PAGE_SIZE) !=
+	    ALIGN(token + 1, PAGE_SIZE))
+		order = 1;
+
+	token_page = __get_free_pages(GFP_KERNEL | __GFP_ZERO, order);
+	if (!token_page)
+		panic("IEE: failed to allocate token page for init_task\n");
+
+	/* The remap works on whole pages; align the token address down. */
+	iee_set_token_page_valid(ALIGN_DOWN(token, PAGE_SIZE), token_page,
+				 order);
+
+	/* init_task is already running, validate its token directly. */
+	iee_validate_token(&init_task);
+}
+
 void iee_set_token_page_valid(unsigned long token, unsigned long token_page,
 			      unsigned int order)
 {
diff --git a/kernel/cred.c b/kernel/cred.c
index d7042a4633408..df02c0290c1d7 100644
--- a/kernel/cred.c
+++ b/kernel/cred.c
@@ -619,11 +619,11 @@ void abort_creds(struct cred *new)
 	       atomic_long_read(&new->usage));
 
 	BUG_ON(atomic_long_read(&new->usage) < 1);
-	put_cred(new);
 #ifdef CONFIG_CREDP
 	if (haoc_enabled)
 		iee_abort_creds(new);
 #endif
+	put_cred(new);
 }
 EXPORT_SYMBOL(abort_creds);