kunit: Fix race condition in try-catch completion

KUnit's try-catch infrastructure now uses vfork_done, which is always
set to a valid completion when a kthread is created, but which is set to
NULL once the thread terminates. This creates a race condition, where
the kthread exits before we can wait on it.

Keep a copy of vfork_done, which is taken before we wake_up_process()
and so valid, and wait on that instead.

Fixes: 93533996100c ("kunit: Handle test faults")
Reported-by: Linux Kernel Functional Testing <[e-mail]>
Closes: https://lore.kernel.org/lkml/[e-mail]/
Tested-by: Linux Kernel Functional Testing <[e-mail]>
Acked-by: Mickaël Salaün <[e-mail]>
Signed-off-by: David Gow <[e-mail]>
Reviewed-by: Rae Moar <[e-mail]>
Tested-by: Miguel Ojeda <[e-mail]>
Signed-off-by: Shuah Khan <[e-mail]>

lib/kunit/try-catch.c

···
 {
 	struct kunit *test = try_catch->test;
 	struct task_struct *task_struct;
+	struct completion *task_done;
 	int exit_code, time_remaining;
 
 	try_catch->context = context;
···
 		return;
 	}
 	get_task_struct(task_struct);
-	wake_up_process(task_struct);
 	/*
 	 * As for a vfork(2), task_struct->vfork_done (pointing to the
 	 * underlying kthread->exited) can be used to wait for the end of a
-	 * kernel thread.
+	 * kernel thread. It is set to NULL when the thread exits, so we
+	 * keep a copy here.
 	 */
-	time_remaining = wait_for_completion_timeout(task_struct->vfork_done,
+	task_done = task_struct->vfork_done;
+	wake_up_process(task_struct);
+
+	time_remaining = wait_for_completion_timeout(task_done,
 						     kunit_test_timeout());
 	if (time_remaining == 0) {
 		try_catch->try_result = -ETIMEDOUT;