Re: [PATCH bpf 1/3] bpf: fix mm lifecycle in open-coded task_vma iterator

[Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>]

Puranjay Mohan <puranjay12@gmail.com> writes:

> On Thu, Mar 5, 2026 at 4:34 PM Mykyta Yatsenko
> <mykyta.yatsenko5@gmail.com> wrote:
>>
>> Puranjay Mohan <puranjay@kernel.org> writes:
>>
>> > The open-coded task_vma BPF iterator reads task->mm and acquires
>> > mmap_read_trylock() but never calls mmget(). This violates refcount
>> > discipline: the mm can reach mm_users == 0 if the task exits while the
>> > iterator holds the lock.
>> >
>> > Add mmget_not_zero() before mmap_read_trylock(). On the error path
>> > after mmget succeeds, the mm reference must be dropped. mmput() can
>> > sleep (exit_mmap, etc.) so it is unsuitable from BPF context.
>> > mmput_async() is safe from hardirq but not from NMI, because
>> > schedule_work() -> queue_work_on() takes raw_spin_lock(&pool->lock)
>> > which can deadlock if the NMI interrupted code holding that lock.
>> >
>> > Add a dedicated per-CPU irq_work (bpf_iter_mmput_work) and a helper
>> > bpf_iter_mmput() that calls mmput_async() directly when not in NMI,
>> > or defers to the irq_work callback when in NMI context. Use it in
>> > both the _new() error path and _destroy(). Add bpf_iter_mmput_busy()
>> > to check irq_work slot availability, and use it alongside
>> > bpf_mmap_unlock_get_irq_work() in _new() to verify both slots are
>> > free before acquiring references.
>> >
>> > Fixes: 4ac454682158 ("bpf: Introduce task_vma open-coded iterator kfuncs")
>> > Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
>> > ---
>> >  kernel/bpf/task_iter.c | 77 +++++++++++++++++++++++++++++++++++++++---
>> >  1 file changed, 73 insertions(+), 4 deletions(-)
>> >
>> > diff --git a/kernel/bpf/task_iter.c b/kernel/bpf/task_iter.c
>> > index 98d9b4c0daff..d3fa8ba0a896 100644
>> > --- a/kernel/bpf/task_iter.c
>> > +++ b/kernel/bpf/task_iter.c
>> > @@ -813,6 +813,55 @@ struct bpf_iter_task_vma_kern {
>> >       struct bpf_iter_task_vma_kern_data *data;
>> >  } __attribute__((aligned(8)));
>> >
>> > +/*
>> > + * Per-CPU irq_work for NMI-safe mmput.
>> > + *
>> > + * mmput_async() is safe from hardirq context but not from NMI, because
>> > + * schedule_work() -> queue_work_on() takes raw_spin_lock(&pool->lock)
>> > + * which can deadlock if the NMI interrupted code holding that lock.
>> > + *
>> > + * This dedicated irq_work defers mmput to hardirq context where
>> > + * mmput_async() is safe. BPF programs are non-preemptible, so one
>> > + * slot per CPU is sufficient.
>> > + */
>> > +struct bpf_iter_mmput_irq_work {
>> > +     struct irq_work irq_work;
>> > +     struct mm_struct *mm;
>> > +};
>> struct mmap_unlock_irq_work is exactly the same struct, perhaps an
>> additional patch renaming it to something like struct bpf_iter_mm_irq_work
>> is needed. Then we can reuse it for mmput.
>
> They are similar but do different things, mmap_unlock_irq_work is used
> to defer from NMI and hard-irq, but mmput_async is safe to run from
> hard-irq and only needs to be deferred from NMI. I had thought of
> combining these but later felt that keeping them seperate would be a
> better approach.
Sure, I just suggest to reuse the structure, not combining the code.
It's very generic: mm_struct + irq_work, basically some async op on the
mm_struct. I agree with you it's better to keep the code separate
in this case.