Re: [RFC PATCH 1/3] x86/mm: fix LAM cr3 mask inconsistency during context switch

[Dave Hansen <dave.hansen@intel.com>]

I know we all have different rules, but any time you could spend absorbing:

	https://www.kernel.org/doc/html/next/process/maintainer-tip.html

would be appreciated, especially:

> The condensed patch description in the subject line should start with
> a uppercase letter and should be written in imperative tone.


On 3/7/24 05:39, Yosry Ahmed wrote:
> In switch_mm_irqs_off(), we read the 'mm->context.lam_cr3_mask' into
> 'new_lam', which is later passed to load_new_mm_cr3(). However, there is
> a call to set_tlbstate_lam_mode() in between which will read
> 'mm->context.lam_cr3_mask' again and set 'cpu_tlbstate.lam' accordingly.
> If we race with another thread updating 'mm->context.lam_cr3_mask', the
> value in 'cpu_tlbstate.lam' could end up being different from CR3.

Your description is fine (modulo the we's).  But I slightly reworded it
to make it more plainly readable:

LAM can only be enabled when a process is single-threaded.  But _kernel_
threads can temporarily use a single-threaded process's mm.  That means
that a context-switching kernel thread can race and observe the mm's LAM
metadata (mm->context.lam_cr3_mask) change.

The context switch code does two logical things with that metadata:
populate CR3 and populate 'cpu_tlbstate.lam'.  If it hits this race,
'cpu_tlbstate.lam' and CR3 can end up out of sync.

This de-synchronization is currently harmless.  But it is confusing and
might lead to warnings or real bugs.

--

> Fix the problem by updating set_tlbstate_lam_mode() to return the LAM
> mask that was set to 'cpu_tlbstate.lam', and use that mask in
> switch_mm_irqs_off() when writing CR3. Use READ_ONCE to make sure we
> read the mask once and use it consistenly.

Spell checking is also appreciated.

...
> -static inline void set_tlbstate_lam_mode(struct mm_struct *mm)
> +static inline unsigned long set_tlbstate_lam_mode(struct mm_struct *mm)
>  {
> -	this_cpu_write(cpu_tlbstate.lam,
> -		       mm->context.lam_cr3_mask >> X86_CR3_LAM_U57_BIT);
> +	unsigned long lam = READ_ONCE(mm->context.lam_cr3_mask);
> +
> +	this_cpu_write(cpu_tlbstate.lam, lam >> X86_CR3_LAM_U57_BIT);
>  	this_cpu_write(tlbstate_untag_mask, mm->context.untag_mask);
> +	return lam;
>  }

The comments about races need to be _here_ so that the purpose of the
READ_ONCE() is clear.

It would also be nice to call out the rule that this can only
meaningfully be called once per context switch.

> @@ -633,7 +628,12 @@ void switch_mm_irqs_off(struct mm_struct *unused, struct mm_struct *next,
>  		barrier();
>  	}
>  
> -	set_tlbstate_lam_mode(next);
> +	/*
> +	 * Even if we are not actually switching mm's, another thread could have
> +	 * updated mm->context.lam_cr3_mask. Make sure tlbstate_lam_cr3_mask()
> +	 * and the loaded CR3 use the up-to-date mask.
> +	 */

I kinda dislike how the comment talks about the details of what
set_tlbstate_lam_mode() does.  It would be much better to put the meat
of this comment at the set_tlbstate_lam_mode() definition.

> +	new_lam = set_tlbstate_lam_mode(next);
>  	if (need_flush) {
>  		this_cpu_write(cpu_tlbstate.ctxs[new_asid].ctx_id, next->context.ctx_id);
>  		this_cpu_write(cpu_tlbstate.ctxs[new_asid].tlb_gen, next_tlb_gen);

This is less a complaint about your change and more of the existing
code, but I wish it was more obvious that set_tlbstate_lam_mode() is
logically shuffling data (once) from 'next' into the tlbstate.

The naming makes it sound like it is modifying the tlbstate of 'next'.

But I don't have any particularly brilliant ideas to fix it either.
Maybe just:

	/* new_lam is effectively cpu_tlbstate.lam */

> @@ -705,7 +705,6 @@ void initialize_tlbstate_and_flush(void)
>  
>  	/* LAM expected to be disabled */
>  	WARN_ON(cr3 & (X86_CR3_LAM_U48 | X86_CR3_LAM_U57));
> -	WARN_ON(mm_lam_cr3_mask(mm));
>  
>  	/*
>  	 * Assert that CR4.PCIDE is set if needed.  (CR4.PCIDE initialization
> @@ -724,7 +723,7 @@ void initialize_tlbstate_and_flush(void)
>  	this_cpu_write(cpu_tlbstate.next_asid, 1);
>  	this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id, mm->context.ctx_id);
>  	this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen, tlb_gen);
> -	set_tlbstate_lam_mode(mm);
> +	WARN_ON(set_tlbstate_lam_mode(mm));

The "set_" naming bugs me in both of the sites that get modified here.
I'd be with a new name that fits better, if we can think of one.