diff mbox series

[RFC,v2,04/29] mm: asi: Add infrastructure for boot-time enablement

Message ID 20250110-asi-rfc-v2-v2-4-8419288bc805@google.com
State New
Headers show
Series Address Space Isolation (ASI) | expand

Commit Message

Brendan Jackman Jan. 10, 2025, 6:40 p.m. UTC
Add a boot time parameter to control the newly added X86_FEATURE_ASI.
"asi=on" or "asi=off" can be used in the kernel command line to enable
or disable ASI at boot time. If not specified, ASI enablement depends
on CONFIG_ADDRESS_SPACE_ISOLATION_DEFAULT_ON, which is off by default.
asi_check_boottime_disable() is modeled after
pti_check_boottime_disable().

The boot parameter is currently ignored until ASI is fully functional.

Once we have a set of ASI features checked in that we have actually
tested, we will stop ignoring the flag. But for now let's just add the
infrastructure so we can implement the usage code.

Ignoring checkpatch.pl CONFIG_DESCRIPTION because the _DEFAULT_ON
Kconfig is trivial to explain.

Checkpatch-args: --ignore CONFIG_DESCRIPTION
Co-developed-by: Junaid Shahid <junaids@google.com>
Signed-off-by: Junaid Shahid <junaids@google.com>
Co-developed-by: Yosry Ahmed <yosryahmed@google.com>
Signed-off-by: Yosry Ahmed <yosryahmed@google.com>
Signed-off-by: Brendan Jackman <jackmanb@google.com>
---
 arch/x86/Kconfig                         |  9 +++++
 arch/x86/include/asm/asi.h               | 19 ++++++++--
 arch/x86/include/asm/cpufeatures.h       |  1 +
 arch/x86/include/asm/disabled-features.h |  8 ++++-
 arch/x86/mm/asi.c                        | 61 +++++++++++++++++++++++++++-----
 arch/x86/mm/init.c                       |  4 ++-
 include/asm-generic/asi.h                |  4 +++
 7 files changed, 92 insertions(+), 14 deletions(-)

Comments

Borislav Petkov March 19, 2025, 5:29 p.m. UTC | #1
On Fri, Jan 10, 2025 at 06:40:30PM +0000, Brendan Jackman wrote:
> Add a boot time parameter to control the newly added X86_FEATURE_ASI.
> "asi=on" or "asi=off" can be used in the kernel command line to enable
> or disable ASI at boot time. If not specified, ASI enablement depends
> on CONFIG_ADDRESS_SPACE_ISOLATION_DEFAULT_ON, which is off by default.

I don't know yet why we need this default-on thing...

> asi_check_boottime_disable() is modeled after
> pti_check_boottime_disable().
> 
> The boot parameter is currently ignored until ASI is fully functional.
> 
> Once we have a set of ASI features checked in that we have actually
> tested, we will stop ignoring the flag. But for now let's just add the
> infrastructure so we can implement the usage code.
> 
> Ignoring checkpatch.pl CONFIG_DESCRIPTION because the _DEFAULT_ON
> Kconfig is trivial to explain.

Those last two paragraphs go...

> Checkpatch-args: --ignore CONFIG_DESCRIPTION
> Co-developed-by: Junaid Shahid <junaids@google.com>
> Signed-off-by: Junaid Shahid <junaids@google.com>
> Co-developed-by: Yosry Ahmed <yosryahmed@google.com>
> Signed-off-by: Yosry Ahmed <yosryahmed@google.com>
> Signed-off-by: Brendan Jackman <jackmanb@google.com>
> ---

... here as that's text not really pertaining to the contents of the patch.

>  arch/x86/Kconfig                         |  9 +++++
>  arch/x86/include/asm/asi.h               | 19 ++++++++--
>  arch/x86/include/asm/cpufeatures.h       |  1 +
>  arch/x86/include/asm/disabled-features.h |  8 ++++-
>  arch/x86/mm/asi.c                        | 61 +++++++++++++++++++++++++++-----
>  arch/x86/mm/init.c                       |  4 ++-
>  include/asm-generic/asi.h                |  4 +++
>  7 files changed, 92 insertions(+), 14 deletions(-)

...

>   * the N ASI classes.
>   */
>  
> +#define static_asi_enabled() cpu_feature_enabled(X86_FEATURE_ASI)

Yeah, as already mentioned somewhere else, whack that thing pls.

> +
>  /*
>   * ASI uses a per-CPU tainting model to track what mitigation actions are
>   * required on domain transitions. Taints exist along two dimensions:
> @@ -131,6 +134,8 @@ struct asi {
>  
>  DECLARE_PER_CPU_ALIGNED(struct asi *, curr_asi);
>  
> +void asi_check_boottime_disable(void);
> +
>  void asi_init_mm_state(struct mm_struct *mm);
>  
>  int asi_init_class(enum asi_class_id class_id, struct asi_taint_policy *taint_policy);
> @@ -155,7 +160,9 @@ void asi_exit(void);
>  /* The target is the domain we'll enter when returning to process context. */
>  static __always_inline struct asi *asi_get_target(struct task_struct *p)
>  {
> -	return p->thread.asi_state.target;
> +	return static_asi_enabled()
> +	       ? p->thread.asi_state.target
> +	       : NULL;

Waaay too fancy for old people:

	if ()
		return...
	else
		return NULL;

:-)

The others too pls.

>  static __always_inline void asi_set_target(struct task_struct *p,
> @@ -166,7 +173,9 @@ static __always_inline void asi_set_target(struct task_struct *p,
>  
>  static __always_inline struct asi *asi_get_current(void)
>  {
> -	return this_cpu_read(curr_asi);
> +	return static_asi_enabled()
> +	       ? this_cpu_read(curr_asi)
> +	       : NULL;
>  }
>  
>  /* Are we currently in a restricted address space? */
> @@ -175,7 +184,11 @@ static __always_inline bool asi_is_restricted(void)
>  	return (bool)asi_get_current();
>  }
>  
> -/* If we exit/have exited, can we stay that way until the next asi_enter? */
> +/*
> + * If we exit/have exited, can we stay that way until the next asi_enter?

What is that supposed to mean here?

> + *
> + * When ASI is disabled, this returns true.
> + */
>  static __always_inline bool asi_is_relaxed(void)
>  {
>  	return !asi_get_target(current);
> diff --git a/arch/x86/include/asm/cpufeatures.h b/arch/x86/include/asm/cpufeatures.h
> index 913fd3a7bac6506141de65f33b9ee61c615c7d7d..d6a808d10c3b8900d190ea01c66fc248863f05e2 100644
> --- a/arch/x86/include/asm/cpufeatures.h
> +++ b/arch/x86/include/asm/cpufeatures.h
> @@ -474,6 +474,7 @@
>  #define X86_FEATURE_CLEAR_BHB_HW	(21*32+ 3) /* BHI_DIS_S HW control enabled */
>  #define X86_FEATURE_CLEAR_BHB_LOOP_ON_VMEXIT (21*32+ 4) /* Clear branch history at vmexit using SW loop */
>  #define X86_FEATURE_FAST_CPPC		(21*32 + 5) /* AMD Fast CPPC */
> +#define X86_FEATURE_ASI			(21*32+6) /* Kernel Address Space Isolation */
>  
>  /*
>   * BUG word(s)
> diff --git a/arch/x86/include/asm/disabled-features.h b/arch/x86/include/asm/disabled-features.h
> index c492bdc97b0595ec77f89dc9b0cefe5e3e64be41..c7964ed4fef8b9441e1c0453da587787d8008d9d 100644
> --- a/arch/x86/include/asm/disabled-features.h
> +++ b/arch/x86/include/asm/disabled-features.h
> @@ -50,6 +50,12 @@
>  # define DISABLE_PTI		(1 << (X86_FEATURE_PTI & 31))
>  #endif
>  
> +#ifdef CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION
> +# define DISABLE_ASI		0
> +#else
> +# define DISABLE_ASI		(1 << (X86_FEATURE_ASI & 31))
> +#endif
> +
>  #ifdef CONFIG_MITIGATION_RETPOLINE
>  # define DISABLE_RETPOLINE	0
>  #else
> @@ -154,7 +160,7 @@
>  #define DISABLED_MASK17	0
>  #define DISABLED_MASK18	(DISABLE_IBT)
>  #define DISABLED_MASK19	(DISABLE_SEV_SNP)
> -#define DISABLED_MASK20	0
> +#define DISABLED_MASK20	(DISABLE_ASI)
>  #define DISABLED_MASK21	0
>  #define DISABLED_MASK_CHECK BUILD_BUG_ON_ZERO(NCAPINTS != 22)
>  

Right, that hunk is done this way now:

diff --git a/arch/x86/Kconfig.cpufeatures b/arch/x86/Kconfig.cpufeatures
index e12d5b7e39a2..f219eaf664fb 100644
--- a/arch/x86/Kconfig.cpufeatures
+++ b/arch/x86/Kconfig.cpufeatures
@@ -199,3 +199,7 @@ config X86_DISABLED_FEATURE_SEV_SNP
 config X86_DISABLED_FEATURE_INVLPGB
 	def_bool y
 	depends on !BROADCAST_TLB_FLUSH
+
+config X86_DISABLED_FEATURE_ASI
+	def_bool y
+	depends on !MITIGATION_ADDRESS_SPACE_ISOLATION


> diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
> index 105cd8b43eaf5c20acc80d4916b761559fb95d74..5baf563a078f5b3a6cd4b9f5e92baaf81b0774c4 100644
> --- a/arch/x86/mm/asi.c
> +++ b/arch/x86/mm/asi.c
> @@ -4,6 +4,7 @@
>  #include <linux/percpu.h>
>  #include <linux/spinlock.h>
>  
> +#include <linux/init.h>
>  #include <asm/asi.h>
>  #include <asm/cmdline.h>
>  #include <asm/cpufeature.h>
> @@ -29,6 +30,9 @@ static inline bool asi_class_id_valid(enum asi_class_id class_id)
>  
>  static inline bool asi_class_initialized(enum asi_class_id class_id)
>  {
> +	if (!boot_cpu_has(X86_FEATURE_ASI))

check_for_deprecated_apis: WARNING: arch/x86/mm/asi.c:33: Do not use boot_cpu_has() - use cpu_feature_enabled() instead.

Check your whole set pls.

> +		return 0;
> +
>  	if (WARN_ON(!asi_class_id_valid(class_id)))
>  		return false;
>  
> @@ -51,6 +55,9 @@ EXPORT_SYMBOL_GPL(asi_init_class);
>  
>  void asi_uninit_class(enum asi_class_id class_id)
>  {
> +	if (!boot_cpu_has(X86_FEATURE_ASI))
> +		return;
> +
>  	if (!asi_class_initialized(class_id))
>  		return;
>  
> @@ -66,10 +73,36 @@ const char *asi_class_name(enum asi_class_id class_id)
>  	return asi_class_names[class_id];
>  }
>  
> +void __init asi_check_boottime_disable(void)
> +{
> +	bool enabled = IS_ENABLED(CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION_DEFAULT_ON);
> +	char arg[4];
> +	int ret;
> +
> +	ret = cmdline_find_option(boot_command_line, "asi", arg, sizeof(arg));
> +	if (ret == 3 && !strncmp(arg, "off", 3)) {
> +		enabled = false;
> +		pr_info("ASI disabled through kernel command line.\n");
> +	} else if (ret == 2 && !strncmp(arg, "on", 2)) {
> +		enabled = true;
> +		pr_info("Ignoring asi=on param while ASI implementation is incomplete.\n");
> +	} else {
> +		pr_info("ASI %s by default.\n",
> +			enabled ? "enabled" : "disabled");
> +	}
> +
> +	if (enabled)
> +		pr_info("ASI enablement ignored due to incomplete implementation.\n");

Incomplete how?

> +}
> +
>  static void __asi_destroy(struct asi *asi)
>  {
> -	lockdep_assert_held(&asi->mm->asi_init_lock);
> +	WARN_ON_ONCE(asi->ref_count <= 0);
> +	if (--(asi->ref_count) > 0)

Switch that to

include/linux/kref.h

It gives you a sanity-checking functionality too so you don't need the WARN...

> +		return;
>  
> +	free_pages((ulong)asi->pgd, PGD_ALLOCATION_ORDER);
> +	memset(asi, 0, sizeof(struct asi));

And then you can do:

	if (kref_put())
		free_pages...

and so on.

Thx.
Yosry Ahmed March 19, 2025, 6:47 p.m. UTC | #2
On Wed, Mar 19, 2025 at 06:29:35PM +0100, Borislav Petkov wrote:
> On Fri, Jan 10, 2025 at 06:40:30PM +0000, Brendan Jackman wrote:
> > Add a boot time parameter to control the newly added X86_FEATURE_ASI.
> > "asi=on" or "asi=off" can be used in the kernel command line to enable
> > or disable ASI at boot time. If not specified, ASI enablement depends
> > on CONFIG_ADDRESS_SPACE_ISOLATION_DEFAULT_ON, which is off by default.
> 
> I don't know yet why we need this default-on thing...

It's a convenience to avoid needing to set asi=on if you want ASI to be
on by default. It's similar to HUGETLB_PAGE_OPTIMIZE_VMEMMAP_DEFAULT_ON
or ZSWAP_DEFAULT_ON.

[..]
> > @@ -175,7 +184,11 @@ static __always_inline bool asi_is_restricted(void)
> >  	return (bool)asi_get_current();
> >  }
> >  
> > -/* If we exit/have exited, can we stay that way until the next asi_enter? */
> > +/*
> > + * If we exit/have exited, can we stay that way until the next asi_enter?
> 
> What is that supposed to mean here?

asi_is_relaxed() checks if the thread is outside an ASI critical
section.

I say "the thread" because it will also return true if we are executing
an interrupt that arrived during the critical section, even though the
interrupt handler is not technically part of the critical section.

Now the reason it says "if we exit we stay that way" is probably
referring to the fact that an asi_exit() when interrupting a critical
section will be undone in the interrupt epilogue by re-entering ASI.

I agree the wording here is confusing. We should probably describe this
more explicitly and probably rename the function after the API
discussions you had in the previous patch.

> 
> > + *
> > + * When ASI is disabled, this returns true.
> > + */
> >  static __always_inline bool asi_is_relaxed(void)
> >  {
> >  	return !asi_get_target(current);
[..]
> > @@ -66,10 +73,36 @@ const char *asi_class_name(enum asi_class_id class_id)
> >  	return asi_class_names[class_id];
> >  }
> >  
> > +void __init asi_check_boottime_disable(void)
> > +{
> > +	bool enabled = IS_ENABLED(CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION_DEFAULT_ON);
> > +	char arg[4];
> > +	int ret;
> > +
> > +	ret = cmdline_find_option(boot_command_line, "asi", arg, sizeof(arg));
> > +	if (ret == 3 && !strncmp(arg, "off", 3)) {
> > +		enabled = false;
> > +		pr_info("ASI disabled through kernel command line.\n");
> > +	} else if (ret == 2 && !strncmp(arg, "on", 2)) {
> > +		enabled = true;
> > +		pr_info("Ignoring asi=on param while ASI implementation is incomplete.\n");
> > +	} else {
> > +		pr_info("ASI %s by default.\n",
> > +			enabled ? "enabled" : "disabled");
> > +	}
> > +
> > +	if (enabled)
> > +		pr_info("ASI enablement ignored due to incomplete implementation.\n");
> 
> Incomplete how?

This is referring to the fact that ASI is still not fully/correctly
functional, but it will be after the following patches.

I think it will be clearer if we just add the feature flag here so that
we have something to check for in the following patches, but add the
infrastructure for boot-time enablement at the end of the series when
the impelemntation is complete.

Basically start by a feature flag that has no way of being enabled, use
it in the implmentation, then add means of enabling it.

> 
> > +}
> > +
> >  static void __asi_destroy(struct asi *asi)
> >  {
> > -	lockdep_assert_held(&asi->mm->asi_init_lock);
> > +	WARN_ON_ONCE(asi->ref_count <= 0);
> > +	if (--(asi->ref_count) > 0)
> 
> Switch that to
> 
> include/linux/kref.h
> 
> It gives you a sanity-checking functionality too so you don't need the WARN...

I think we hve internal changes that completely get rid of this
ref_count and simplifies the lifetime handling that we can squash here.
We basically keep ASI objects around until the process is torn down,
which makes this simpler and avoids the need for complex synchronization
when we try to context switch or run userspace without exiting ASI
(spoiler alert :) ).

> 
> > +		return;
> >  
> > +	free_pages((ulong)asi->pgd, PGD_ALLOCATION_ORDER);
> > +	memset(asi, 0, sizeof(struct asi));
> 
> And then you can do:
> 
> 	if (kref_put())
> 		free_pages...
> 
> and so on.
> 
> Thx.
> 
> -- 
> Regards/Gruss,
>     Boris.
> 
> https://people.kernel.org/tglx/notes-about-netiquette
>
Brendan Jackman March 20, 2025, 10:44 a.m. UTC | #3
On Wed Mar 19, 2025 at 6:47 PM UTC, Yosry Ahmed wrote:
> On Wed, Mar 19, 2025 at 06:29:35PM +0100, Borislav Petkov wrote:
> > On Fri, Jan 10, 2025 at 06:40:30PM +0000, Brendan Jackman wrote:
> > > Add a boot time parameter to control the newly added X86_FEATURE_ASI.
> > > "asi=on" or "asi=off" can be used in the kernel command line to enable
> > > or disable ASI at boot time. If not specified, ASI enablement depends
> > > on CONFIG_ADDRESS_SPACE_ISOLATION_DEFAULT_ON, which is off by default.
> > 
> > I don't know yet why we need this default-on thing...
>
> It's a convenience to avoid needing to set asi=on if you want ASI to be
> on by default. It's similar to HUGETLB_PAGE_OPTIMIZE_VMEMMAP_DEFAULT_ON
> or ZSWAP_DEFAULT_ON.
>
> [..]
> > > @@ -175,7 +184,11 @@ static __always_inline bool asi_is_restricted(void)
> > >  	return (bool)asi_get_current();
> > >  }
> > >  
> > > -/* If we exit/have exited, can we stay that way until the next asi_enter? */
> > > +/*
> > > + * If we exit/have exited, can we stay that way until the next asi_enter?
> > 
> > What is that supposed to mean here?
>
> asi_is_relaxed() checks if the thread is outside an ASI critical
> section.
>
> I say "the thread" because it will also return true if we are executing
> an interrupt that arrived during the critical section, even though the
> interrupt handler is not technically part of the critical section.
>
> Now the reason it says "if we exit we stay that way" is probably
> referring to the fact that an asi_exit() when interrupting a critical
> section will be undone in the interrupt epilogue by re-entering ASI.
>
> I agree the wording here is confusing. We should probably describe this
> more explicitly and probably rename the function after the API
> discussions you had in the previous patch.

Yeah, this is confusing. It's trying to very concisely define the
concept of "relaxed" but now I see it through Boris' eyes I realise
it's really unhelpful to try and do that. And yeah we should probably
just rework the terminology/API.

To re-iterate what Yosry said, aside from my too-clever comment style
the more fundamental thing that's confusing here is that, using the
terminology currently in the code there are two concepts at play:

- The critical section: this is the path from asi_enter() to
  asi_relax(). The critical section can be interrupted, and code
  running in those interupts is not said to be "in the critical
  section".

- Being "tense" vs "relaxed". Being "tense" means the _task_ is in a
  critical section, but the current code might not be.

This distinction is theoretically relevant because e.g. it's a bug to
access sensitive data in a critical section, but it's OK to access it
while in the tense state (we will switch to the restricted address
space, but this is OK because we will have a chance to asi_enter()
again before we get back to the untrusted code). 

BTW, just to be clear:

1. Both of these are only relevant to code that's pretty deeply aware
   of ASI. (TLB flushing code, entry code, stuff like that).

2. To be honest whenever you write:

     if (asi_in_critical_section())

   You probably mean:

     if (WARN_ON(asi_in_critical_section()))

   For example if we try to flush the TLB in the critical section,
   there's a thing we can do to handle it. But that really shouldn't
   be necessary.  We want the critical section code to be very small
   and straight-line code.

   And indeed in the present code we don't use
   asi_in_critical_section() for anything bur WARNing.

> asi_is_relaxed() checks if the thread is outside an ASI critical
> section.

Now I see it written this way, this is probably the best way to
conceptualise it. Instead of having two concepts "tense/relaxed" vs
"ASI critical section" we could just say "the task is in a critical
section" vs "the CPU is in a critical section". So we could have
something like:

bool asi_task_critical(void);
bool asi_cpu_critical(void);

(They could also accept an argument for the task/CPU, but I can't see
any reason why you'd peek at another context like that).

--

For everything else, Ack to Boris or +1 to Yosry respectively.
diff mbox series

Patch

diff --git a/arch/x86/Kconfig b/arch/x86/Kconfig
index 5a50582eb210e9d1309856a737d32b76fa1bfc85..1fcb52cb8cd5084ac3cef04af61b7d1653162bdb 100644
--- a/arch/x86/Kconfig
+++ b/arch/x86/Kconfig
@@ -2533,6 +2533,15 @@  config MITIGATION_ADDRESS_SPACE_ISOLATION
 	  there are likely to be unhandled cases, in particular concerning TLB
 	  flushes.
 
+
+config ADDRESS_SPACE_ISOLATION_DEFAULT_ON
+	bool "Enable address space isolation by default"
+	default n
+	depends on MITIGATION_ADDRESS_SPACE_ISOLATION
+	help
+	  If selected, ASI is enabled by default at boot if the asi=on or
+	  asi=off are not specified.
+
 config MITIGATION_RETPOLINE
 	bool "Avoid speculative indirect branches in kernel"
 	select OBJTOOL if HAVE_OBJTOOL
diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
index 7cc635b6653a3970ba9dbfdc9c828a470e27bd44..b9671ef2dd3278adceed18507fd260e21954d574 100644
--- a/arch/x86/include/asm/asi.h
+++ b/arch/x86/include/asm/asi.h
@@ -8,6 +8,7 @@ 
 
 #include <asm/pgtable_types.h>
 #include <asm/percpu.h>
+#include <asm/cpufeature.h>
 #include <asm/processor.h>
 
 #ifdef CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION
@@ -66,6 +67,8 @@ 
  * the N ASI classes.
  */
 
+#define static_asi_enabled() cpu_feature_enabled(X86_FEATURE_ASI)
+
 /*
  * ASI uses a per-CPU tainting model to track what mitigation actions are
  * required on domain transitions. Taints exist along two dimensions:
@@ -131,6 +134,8 @@  struct asi {
 
 DECLARE_PER_CPU_ALIGNED(struct asi *, curr_asi);
 
+void asi_check_boottime_disable(void);
+
 void asi_init_mm_state(struct mm_struct *mm);
 
 int asi_init_class(enum asi_class_id class_id, struct asi_taint_policy *taint_policy);
@@ -155,7 +160,9 @@  void asi_exit(void);
 /* The target is the domain we'll enter when returning to process context. */
 static __always_inline struct asi *asi_get_target(struct task_struct *p)
 {
-	return p->thread.asi_state.target;
+	return static_asi_enabled()
+	       ? p->thread.asi_state.target
+	       : NULL;
 }
 
 static __always_inline void asi_set_target(struct task_struct *p,
@@ -166,7 +173,9 @@  static __always_inline void asi_set_target(struct task_struct *p,
 
 static __always_inline struct asi *asi_get_current(void)
 {
-	return this_cpu_read(curr_asi);
+	return static_asi_enabled()
+	       ? this_cpu_read(curr_asi)
+	       : NULL;
 }
 
 /* Are we currently in a restricted address space? */
@@ -175,7 +184,11 @@  static __always_inline bool asi_is_restricted(void)
 	return (bool)asi_get_current();
 }
 
-/* If we exit/have exited, can we stay that way until the next asi_enter? */
+/*
+ * If we exit/have exited, can we stay that way until the next asi_enter?
+ *
+ * When ASI is disabled, this returns true.
+ */
 static __always_inline bool asi_is_relaxed(void)
 {
 	return !asi_get_target(current);
diff --git a/arch/x86/include/asm/cpufeatures.h b/arch/x86/include/asm/cpufeatures.h
index 913fd3a7bac6506141de65f33b9ee61c615c7d7d..d6a808d10c3b8900d190ea01c66fc248863f05e2 100644
--- a/arch/x86/include/asm/cpufeatures.h
+++ b/arch/x86/include/asm/cpufeatures.h
@@ -474,6 +474,7 @@ 
 #define X86_FEATURE_CLEAR_BHB_HW	(21*32+ 3) /* BHI_DIS_S HW control enabled */
 #define X86_FEATURE_CLEAR_BHB_LOOP_ON_VMEXIT (21*32+ 4) /* Clear branch history at vmexit using SW loop */
 #define X86_FEATURE_FAST_CPPC		(21*32 + 5) /* AMD Fast CPPC */
+#define X86_FEATURE_ASI			(21*32+6) /* Kernel Address Space Isolation */
 
 /*
  * BUG word(s)
diff --git a/arch/x86/include/asm/disabled-features.h b/arch/x86/include/asm/disabled-features.h
index c492bdc97b0595ec77f89dc9b0cefe5e3e64be41..c7964ed4fef8b9441e1c0453da587787d8008d9d 100644
--- a/arch/x86/include/asm/disabled-features.h
+++ b/arch/x86/include/asm/disabled-features.h
@@ -50,6 +50,12 @@ 
 # define DISABLE_PTI		(1 << (X86_FEATURE_PTI & 31))
 #endif
 
+#ifdef CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION
+# define DISABLE_ASI		0
+#else
+# define DISABLE_ASI		(1 << (X86_FEATURE_ASI & 31))
+#endif
+
 #ifdef CONFIG_MITIGATION_RETPOLINE
 # define DISABLE_RETPOLINE	0
 #else
@@ -154,7 +160,7 @@ 
 #define DISABLED_MASK17	0
 #define DISABLED_MASK18	(DISABLE_IBT)
 #define DISABLED_MASK19	(DISABLE_SEV_SNP)
-#define DISABLED_MASK20	0
+#define DISABLED_MASK20	(DISABLE_ASI)
 #define DISABLED_MASK21	0
 #define DISABLED_MASK_CHECK BUILD_BUG_ON_ZERO(NCAPINTS != 22)
 
diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
index 105cd8b43eaf5c20acc80d4916b761559fb95d74..5baf563a078f5b3a6cd4b9f5e92baaf81b0774c4 100644
--- a/arch/x86/mm/asi.c
+++ b/arch/x86/mm/asi.c
@@ -4,6 +4,7 @@ 
 #include <linux/percpu.h>
 #include <linux/spinlock.h>
 
+#include <linux/init.h>
 #include <asm/asi.h>
 #include <asm/cmdline.h>
 #include <asm/cpufeature.h>
@@ -29,6 +30,9 @@  static inline bool asi_class_id_valid(enum asi_class_id class_id)
 
 static inline bool asi_class_initialized(enum asi_class_id class_id)
 {
+	if (!boot_cpu_has(X86_FEATURE_ASI))
+		return 0;
+
 	if (WARN_ON(!asi_class_id_valid(class_id)))
 		return false;
 
@@ -51,6 +55,9 @@  EXPORT_SYMBOL_GPL(asi_init_class);
 
 void asi_uninit_class(enum asi_class_id class_id)
 {
+	if (!boot_cpu_has(X86_FEATURE_ASI))
+		return;
+
 	if (!asi_class_initialized(class_id))
 		return;
 
@@ -66,10 +73,36 @@  const char *asi_class_name(enum asi_class_id class_id)
 	return asi_class_names[class_id];
 }
 
+void __init asi_check_boottime_disable(void)
+{
+	bool enabled = IS_ENABLED(CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION_DEFAULT_ON);
+	char arg[4];
+	int ret;
+
+	ret = cmdline_find_option(boot_command_line, "asi", arg, sizeof(arg));
+	if (ret == 3 && !strncmp(arg, "off", 3)) {
+		enabled = false;
+		pr_info("ASI disabled through kernel command line.\n");
+	} else if (ret == 2 && !strncmp(arg, "on", 2)) {
+		enabled = true;
+		pr_info("Ignoring asi=on param while ASI implementation is incomplete.\n");
+	} else {
+		pr_info("ASI %s by default.\n",
+			enabled ? "enabled" : "disabled");
+	}
+
+	if (enabled)
+		pr_info("ASI enablement ignored due to incomplete implementation.\n");
+}
+
 static void __asi_destroy(struct asi *asi)
 {
-	lockdep_assert_held(&asi->mm->asi_init_lock);
+	WARN_ON_ONCE(asi->ref_count <= 0);
+	if (--(asi->ref_count) > 0)
+		return;
 
+	free_pages((ulong)asi->pgd, PGD_ALLOCATION_ORDER);
+	memset(asi, 0, sizeof(struct asi));
 }
 
 int asi_init(struct mm_struct *mm, enum asi_class_id class_id, struct asi **out_asi)
@@ -79,6 +112,9 @@  int asi_init(struct mm_struct *mm, enum asi_class_id class_id, struct asi **out_
 
 	*out_asi = NULL;
 
+	if (!boot_cpu_has(X86_FEATURE_ASI))
+		return 0;
+
 	if (WARN_ON(!asi_class_initialized(class_id)))
 		return -EINVAL;
 
@@ -122,7 +158,7 @@  void asi_destroy(struct asi *asi)
 {
 	struct mm_struct *mm;
 
-	if (!asi)
+	if (!boot_cpu_has(X86_FEATURE_ASI) || !asi)
 		return;
 
 	if (WARN_ON(!asi_class_initialized(asi->class_id)))
@@ -134,11 +170,7 @@  void asi_destroy(struct asi *asi)
 	 * to block concurrent asi_init calls.
 	 */
 	mutex_lock(&mm->asi_init_lock);
-	WARN_ON_ONCE(asi->ref_count <= 0);
-	if (--(asi->ref_count) == 0) {
-		free_pages((ulong)asi->pgd, PGD_ALLOCATION_ORDER);
-		memset(asi, 0, sizeof(struct asi));
-	}
+	__asi_destroy(asi);
 	mutex_unlock(&mm->asi_init_lock);
 }
 EXPORT_SYMBOL_GPL(asi_destroy);
@@ -255,6 +287,9 @@  static noinstr void __asi_enter(void)
 
 noinstr void asi_enter(struct asi *asi)
 {
+	if (!static_asi_enabled())
+		return;
+
 	VM_WARN_ON_ONCE(!asi);
 
 	/* Should not have an asi_enter() without a prior asi_relax(). */
@@ -269,8 +304,10 @@  EXPORT_SYMBOL_GPL(asi_enter);
 
 noinstr void asi_relax(void)
 {
-	barrier();
-	asi_set_target(current, NULL);
+	if (static_asi_enabled()) {
+		barrier();
+		asi_set_target(current, NULL);
+	}
 }
 EXPORT_SYMBOL_GPL(asi_relax);
 
@@ -279,6 +316,9 @@  noinstr void asi_exit(void)
 	u64 unrestricted_cr3;
 	struct asi *asi;
 
+	if (!static_asi_enabled())
+		return;
+
 	preempt_disable_notrace();
 
 	VM_BUG_ON(this_cpu_read(cpu_tlbstate.loaded_mm) ==
@@ -310,6 +350,9 @@  EXPORT_SYMBOL_GPL(asi_exit);
 
 void asi_init_mm_state(struct mm_struct *mm)
 {
+	if (!boot_cpu_has(X86_FEATURE_ASI))
+		return;
+
 	memset(mm->asi, 0, sizeof(mm->asi));
 	mutex_init(&mm->asi_init_lock);
 }
diff --git a/arch/x86/mm/init.c b/arch/x86/mm/init.c
index de4227ed5169ff84d0ce80b677caffc475198fa6..ded3a47f2a9c1f554824d4ad19f3b48bce271274 100644
--- a/arch/x86/mm/init.c
+++ b/arch/x86/mm/init.c
@@ -28,6 +28,7 @@ 
 #include <asm/text-patching.h>
 #include <asm/memtype.h>
 #include <asm/paravirt.h>
+#include <asm/asi.h>
 
 /*
  * We need to define the tracepoints somewhere, and tlb.c
@@ -251,7 +252,7 @@  static void __init probe_page_size_mask(void)
 	__default_kernel_pte_mask = __supported_pte_mask;
 	/* Except when with PTI where the kernel is mostly non-Global: */
 	if (cpu_feature_enabled(X86_FEATURE_PTI) ||
-	    IS_ENABLED(CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION))
+	    cpu_feature_enabled(X86_FEATURE_ASI))
 		__default_kernel_pte_mask &= ~_PAGE_GLOBAL;
 
 	/* Enable 1 GB linear kernel mappings if available: */
@@ -754,6 +755,7 @@  void __init init_mem_mapping(void)
 	unsigned long end;
 
 	pti_check_boottime_disable();
+	asi_check_boottime_disable();
 	probe_page_size_mask();
 	setup_pcid();
 
diff --git a/include/asm-generic/asi.h b/include/asm-generic/asi.h
index 6b84202837605fa57e4a910318c8353b3f816f06..eedc961ee916a9e1da631ca489ea4a7bc9e6089f 100644
--- a/include/asm-generic/asi.h
+++ b/include/asm-generic/asi.h
@@ -65,6 +65,10 @@  static inline pgd_t *asi_pgd(struct asi *asi) { return NULL; }
 
 static inline void asi_handle_switch_mm(void) { }
 
+#define static_asi_enabled() false
+
+static inline void asi_check_boottime_disable(void) { }
+
 #endif /* !CONFIG_MITIGATION_ADDRESS_SPACE_ISOLATION */
 
 #endif  /* !_ASSEMBLY_ */