Re: [PATCH RFC v2 03/29] mm: asi: Introduce ASI core API

[Brendan Jackman <jackmanb@google.com>]

On Sat, Mar 15, 2025 at 01:36:21PM +0100, Borislav Petkov wrote:
> On Fri, Mar 14, 2025 at 06:34:32PM -0700, Junaid Shahid wrote:
> > The reason this isn't a problem with the current asi_enter() is because
> > there the equivalent of asi_start_critical() happens _before_ the address
> > space switch. That ensures that even if an NMI arrives in the middle of
> > asi_enter(), the NMI epilog will switch to the restricted address space and
> > there is no window where an NMI (or any other interrupt/exception for that
> > matter) would result in going into vmenter with an unrestricted address
> > space.
> 
> Aha.
> 
> > So
> > 	asi_enter();
> > 	asi_start_critical();
> > 	vmenter();
> > 	asi_end_critical();
> > 
> > is broken as there is a problematic window between asi_enter() and
> > asi_start_critical() as Brendan pointed out.
> > 
> > However,
> > 	asi_start_critical();
> > 	asi_enter();
> > 	vmenter();
> > 	asi_end_critical();
> > 
> > would work perfectly fine.
> > 
> > Perhaps that might be the way to refactor the API?
> 
> Ok, let's see if I understand the API better now. And I'm using function names
> which say what they do:
> 
> I guess the flow and needed ingredients should be:
> 
> 	1. asi_lock() or asi_start() or whatnot which does that atomic switch
> 	   of asi target. That tells other code like the NMI glue where in the
> 	   asi context the CPU is so that glue code can know what to do on
> 	   return
> 
> 	2. asi_switch_address_space() - the expensive pagetables build and CR3
> 	   switch

(Not relevant to the present discussion but just taking a chance to
 add some colour: there are no pagetables to build here, they are ready
 to go - e.g. see my separate RFC from last week[0] on how the physmap
 could be managed.

 Also, it's perhaps worth noting that the CR3 write itself isn't
 expected to be the really expensive thing, in the general case we
 expect that to be dwarfed by the cost of the actual security
 mitigations - IBPB etc)

[0] https://lore.kernel.org/all/20250313-asi-page-alloc-v1-0-04972e046cea@google.com/

> 	3. asi_enter_critical_region() - this could be NOP but basically marks
> 	   the beginning of the CPU executing "unsafe" code
> 
> 		<... executes unsafe code... >
> 
> 	4. asi_exit_critical_region() - sets ASI target to NULL, i.e., what
> 	   asi_relax) does now. This also atomic and tells other code, we're
> 	   done with executing unsafe code but we're still running in the
> 	   restricted address space.
> 
> 	   This here can go back to 3 as often as needed.
> 
> 	5. asi_switch_address_space() - this goes back to the unrestricted
> 	   adddress space
> 
> 	6. asi_unlock()
> 
> Close?

I don't understand having both asi_[un]lock() _and_
asi_{start,enter}_critical_region(). The only reason we need the
critical section concept is for the purposes of the NMI glue code you
mentioned in part 1, and that setup must happen before the switch into
the restricted address space.

Also, I don't think we want part 5 inside the asi_lock()->asi_unlock()
region. That seems like the region betwen part 5 and 6, we are in the
unrestricted address space, but the NMI entry code is still set up to
return to the restricted address space on exception return. I think
that would actually be harmless, but it doesn't achieve anything.

The more I talk about it, the more convinced I am that the proper API
should only have two elements, one that says "I'm about to run
untrusted code" and one that says "I've finished running untrusted
code". But...

> 1. you can do empty calls to keep the interface balanced and easy to use
>
> 2. once you can remove asi_exit(), you should be able to replace all in-tree
>    users in one atomic change so that they're all switched to the new,
>    simplified interface

Then what about if we did this:

/* 
 * Begin a region where ASI restricted address spaces _may_ be used.
 *
 * Preemption must be off throughout this region.
 */
static inline void asi_start(void)
{
	/* 
	 * Cannot currently context switch in the restricted adddress
	 * space.
	 */
	lockdep_assert_preemption_disabled();

	/* 
	 * (Actually, this doesn't do anything besides assert, it's
	 * just to help the API make sense).
	 */
}

/*
 * End a region begun by asi_start(). After this, the CPU cannot be in
 * the restricted address space until the next asi_start(). 
 */
static inline void asi_end(void)
{
	/* Leave the restricted address space if we're in it. */
	...
}

/*
 * About to run untrusted code, begin a region that _must_ run in the
 * restricted address space.
 */
void asi_start_critical(void);

/* End a region begun by asi_start_critical(). */
void asi_end_critical(void);

ioctl(KVM_RUN) {
    enter_from_user_mode()
    asi_start()
    while !need_userspace_handling()
        asi_start_critical();
        vmenter();
        asi_end_critical();
    }
    asi_end()
    exit_to_user_mode()
}

Then the API is balanced, and we have a clear migration path towards
the two-element API, i.e. we need to just remove asi_start() and
asi_end(). It also better captures the point about the temporary
simplification: basically the reason why the API is currently
overcomplicated is: if totally arbitrary parts of the kernel can find
themselves in the restricted address space, we have more work to do.
(It's totally possible, but we don't wanna block initial submission on
that work). The simplification is about demarcating what code is and
isn't affected by ASI, so having this "region" kinda helps with that.
Although, because NMIs can also be affected it's a bit of a fuzzy
demarcation...