Re: [PATCH v3 06/27] shmem/userfaultfd: Handle uffd-wp special pte in page fault handler

[Alistair Popple <apopple@nvidia.com>]

On Friday, 28 May 2021 6:21:22 AM AEST Peter Xu wrote:
> File-backed memories are prone to unmap/swap so the ptes are always unstable.
> This could lead to userfaultfd-wp information got lost when unmapped or swapped
> out on such types of memory, for example, shmem.  To keep such an information
> persistent, we will start to use the newly introduced swap-like special ptes to
> replace a null pte when those ptes were removed.
> 
> Prepare this by handling such a special pte first before it is applied in the
> general page fault handler.
> 
> The handling of this special pte page fault is similar to missing fault, but it
> should happen after the pte missing logic since the special pte is designed to
> be a swap-like pte.  Meanwhile it should be handled before do_swap_page() so
> that the swap core logic won't be confused to see such an illegal swap pte.
> 
> This is a slow path of uffd-wp handling, because unmap of wr-protected shmem
> ptes should be rare.  So far it should only trigger in two conditions:
> 
>   (1) When trying to punch holes in shmem_fallocate(), there will be a
>       pre-unmap optimization before evicting the page.  That will create
>       unmapped shmem ptes with wr-protected pages covered.
> 
>   (2) Swapping out of shmem pages
> 
> Because of this, the page fault handling is simplifed too by not sending the
> wr-protect message in the 1st page fault, instead the page will be installed
> read-only, so the message will be generated until the next write, which will
> trigger the do_wp_page() path of general uffd-wp handling.
> 
> Disable fault-around for all uffd-wp registered ranges for extra safety, and
> clean the code up a bit after we introduced MINOR fault.
> 
> Signed-off-by: Peter Xu <peterx@redhat.com>
> ---
>  include/linux/userfaultfd_k.h | 12 +++++
>  mm/memory.c                   | 88 +++++++++++++++++++++++++++++++----
>  2 files changed, 90 insertions(+), 10 deletions(-)
> 
> diff --git a/include/linux/userfaultfd_k.h b/include/linux/userfaultfd_k.h
> index 93f932b53a71..ca3f794d07e9 100644
> --- a/include/linux/userfaultfd_k.h
> +++ b/include/linux/userfaultfd_k.h
> @@ -94,6 +94,18 @@ static inline bool uffd_disable_huge_pmd_share(struct vm_area_struct *vma)
>  	return vma->vm_flags & (VM_UFFD_WP | VM_UFFD_MINOR);
>  }
>  
> +/*
> + * Don't do fault around for either WP or MINOR registered uffd range.  For
> + * MINOR registered range, fault around will be a total disaster and ptes can
> + * be installed without notifications; for WP it should mostly be fine as long
> + * as the fault around checks for pte_none() before the installation, however
> + * to be super safe we just forbid it.
> + */
> +static inline bool uffd_disable_fault_around(struct vm_area_struct *vma)
> +{
> +	return vma->vm_flags & (VM_UFFD_WP | VM_UFFD_MINOR);
> +}
> +
>  static inline bool userfaultfd_missing(struct vm_area_struct *vma)
>  {
>  	return vma->vm_flags & VM_UFFD_MISSING;
> diff --git a/mm/memory.c b/mm/memory.c
> index 2b24af4616df..45a2f71e447a 100644
> --- a/mm/memory.c
> +++ b/mm/memory.c
> @@ -3917,6 +3917,7 @@ vm_fault_t do_set_pmd(struct vm_fault *vmf, struct page *page)
>  void do_set_pte(struct vm_fault *vmf, struct page *page, unsigned long addr)
>  {
>  	struct vm_area_struct *vma = vmf->vma;
> +	bool uffd_wp = pte_swp_uffd_wp_special(vmf->orig_pte);
>  	bool write = vmf->flags & FAULT_FLAG_WRITE;
>  	bool prefault = vmf->address != addr;
>  	pte_t entry;
> @@ -3929,6 +3930,8 @@ void do_set_pte(struct vm_fault *vmf, struct page *page, unsigned long addr)
>  
>  	if (write)
>  		entry = maybe_mkwrite(pte_mkdirty(entry), vma);
> +	if (unlikely(uffd_wp))
> +		entry = pte_mkuffd_wp(pte_wrprotect(entry));
>  	/* copy-on-write page */
>  	if (write && !(vma->vm_flags & VM_SHARED)) {
>  		inc_mm_counter_fast(vma->vm_mm, MM_ANONPAGES);
> @@ -3996,8 +3999,12 @@ vm_fault_t finish_fault(struct vm_fault *vmf)
>  	vmf->pte = pte_offset_map_lock(vma->vm_mm, vmf->pmd,
>  				      vmf->address, &vmf->ptl);
>  	ret = 0;
> -	/* Re-check under ptl */
> -	if (likely(pte_none(*vmf->pte)))
> +
> +	/*
> +	 * Re-check under ptl.  Note: this will cover both none pte and
> +	 * uffd-wp-special swap pte
> +	 */
> +	if (likely(pte_same(*vmf->pte, vmf->orig_pte)))
>  		do_set_pte(vmf, page, vmf->address);
>  	else
>  		ret = VM_FAULT_NOPAGE;
> @@ -4101,9 +4108,21 @@ static vm_fault_t do_fault_around(struct vm_fault *vmf)
>  	return vmf->vma->vm_ops->map_pages(vmf, start_pgoff, end_pgoff);
>  }
>  
> +/* Return true if we should do read fault-around, false otherwise */
> +static inline bool should_fault_around(struct vm_fault *vmf)
> +{
> +	/* No ->map_pages?  No way to fault around... */
> +	if (!vmf->vma->vm_ops->map_pages)
> +		return false;
> +
> +	if (uffd_disable_fault_around(vmf->vma))
> +		return false;
> +
> +	return fault_around_bytes >> PAGE_SHIFT > 1;
> +}
> +
>  static vm_fault_t do_read_fault(struct vm_fault *vmf)
>  {
> -	struct vm_area_struct *vma = vmf->vma;
>  	vm_fault_t ret = 0;
>  
>  	/*
> @@ -4111,12 +4130,10 @@ static vm_fault_t do_read_fault(struct vm_fault *vmf)
>  	 * if page by the offset is not ready to be mapped (cold cache or
>  	 * something).
>  	 */
> -	if (vma->vm_ops->map_pages && fault_around_bytes >> PAGE_SHIFT > 1) {
> -		if (likely(!userfaultfd_minor(vmf->vma))) {
> -			ret = do_fault_around(vmf);
> -			if (ret)
> -				return ret;
> -		}
> +	if (should_fault_around(vmf)) {
> +		ret = do_fault_around(vmf);
> +		if (ret)
> +			return ret;
>  	}
>  
>  	ret = __do_fault(vmf);
> @@ -4435,6 +4452,57 @@ static vm_fault_t wp_huge_pud(struct vm_fault *vmf, pud_t orig_pud)
>  	return VM_FAULT_FALLBACK;
>  }
>  
> +static vm_fault_t uffd_wp_clear_special(struct vm_fault *vmf)
> +{
> +	vmf->pte = pte_offset_map_lock(vmf->vma->vm_mm, vmf->pmd,
> +				       vmf->address, &vmf->ptl);
> +	/*
> +	 * Be careful so that we will only recover a special uffd-wp pte into a
> +	 * none pte.  Otherwise it means the pte could have changed, so retry.
> +	 */
> +	if (pte_swp_uffd_wp_special(*vmf->pte))
> +		pte_clear(vmf->vma->vm_mm, vmf->address, vmf->pte);
> +	pte_unmap_unlock(vmf->pte, vmf->ptl);
> +	return 0;
> +}
> +
> +/*
> + * This is actually a page-missing access, but with uffd-wp special pte
> + * installed.  It means this pte was wr-protected before being unmapped.
> + */
> +static vm_fault_t uffd_wp_handle_special(struct vm_fault *vmf)
> +{
> +	/* Careful!  vmf->pte unmapped after return */
> +	if (!pte_unmap_same(vmf))
> +		return 0;
> +
> +	/*
> +	 * Just in case there're leftover special ptes even after the region
> +	 * got unregistered - we can simply clear them.
> +	 */
> +	if (unlikely(!userfaultfd_wp(vmf->vma) || vma_is_anonymous(vmf->vma)))
> +		return uffd_wp_clear_special(vmf);
> +
> +	/*
> +	 * Here we share most code with do_fault(), in which we can identify
> +	 * whether this is "none pte fault" or "uffd-wp-special fault" by
> +	 * checking the vmf->orig_pte.
> +	 */
> +	return do_fault(vmf);
> +}
> +
> +static vm_fault_t do_swap_pte(struct vm_fault *vmf)
> +{
> +	/*
> +	 * We need to handle special swap ptes before handling ptes that
> +	 * contain swap entries, always.
> +	 */
> +	if (unlikely(pte_swp_uffd_wp_special(vmf->orig_pte)))
> +		return uffd_wp_handle_special(vmf);
> +
> +	return do_swap_page(vmf);

Probably pretty minor in the scheme of things but why not add this special
case directly to do_swap_page()? Your earlier "shmem/userfaultfd: Handle
uffd-wp special pte in page fault handler" adds this to do_swap_page()
anyway:

	/*
	 * We should never call do_swap_page upon a swap special pte; just be
	 * safe to bail out if it happens.
	 */
	if (WARN_ON_ONCE(is_swap_special_pte(vmf->orig_pte)))
		goto out;

So this patch could instead replace the warning with the call to
uffd_wp_handle_special(), which also means you can remove the extra
pte_unmap_same(vmf) check in uffd_wp_handle_special().

I suppose you might have to worry about other callers of do_swap_page(),
but the only other one I could see was __collapse_huge_page_swapin().
Initially I thought that might be able to trigger the warning here but I
see it checks pte_has_swap_entry() first which should skip it if it finds
the special pte.

 - Alistair

> +}
> +
>  /*
>   * These routines also need to handle stuff like marking pages dirty
>   * and/or accessed for architectures that don't do it in hardware (most
> @@ -4509,7 +4577,7 @@ static vm_fault_t handle_pte_fault(struct vm_fault *vmf)
>  	}
>  
>  	if (!pte_present(vmf->orig_pte))
> -		return do_swap_page(vmf);
> +		return do_swap_pte(vmf);
>  
>  	if (pte_protnone(vmf->orig_pte) && vma_is_accessible(vmf->vma))
>  		return do_numa_page(vmf);
>