Re: [PATCH v3 3/3] userfaultfd: use per-vma locks in userfaultfd operations

[Jann Horn <jannh@google.com>]

On Tue, Feb 6, 2024 at 2:09 AM Lokesh Gidra <lokeshgidra@google.com> wrote:
> All userfaultfd operations, except write-protect, opportunistically use
> per-vma locks to lock vmas. On failure, attempt again inside mmap_lock
> critical section.
>
> Write-protect operation requires mmap_lock as it iterates over multiple
> vmas.
>
> Signed-off-by: Lokesh Gidra <lokeshgidra@google.com>
[...]
> diff --git a/mm/memory.c b/mm/memory.c
> index b05fd28dbce1..393ab3b0d6f3 100644
> --- a/mm/memory.c
> +++ b/mm/memory.c
[...]
> +/*
> + * lock_vma() - Lookup and lock VMA corresponding to @address.
> + * @prepare_anon: If true, then prepare the VMA (if anonymous) with anon_vma.
> + *
> + * Should be called without holding mmap_lock. VMA should be unlocked after use
> + * with unlock_vma().
> + *
> + * Return: A locked VMA containing @address, NULL of no VMA is found, or
> + * -ENOMEM if anon_vma couldn't be allocated.
> + */
> +struct vm_area_struct *lock_vma(struct mm_struct *mm,
> +                               unsigned long address,
> +                               bool prepare_anon)
> +{
> +       struct vm_area_struct *vma;
> +
> +       vma = lock_vma_under_rcu(mm, address);
> +
> +       if (vma)
> +               return vma;
> +
> +       mmap_read_lock(mm);
> +       vma = vma_lookup(mm, address);
> +       if (vma) {
> +               if (prepare_anon && vma_is_anonymous(vma) &&
> +                   anon_vma_prepare(vma))
> +                       vma = ERR_PTR(-ENOMEM);
> +               else
> +                       vma_acquire_read_lock(vma);

This new code only calls anon_vma_prepare() for VMAs where
vma_is_anonymous() is true (meaning they are private anonymous).

[...]
> diff --git a/mm/userfaultfd.c b/mm/userfaultfd.c
> index 74aad0831e40..64e22e467e4f 100644
> --- a/mm/userfaultfd.c
> +++ b/mm/userfaultfd.c
> @@ -19,20 +19,25 @@
>  #include <asm/tlb.h>
>  #include "internal.h"
>
> -static __always_inline
> -struct vm_area_struct *find_dst_vma(struct mm_struct *dst_mm,
> -                                   unsigned long dst_start,
> -                                   unsigned long len)
> +/* Search for VMA and make sure it is valid. */
> +static struct vm_area_struct *find_and_lock_dst_vma(struct mm_struct *dst_mm,
> +                                                   unsigned long dst_start,
> +                                                   unsigned long len)
>  {
> -       /*
> -        * Make sure that the dst range is both valid and fully within a
> -        * single existing vma.
> -        */
>         struct vm_area_struct *dst_vma;
>
> -       dst_vma = find_vma(dst_mm, dst_start);
> -       if (!range_in_vma(dst_vma, dst_start, dst_start + len))
> -               return NULL;
> +       /* Ensure anon_vma is assigned for anonymous vma */
> +       dst_vma = lock_vma(dst_mm, dst_start, true);

lock_vma() is now used by find_and_lock_dst_vma(), which is used by
mfill_atomic().

> +       if (!dst_vma)
> +               return ERR_PTR(-ENOENT);
> +
> +       if (PTR_ERR(dst_vma) == -ENOMEM)
> +               return dst_vma;
> +
> +       /* Make sure that the dst range is fully within dst_vma. */
> +       if (dst_start + len > dst_vma->vm_end)
> +               goto out_unlock;
>
>         /*
>          * Check the vma is registered in uffd, this is required to
[...]
> @@ -597,7 +599,15 @@ static __always_inline ssize_t mfill_atomic(struct userfaultfd_ctx *ctx,
>         copied = 0;
>         folio = NULL;
>  retry:
> -       mmap_read_lock(dst_mm);
> +       /*
> +        * Make sure the vma is not shared, that the dst range is
> +        * both valid and fully within a single existing vma.
> +        */
> +       dst_vma = find_and_lock_dst_vma(dst_mm, dst_start, len);
> +       if (IS_ERR(dst_vma)) {
> +               err = PTR_ERR(dst_vma);
> +               goto out;
> +       }
>
>         /*
>          * If memory mappings are changing because of non-cooperative
> @@ -609,15 +619,6 @@ static __always_inline ssize_t mfill_atomic(struct userfaultfd_ctx *ctx,
>         if (atomic_read(&ctx->mmap_changing))
>                 goto out_unlock;
>
> -       /*
> -        * Make sure the vma is not shared, that the dst range is
> -        * both valid and fully within a single existing vma.
> -        */
> -       err = -ENOENT;
> -       dst_vma = find_dst_vma(dst_mm, dst_start, len);
> -       if (!dst_vma)
> -               goto out_unlock;
> -
>         err = -EINVAL;
>         /*
>          * shmem_zero_setup is invoked in mmap for MAP_ANONYMOUS|MAP_SHARED but
> @@ -647,16 +648,6 @@ static __always_inline ssize_t mfill_atomic(struct userfaultfd_ctx *ctx,
>             uffd_flags_mode_is(flags, MFILL_ATOMIC_CONTINUE))
>                 goto out_unlock;
>
> -       /*
> -        * Ensure the dst_vma has a anon_vma or this page
> -        * would get a NULL anon_vma when moved in the
> -        * dst_vma.
> -        */
> -       err = -ENOMEM;
> -       if (!(dst_vma->vm_flags & VM_SHARED) &&
> -           unlikely(anon_vma_prepare(dst_vma)))
> -               goto out_unlock;

But the check mfill_atomic() used to do was different, it checked for VM_SHARED.

Each VMA has one of these three types:

1. shared (marked by VM_SHARED; does not have an anon_vma)
2. private file-backed (needs to have anon_vma when storing PTEs)
3. private anonymous (what vma_is_anonymous() detects; needs to have
anon_vma when storing PTEs)

This old code would call anon_vma_prepare() for both private VMA types
(which is correct). The new code only calls anon_vma_prepare() for
private anonymous VMAs, not for private file-backed ones. I think this
code will probably crash with a BUG_ON() in __folio_set_anon() if you
try to use userfaultfd to insert a PTE into a private file-backed VMA
of a shmem file. (Which you should be able to get by creating a file
in /dev/shm/ and then mapping that file with mmap(NULL, <size>,
PROT_READ|PROT_WRITE, MAP_PRIVATE, <fd>, 0).)