diff mbox series

[RFC,v3,2/7] mm/asi: ASI entry/exit interface

Message ID 1582734120-26757-3-git-send-email-alexandre.chartre@oracle.com (mailing list archive)
State New, archived
Headers show
Series Kernel Address Space Isolation | expand

Commit Message

Alexandre Chartre Feb. 26, 2020, 4:21 p.m. UTC 
Address Space Isolation (ASI) is entered by calling asi_enter() which
switches the kernel page-table to the ASI page-table. Isolation is then
exited by calling asi_exit() which switches the page-table back to the
original kernel page-table.

The ASI being used and its state is tracked in a per-cpu ASI session
structure (struct asi_session).

Signed-off-by: Alexandre Chartre <alexandre.chartre@oracle.com>
---
 arch/x86/include/asm/asi.h         |    4 ++
 arch/x86/include/asm/asi_session.h |   17 +++++++
 arch/x86/include/asm/mmu_context.h |   19 +++++++-
 arch/x86/include/asm/tlbflush.h    |   12 +++++
 arch/x86/mm/asi.c                  |   90 ++++++++++++++++++++++++++++++++++++
 5 files changed, 140 insertions(+), 2 deletions(-)
 create mode 100644 arch/x86/include/asm/asi_session.h

Comments

Ira Weiny Feb. 26, 2020, 10:02 p.m. UTC | #1 
On Wed, Feb 26, 2020 at 05:21:55PM +0100, Alexandre Chartre wrote:
> Address Space Isolation (ASI) is entered by calling asi_enter() which
> switches the kernel page-table to the ASI page-table. Isolation is then
> exited by calling asi_exit() which switches the page-table back to the
> original kernel page-table.
> 
> The ASI being used and its state is tracked in a per-cpu ASI session
> structure (struct asi_session).
> 
> Signed-off-by: Alexandre Chartre <alexandre.chartre@oracle.com>
> ---
>  arch/x86/include/asm/asi.h         |    4 ++
>  arch/x86/include/asm/asi_session.h |   17 +++++++
>  arch/x86/include/asm/mmu_context.h |   19 +++++++-
>  arch/x86/include/asm/tlbflush.h    |   12 +++++
>  arch/x86/mm/asi.c                  |   90 ++++++++++++++++++++++++++++++++++++
>  5 files changed, 140 insertions(+), 2 deletions(-)
>  create mode 100644 arch/x86/include/asm/asi_session.h
> 
> diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
> index 844a81f..29b745a 100644
> --- a/arch/x86/include/asm/asi.h
> +++ b/arch/x86/include/asm/asi.h
> @@ -44,6 +44,8 @@
>  
>  #include <linux/export.h>
>  
> +#include <asm/asi_session.h>
> +
>  struct asi_type {
>  	int			pcid_prefix;	/* PCID prefix */
>  };
> @@ -80,6 +82,8 @@ struct asi {
>  extern struct asi *asi_create(struct asi_type *type);
>  extern void asi_destroy(struct asi *asi);
>  extern void asi_set_pagetable(struct asi *asi, pgd_t *pagetable);
> +extern int asi_enter(struct asi *asi);
> +extern void asi_exit(struct asi *asi);
>  
>  #endif	/* __ASSEMBLY__ */
>  
> diff --git a/arch/x86/include/asm/asi_session.h b/arch/x86/include/asm/asi_session.h
> new file mode 100644
> index 0000000..9d39c93
> --- /dev/null
> +++ b/arch/x86/include/asm/asi_session.h
> @@ -0,0 +1,17 @@
> +/* SPDX-License-Identifier: GPL-2.0 */
> +#ifndef ARCH_X86_MM_ASI_SESSION_H
> +#define ARCH_X86_MM_ASI_SESSION_H
> +
> +#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
> +
> +struct asi;
> +
> +struct asi_session {
> +	struct asi		*asi;		/* ASI for this session */
> +	unsigned long		isolation_cr3;	/* cr3 when ASI is active */
> +	unsigned long		original_cr3;	/* cr3 before entering ASI */
> +};
> +
> +#endif	/* CONFIG_ADDRESS_SPACE_ISOLATION */
> +
> +#endif
> diff --git a/arch/x86/include/asm/mmu_context.h b/arch/x86/include/asm/mmu_context.h
> index 5f33924..2d65443 100644
> --- a/arch/x86/include/asm/mmu_context.h
> +++ b/arch/x86/include/asm/mmu_context.h
> @@ -14,6 +14,7 @@
>  #include <asm/paravirt.h>
>  #include <asm/mpx.h>
>  #include <asm/debugreg.h>
> +#include <asm/asi.h>
>  
>  extern atomic64_t last_mm_ctx_id;
>  
> @@ -349,8 +350,22 @@ static inline bool arch_vma_access_permitted(struct vm_area_struct *vma,
>   */
>  static inline unsigned long __get_current_cr3_fast(void)
>  {
> -	unsigned long cr3 = build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
> -		this_cpu_read(cpu_tlbstate.loaded_mm_asid));
> +	unsigned long cr3;
> +
> +	/*
> +	 * If isolation is active then we need to return the CR3 for the
> +	 * currently active ASI. This value is stored in the isolation_cr3
> +	 * field of the ASI session.
> +	 */
> +	if (IS_ENABLED(CONFIG_ADDRESS_SPACE_ISOLATION) &&
> +	    this_cpu_read(cpu_asi_session.asi)) {
> +		cr3 = this_cpu_read(cpu_asi_session.isolation_cr3);
> +		/* CR3 read never returns with the NOFLUSH bit */
> +		cr3 &= ~X86_CR3_PCID_NOFLUSH;
> +	} else {
> +		cr3 = build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
> +				this_cpu_read(cpu_tlbstate.loaded_mm_asid));
> +	}
>  
>  	/* For now, be very restrictive about when this can be called. */
>  	VM_WARN_ON(in_nmi() || preemptible());
> diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
> index 6f66d84..241058f 100644
> --- a/arch/x86/include/asm/tlbflush.h
> +++ b/arch/x86/include/asm/tlbflush.h
> @@ -12,6 +12,7 @@
>  #include <asm/invpcid.h>
>  #include <asm/pti.h>
>  #include <asm/processor-flags.h>
> +#include <asm/asi.h>
>  
>  /*
>   * The x86 feature is called PCID (Process Context IDentifier). It is similar
> @@ -239,9 +240,20 @@ struct tlb_state {
>  	 * context 0.
>  	 */
>  	struct tlb_context ctxs[TLB_NR_DYN_ASIDS];
> +
> +#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
> +	/*
> +	 * The ASI session tracks the ASI being used and its state.
> +	 */
> +	struct asi_session asi_session;
> +#endif
>  };
>  DECLARE_PER_CPU_SHARED_ALIGNED(struct tlb_state, cpu_tlbstate);
>  
> +#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
> +#define cpu_asi_session	(cpu_tlbstate.asi_session)
> +#endif
> +
>  /*
>   * Blindly accessing user memory from NMI context can be dangerous
>   * if we're in the middle of switching the current user task or
> diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
> index 0a0ac9d..9fbc921 100644
> --- a/arch/x86/mm/asi.c
> +++ b/arch/x86/mm/asi.c
> @@ -10,6 +10,8 @@
>  
>  #include <asm/asi.h>
>  #include <asm/bug.h>
> +#include <asm/mmu_context.h>
> +#include <asm/tlbflush.h>
>  
>  struct asi *asi_create(struct asi_type *type)
>  {
> @@ -58,3 +60,91 @@ void asi_set_pagetable(struct asi *asi, pgd_t *pagetable)
>  
>  }
>  EXPORT_SYMBOL(asi_set_pagetable);
> +
> +static void asi_switch_to_asi_cr3(struct asi *asi)
> +{
> +	unsigned long original_cr3, asi_cr3;
> +	struct asi_session *asi_session;
> +	u16 pcid;
> +
> +	WARN_ON(!irqs_disabled());
> +
> +	original_cr3 = __get_current_cr3_fast();
> +
> +	/* build the ASI cr3 value */
> +	asi_cr3 = asi->base_cr3;
> +	if (boot_cpu_has(X86_FEATURE_PCID)) {
> +		pcid = original_cr3 & ASI_KERNEL_PCID_MASK;
> +		asi_cr3 |= pcid;
> +	}
> +
> +	/* get the ASI session ready for entering ASI */
> +	asi_session = &get_cpu_var(cpu_asi_session);
> +	asi_session->asi = asi;
> +	asi_session->original_cr3 = original_cr3;
> +	asi_session->isolation_cr3 = asi_cr3;
> +
> +	/* Update CR3 to immediately enter ASI */
> +	native_write_cr3(asi_cr3);
> +}
> +
> +static void asi_switch_to_kernel_cr3(struct asi *asi)

asi is not used in this function?

> +{
> +	struct asi_session *asi_session;
> +	unsigned long original_cr3;
> +
> +	WARN_ON(!irqs_disabled());
> +
> +	original_cr3 = this_cpu_read(cpu_asi_session.original_cr3);
> +	if (boot_cpu_has(X86_FEATURE_PCID))
> +		original_cr3 |= X86_CR3_PCID_NOFLUSH;
> +	native_write_cr3(original_cr3);
> +
> +	asi_session = &get_cpu_var(cpu_asi_session);
> +	asi_session->asi = NULL;
> +}
> +
> +int asi_enter(struct asi *asi)
> +{
> +	struct asi *current_asi;
> +	unsigned long flags;
> +
> +	/*
> +	 * We can re-enter isolation, but only with the same ASI (we don't
> +	 * support nesting isolation).
> +	 */
> +	current_asi = this_cpu_read(cpu_asi_session.asi);
> +	if (current_asi) {
> +		if (current_asi != asi) {
> +			WARN_ON(1);
> +			return -EBUSY;
> +		}

if (WARN_ON(current_asi != asi)) {
	return -EBUSY;

???

Ira

> +		return 0;
> +	}
> +
> +	local_irq_save(flags);
> +	asi_switch_to_asi_cr3(asi);
> +	local_irq_restore(flags);
> +
> +	return 0;
> +}
> +EXPORT_SYMBOL(asi_enter);
> +
> +void asi_exit(struct asi *asi)
> +{
> +	struct asi *current_asi;
> +	unsigned long flags;
> +
> +	current_asi = this_cpu_read(cpu_asi_session.asi);
> +	if (!current_asi) {
> +		/* already exited */
> +		return;
> +	}
> +
> +	WARN_ON(current_asi != asi);
> +
> +	local_irq_save(flags);
> +	asi_switch_to_kernel_cr3(asi);
> +	local_irq_restore(flags);
> +}
> +EXPORT_SYMBOL(asi_exit);
> -- 
> 1.7.1
> 
>
Alexandre Chartre Feb. 27, 2020, 8:51 a.m. UTC | #2 
On 2/26/20 11:02 PM, Ira Weiny wrote:
> On Wed, Feb 26, 2020 at 05:21:55PM +0100, Alexandre Chartre wrote:
>> Address Space Isolation (ASI) is entered by calling asi_enter() which
>> switches the kernel page-table to the ASI page-table. Isolation is then
>> exited by calling asi_exit() which switches the page-table back to the
>> original kernel page-table.
>>
>> The ASI being used and its state is tracked in a per-cpu ASI session
>> structure (struct asi_session).
>>
>> Signed-off-by: Alexandre Chartre <alexandre.chartre@oracle.com>
>> ---
>>   arch/x86/include/asm/asi.h         |    4 ++
>>   arch/x86/include/asm/asi_session.h |   17 +++++++
>>   arch/x86/include/asm/mmu_context.h |   19 +++++++-
>>   arch/x86/include/asm/tlbflush.h    |   12 +++++
>>   arch/x86/mm/asi.c                  |   90 ++++++++++++++++++++++++++++++++++++
>>   5 files changed, 140 insertions(+), 2 deletions(-)
>>   create mode 100644 arch/x86/include/asm/asi_session.h
>>
>> diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
>> index 844a81f..29b745a 100644
>> --- a/arch/x86/include/asm/asi.h
>> +++ b/arch/x86/include/asm/asi.h
>> @@ -44,6 +44,8 @@
>>   
>>   #include <linux/export.h>
>>   
>> +#include <asm/asi_session.h>
>> +
>>   struct asi_type {
>>   	int			pcid_prefix;	/* PCID prefix */
>>   };
>> @@ -80,6 +82,8 @@ struct asi {
>>   extern struct asi *asi_create(struct asi_type *type);
>>   extern void asi_destroy(struct asi *asi);
>>   extern void asi_set_pagetable(struct asi *asi, pgd_t *pagetable);
>> +extern int asi_enter(struct asi *asi);
>> +extern void asi_exit(struct asi *asi);
>>   
>>   #endif	/* __ASSEMBLY__ */
>>   
>> diff --git a/arch/x86/include/asm/asi_session.h b/arch/x86/include/asm/asi_session.h
>> new file mode 100644
>> index 0000000..9d39c93
>> --- /dev/null
>> +++ b/arch/x86/include/asm/asi_session.h
>> @@ -0,0 +1,17 @@
>> +/* SPDX-License-Identifier: GPL-2.0 */
>> +#ifndef ARCH_X86_MM_ASI_SESSION_H
>> +#define ARCH_X86_MM_ASI_SESSION_H
>> +
>> +#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
>> +
>> +struct asi;
>> +
>> +struct asi_session {
>> +	struct asi		*asi;		/* ASI for this session */
>> +	unsigned long		isolation_cr3;	/* cr3 when ASI is active */
>> +	unsigned long		original_cr3;	/* cr3 before entering ASI */
>> +};
>> +
>> +#endif	/* CONFIG_ADDRESS_SPACE_ISOLATION */
>> +
>> +#endif
>> diff --git a/arch/x86/include/asm/mmu_context.h b/arch/x86/include/asm/mmu_context.h
>> index 5f33924..2d65443 100644
>> --- a/arch/x86/include/asm/mmu_context.h
>> +++ b/arch/x86/include/asm/mmu_context.h
>> @@ -14,6 +14,7 @@
>>   #include <asm/paravirt.h>
>>   #include <asm/mpx.h>
>>   #include <asm/debugreg.h>
>> +#include <asm/asi.h>
>>   
>>   extern atomic64_t last_mm_ctx_id;
>>   
>> @@ -349,8 +350,22 @@ static inline bool arch_vma_access_permitted(struct vm_area_struct *vma,
>>    */
>>   static inline unsigned long __get_current_cr3_fast(void)
>>   {
>> -	unsigned long cr3 = build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
>> -		this_cpu_read(cpu_tlbstate.loaded_mm_asid));
>> +	unsigned long cr3;
>> +
>> +	/*
>> +	 * If isolation is active then we need to return the CR3 for the
>> +	 * currently active ASI. This value is stored in the isolation_cr3
>> +	 * field of the ASI session.
>> +	 */
>> +	if (IS_ENABLED(CONFIG_ADDRESS_SPACE_ISOLATION) &&
>> +	    this_cpu_read(cpu_asi_session.asi)) {
>> +		cr3 = this_cpu_read(cpu_asi_session.isolation_cr3);
>> +		/* CR3 read never returns with the NOFLUSH bit */
>> +		cr3 &= ~X86_CR3_PCID_NOFLUSH;
>> +	} else {
>> +		cr3 = build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
>> +				this_cpu_read(cpu_tlbstate.loaded_mm_asid));
>> +	}
>>   
>>   	/* For now, be very restrictive about when this can be called. */
>>   	VM_WARN_ON(in_nmi() || preemptible());
>> diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
>> index 6f66d84..241058f 100644
>> --- a/arch/x86/include/asm/tlbflush.h
>> +++ b/arch/x86/include/asm/tlbflush.h
>> @@ -12,6 +12,7 @@
>>   #include <asm/invpcid.h>
>>   #include <asm/pti.h>
>>   #include <asm/processor-flags.h>
>> +#include <asm/asi.h>
>>   
>>   /*
>>    * The x86 feature is called PCID (Process Context IDentifier). It is similar
>> @@ -239,9 +240,20 @@ struct tlb_state {
>>   	 * context 0.
>>   	 */
>>   	struct tlb_context ctxs[TLB_NR_DYN_ASIDS];
>> +
>> +#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
>> +	/*
>> +	 * The ASI session tracks the ASI being used and its state.
>> +	 */
>> +	struct asi_session asi_session;
>> +#endif
>>   };
>>   DECLARE_PER_CPU_SHARED_ALIGNED(struct tlb_state, cpu_tlbstate);
>>   
>> +#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
>> +#define cpu_asi_session	(cpu_tlbstate.asi_session)
>> +#endif
>> +
>>   /*
>>    * Blindly accessing user memory from NMI context can be dangerous
>>    * if we're in the middle of switching the current user task or
>> diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
>> index 0a0ac9d..9fbc921 100644
>> --- a/arch/x86/mm/asi.c
>> +++ b/arch/x86/mm/asi.c
>> @@ -10,6 +10,8 @@
>>   
>>   #include <asm/asi.h>
>>   #include <asm/bug.h>
>> +#include <asm/mmu_context.h>
>> +#include <asm/tlbflush.h>
>>   
>>   struct asi *asi_create(struct asi_type *type)
>>   {
>> @@ -58,3 +60,91 @@ void asi_set_pagetable(struct asi *asi, pgd_t *pagetable)
>>   
>>   }
>>   EXPORT_SYMBOL(asi_set_pagetable);
>> +
>> +static void asi_switch_to_asi_cr3(struct asi *asi)
>> +{
>> +	unsigned long original_cr3, asi_cr3;
>> +	struct asi_session *asi_session;
>> +	u16 pcid;
>> +
>> +	WARN_ON(!irqs_disabled());
>> +
>> +	original_cr3 = __get_current_cr3_fast();
>> +
>> +	/* build the ASI cr3 value */
>> +	asi_cr3 = asi->base_cr3;
>> +	if (boot_cpu_has(X86_FEATURE_PCID)) {
>> +		pcid = original_cr3 & ASI_KERNEL_PCID_MASK;
>> +		asi_cr3 |= pcid;
>> +	}
>> +
>> +	/* get the ASI session ready for entering ASI */
>> +	asi_session = &get_cpu_var(cpu_asi_session);
>> +	asi_session->asi = asi;
>> +	asi_session->original_cr3 = original_cr3;
>> +	asi_session->isolation_cr3 = asi_cr3;
>> +
>> +	/* Update CR3 to immediately enter ASI */
>> +	native_write_cr3(asi_cr3);
>> +}
>> +
>> +static void asi_switch_to_kernel_cr3(struct asi *asi)
> 
> asi is not used in this function?
> 

You are right, I was probably using it in a previous version of the code
and forgot to remote it. I will remove it.
  

>> +{
>> +	struct asi_session *asi_session;
>> +	unsigned long original_cr3;
>> +
>> +	WARN_ON(!irqs_disabled());
>> +
>> +	original_cr3 = this_cpu_read(cpu_asi_session.original_cr3);
>> +	if (boot_cpu_has(X86_FEATURE_PCID))
>> +		original_cr3 |= X86_CR3_PCID_NOFLUSH;
>> +	native_write_cr3(original_cr3);
>> +
>> +	asi_session = &get_cpu_var(cpu_asi_session);
>> +	asi_session->asi = NULL;
>> +}
>> +
>> +int asi_enter(struct asi *asi)
>> +{
>> +	struct asi *current_asi;
>> +	unsigned long flags;
>> +
>> +	/*
>> +	 * We can re-enter isolation, but only with the same ASI (we don't
>> +	 * support nesting isolation).
>> +	 */
>> +	current_asi = this_cpu_read(cpu_asi_session.asi);
>> +	if (current_asi) {
>> +		if (current_asi != asi) {
>> +			WARN_ON(1);
>> +			return -EBUSY;
>> +		}
> 
> if (WARN_ON(current_asi != asi)) {
> 	return -EBUSY;
> 

Yes, much better, will change.

Thanks,

alex.

> 
>> +		return 0;
>> +	}
>> +
>> +	local_irq_save(flags);
>> +	asi_switch_to_asi_cr3(asi);
>> +	local_irq_restore(flags);
>> +
>> +	return 0;
>> +}
>> +EXPORT_SYMBOL(asi_enter);
>> +
>> +void asi_exit(struct asi *asi)
>> +{
>> +	struct asi *current_asi;
>> +	unsigned long flags;
>> +
>> +	current_asi = this_cpu_read(cpu_asi_session.asi);
>> +	if (!current_asi) {
>> +		/* already exited */
>> +		return;
>> +	}
>> +
>> +	WARN_ON(current_asi != asi);
>> +
>> +	local_irq_save(flags);
>> +	asi_switch_to_kernel_cr3(asi);
>> +	local_irq_restore(flags);
>> +}
>> +EXPORT_SYMBOL(asi_exit);
>> -- 
>> 1.7.1
>>
>>
>
diff mbox series

Patch

diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
index 844a81f..29b745a 100644
--- a/arch/x86/include/asm/asi.h
+++ b/arch/x86/include/asm/asi.h
@@ -44,6 +44,8 @@ 
 
 #include <linux/export.h>
 
+#include <asm/asi_session.h>
+
 struct asi_type {
 	int			pcid_prefix;	/* PCID prefix */
 };
@@ -80,6 +82,8 @@  struct asi {
 extern struct asi *asi_create(struct asi_type *type);
 extern void asi_destroy(struct asi *asi);
 extern void asi_set_pagetable(struct asi *asi, pgd_t *pagetable);
+extern int asi_enter(struct asi *asi);
+extern void asi_exit(struct asi *asi);
 
 #endif	/* __ASSEMBLY__ */
 
diff --git a/arch/x86/include/asm/asi_session.h b/arch/x86/include/asm/asi_session.h
new file mode 100644
index 0000000..9d39c93
--- /dev/null
+++ b/arch/x86/include/asm/asi_session.h
@@ -0,0 +1,17 @@ 
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef ARCH_X86_MM_ASI_SESSION_H
+#define ARCH_X86_MM_ASI_SESSION_H
+
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+
+struct asi;
+
+struct asi_session {
+	struct asi		*asi;		/* ASI for this session */
+	unsigned long		isolation_cr3;	/* cr3 when ASI is active */
+	unsigned long		original_cr3;	/* cr3 before entering ASI */
+};
+
+#endif	/* CONFIG_ADDRESS_SPACE_ISOLATION */
+
+#endif
diff --git a/arch/x86/include/asm/mmu_context.h b/arch/x86/include/asm/mmu_context.h
index 5f33924..2d65443 100644
--- a/arch/x86/include/asm/mmu_context.h
+++ b/arch/x86/include/asm/mmu_context.h
@@ -14,6 +14,7 @@ 
 #include <asm/paravirt.h>
 #include <asm/mpx.h>
 #include <asm/debugreg.h>
+#include <asm/asi.h>
 
 extern atomic64_t last_mm_ctx_id;
 
@@ -349,8 +350,22 @@  static inline bool arch_vma_access_permitted(struct vm_area_struct *vma,
  */
 static inline unsigned long __get_current_cr3_fast(void)
 {
-	unsigned long cr3 = build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
-		this_cpu_read(cpu_tlbstate.loaded_mm_asid));
+	unsigned long cr3;
+
+	/*
+	 * If isolation is active then we need to return the CR3 for the
+	 * currently active ASI. This value is stored in the isolation_cr3
+	 * field of the ASI session.
+	 */
+	if (IS_ENABLED(CONFIG_ADDRESS_SPACE_ISOLATION) &&
+	    this_cpu_read(cpu_asi_session.asi)) {
+		cr3 = this_cpu_read(cpu_asi_session.isolation_cr3);
+		/* CR3 read never returns with the NOFLUSH bit */
+		cr3 &= ~X86_CR3_PCID_NOFLUSH;
+	} else {
+		cr3 = build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
+				this_cpu_read(cpu_tlbstate.loaded_mm_asid));
+	}
 
 	/* For now, be very restrictive about when this can be called. */
 	VM_WARN_ON(in_nmi() || preemptible());
diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
index 6f66d84..241058f 100644
--- a/arch/x86/include/asm/tlbflush.h
+++ b/arch/x86/include/asm/tlbflush.h
@@ -12,6 +12,7 @@ 
 #include <asm/invpcid.h>
 #include <asm/pti.h>
 #include <asm/processor-flags.h>
+#include <asm/asi.h>
 
 /*
  * The x86 feature is called PCID (Process Context IDentifier). It is similar
@@ -239,9 +240,20 @@  struct tlb_state {
 	 * context 0.
 	 */
 	struct tlb_context ctxs[TLB_NR_DYN_ASIDS];
+
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+	/*
+	 * The ASI session tracks the ASI being used and its state.
+	 */
+	struct asi_session asi_session;
+#endif
 };
 DECLARE_PER_CPU_SHARED_ALIGNED(struct tlb_state, cpu_tlbstate);
 
+#ifdef CONFIG_ADDRESS_SPACE_ISOLATION
+#define cpu_asi_session	(cpu_tlbstate.asi_session)
+#endif
+
 /*
  * Blindly accessing user memory from NMI context can be dangerous
  * if we're in the middle of switching the current user task or
diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
index 0a0ac9d..9fbc921 100644
--- a/arch/x86/mm/asi.c
+++ b/arch/x86/mm/asi.c
@@ -10,6 +10,8 @@ 
 
 #include <asm/asi.h>
 #include <asm/bug.h>
+#include <asm/mmu_context.h>
+#include <asm/tlbflush.h>
 
 struct asi *asi_create(struct asi_type *type)
 {
@@ -58,3 +60,91 @@  void asi_set_pagetable(struct asi *asi, pgd_t *pagetable)
 
 }
 EXPORT_SYMBOL(asi_set_pagetable);
+
+static void asi_switch_to_asi_cr3(struct asi *asi)
+{
+	unsigned long original_cr3, asi_cr3;
+	struct asi_session *asi_session;
+	u16 pcid;
+
+	WARN_ON(!irqs_disabled());
+
+	original_cr3 = __get_current_cr3_fast();
+
+	/* build the ASI cr3 value */
+	asi_cr3 = asi->base_cr3;
+	if (boot_cpu_has(X86_FEATURE_PCID)) {
+		pcid = original_cr3 & ASI_KERNEL_PCID_MASK;
+		asi_cr3 |= pcid;
+	}
+
+	/* get the ASI session ready for entering ASI */
+	asi_session = &get_cpu_var(cpu_asi_session);
+	asi_session->asi = asi;
+	asi_session->original_cr3 = original_cr3;
+	asi_session->isolation_cr3 = asi_cr3;
+
+	/* Update CR3 to immediately enter ASI */
+	native_write_cr3(asi_cr3);
+}
+
+static void asi_switch_to_kernel_cr3(struct asi *asi)
+{
+	struct asi_session *asi_session;
+	unsigned long original_cr3;
+
+	WARN_ON(!irqs_disabled());
+
+	original_cr3 = this_cpu_read(cpu_asi_session.original_cr3);
+	if (boot_cpu_has(X86_FEATURE_PCID))
+		original_cr3 |= X86_CR3_PCID_NOFLUSH;
+	native_write_cr3(original_cr3);
+
+	asi_session = &get_cpu_var(cpu_asi_session);
+	asi_session->asi = NULL;
+}
+
+int asi_enter(struct asi *asi)
+{
+	struct asi *current_asi;
+	unsigned long flags;
+
+	/*
+	 * We can re-enter isolation, but only with the same ASI (we don't
+	 * support nesting isolation).
+	 */
+	current_asi = this_cpu_read(cpu_asi_session.asi);
+	if (current_asi) {
+		if (current_asi != asi) {
+			WARN_ON(1);
+			return -EBUSY;
+		}
+		return 0;
+	}
+
+	local_irq_save(flags);
+	asi_switch_to_asi_cr3(asi);
+	local_irq_restore(flags);
+
+	return 0;
+}
+EXPORT_SYMBOL(asi_enter);
+
+void asi_exit(struct asi *asi)
+{
+	struct asi *current_asi;
+	unsigned long flags;
+
+	current_asi = this_cpu_read(cpu_asi_session.asi);
+	if (!current_asi) {
+		/* already exited */
+		return;
+	}
+
+	WARN_ON(current_asi != asi);
+
+	local_irq_save(flags);
+	asi_switch_to_kernel_cr3(asi);
+	local_irq_restore(flags);
+}
+EXPORT_SYMBOL(asi_exit);