It's possible for irq_work_queue() to fail if the work has already been
claimed.  That can happen if a TWA_NMI_CURRENT task work is requested
before a previous TWA_NMI_CURRENT IRQ work on the same CPU has gotten a
chance to run.

The error has to be checked before the write to task->task_works.  Also
the try_cmpxchg() loop isn't needed in NMI context.  The TWA_NMI_CURRENT
case really is special, keep things simple by keeping its code all
together in one place.

Fixes: 466e4d801cd4 ("task_work: Add TWA_NMI_CURRENT as an additional notify 
mode.")
Signed-off-by: Josh Poimboeuf <jpoim...@kernel.org>
---
 kernel/task_work.c | 39 ++++++++++++++++++++++++++-------------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/kernel/task_work.c b/kernel/task_work.c
index c969f1f26be5..92024a8bfe12 100644
--- a/kernel/task_work.c
+++ b/kernel/task_work.c
@@ -58,25 +58,38 @@ int task_work_add(struct task_struct *task, struct 
callback_head *work,
        int flags = notify & TWA_FLAGS;
 
        notify &= ~TWA_FLAGS;
+
        if (notify == TWA_NMI_CURRENT) {
-               if (WARN_ON_ONCE(task != current))
+               if (WARN_ON_ONCE(!in_nmi() || task != current))
                        return -EINVAL;
                if (!IS_ENABLED(CONFIG_IRQ_WORK))
                        return -EINVAL;
-       } else {
-               /*
-                * Record the work call stack in order to print it in KASAN
-                * reports.
-                *
-                * Note that stack allocation can fail if TWAF_NO_ALLOC flag
-                * is set and new page is needed to expand the stack buffer.
-                */
-               if (flags & TWAF_NO_ALLOC)
-                       kasan_record_aux_stack_noalloc(work);
-               else
-                       kasan_record_aux_stack(work);
+#ifdef CONFIG_IRQ_WORK
+               head = task->task_works;
+               if (unlikely(head == &work_exited))
+                       return -ESRCH;
+
+               if (!irq_work_queue(this_cpu_ptr(&irq_work_NMI_resume)))
+                       return -EBUSY;
+
+               work->next = head;
+               task->task_works = work;
+#endif
+               return 0;
        }
 
+       /*
+        * Record the work call stack in order to print it in KASAN
+        * reports.
+        *
+        * Note that stack allocation can fail if TWAF_NO_ALLOC flag
+        * is set and new page is needed to expand the stack buffer.
+        */
+       if (flags & TWAF_NO_ALLOC)
+               kasan_record_aux_stack_noalloc(work);
+       else
+               kasan_record_aux_stack(work);
+
        head = READ_ONCE(task->task_works);
        do {
                if (unlikely(head == &work_exited))
-- 
2.48.1


Reply via email to