lkl: host ops: fix jump buffer API and implementation

longjmp was called after the function which called setjmp exited (i.e.
lkl_ops->jmp_buf_set), which rendered the context set by setjmp invalid.

To avoid this issue this patch implements a different set jump buffer
API by adding the function to execute as an argument to make sure that
the setjmp context remains valid during the exection of the function.
Thus, calles of the function can safely call longjmp.

Signed-off-by: Octavian Purdila <tavi@cs.pub.ro>

arch/lkl/include/asm/sched.h

@@ -0,0 +1,27 @@
+#ifndef _ASM_LKL_SCHED_H
+#define _ASM_LKL_SCHED_H
+
+#include <linux/sched.h>
+
+static inline void thread_sched_jb(void)
+{
+	set_ti_thread_flag(current_thread_info(), TIF_SCHED_JB);
+
+	if (test_ti_thread_flag(current_thread_info(), TIF_HOST_THREAD)) {
+		set_current_state(TASK_UNINTERRUPTIBLE);
+		lkl_ops->jmp_buf_set(&current_thread_info()->sched_jb,
+				     schedule);
+	} else  {
+		lkl_ops->jmp_buf_set(&current_thread_info()->sched_jb,
+				     lkl_idle_tail_schedule);
+	}
+}
+
+static inline void thread_set_sched_exit(void)
+{
+	set_ti_thread_flag(current_thread_info(), TIF_SCHED_EXIT);
+}
+
+void switch_to_host_task(struct task_struct *);
+
+#endif /*  _ASM_LKL_SCHED_H */

arch/lkl/include/asm/thread_info.h

@@ -63,21 +63,6 @@ void threads_cnt_dec(void);
 #define TIF_HOST_THREAD			9
 #define TIF_IDLE			10
 
-static inline void set_ti_thread_flag(struct thread_info *ti, int flag);
-
-static inline int thread_set_sched_jmp(void)
-{
-	set_ti_thread_flag(current_thread_info(), TIF_SCHED_JB);
-	return lkl_ops->jmp_buf_set(&current_thread_info()->sched_jb);
-}
-
-static inline void thread_set_sched_exit(void)
-{
-	set_ti_thread_flag(current_thread_info(), TIF_SCHED_EXIT);
-}
-
-void switch_to_host_task(struct task_struct *);
-
 #define __HAVE_THREAD_FUNCTIONS
 
 #define task_thread_info(task)	((struct thread_info *)(task)->stack)

arch/lkl/include/uapi/asm/host_ops.h

@@ -70,6 +70,17 @@ struct lkl_jmp_buf {
  *
  * @gettid - returns the host thread id of the caller, which need not
  * be the same as the handle returned by thread_create
+ *
+ * @jmp_buf_set - runs the give function and setups a jump back point by saving
+ * the context in the jump buffer; jmp_buf_longjmp can be called from the give
+ * function or any callee in that function to return back to the jump back
+ * point
+ *
+ * NOTE: we can't return from jmp_buf_set before calling jmp_buf_longjmp or
+ * otherwise the saved context (stack) is not going to be valid, so we must pass
+ * the function that will eventually call longjmp here
+ *
+ * @jmp_buf_longjmp - perform a jump back to the saved jump buffer
  */
 struct lkl_host_operations {
 	const char *virtio_devices;
@@ -114,7 +125,7 @@ struct lkl_host_operations {
 
 	long (*gettid)(void);
 
-	int (*jmp_buf_set)(struct lkl_jmp_buf *jmpb);
+	void (*jmp_buf_set)(struct lkl_jmp_buf *jmpb, void (*f)(void));
 	void (*jmp_buf_longjmp)(struct lkl_jmp_buf *jmpb, int val);
 };
 

arch/lkl/kernel/cpu.c

@@ -7,6 +7,7 @@
 #include <asm/cpu.h>
 #include <asm/thread_info.h>
 #include <asm/unistd.h>
+#include <asm/sched.h>
 #include <asm/syscalls.h>
 
 
@@ -137,14 +138,7 @@ void lkl_cpu_put(void)
 		if (in_interrupt())
 			lkl_bug("%s: in interrupt\n", __func__);
 		lkl_ops->mutex_unlock(cpu.lock);
-		if (test_thread_flag(TIF_HOST_THREAD)) {
-			set_current_state(TASK_UNINTERRUPTIBLE);
-			if (!thread_set_sched_jmp())
-				schedule();
-		} else {
-			if (!thread_set_sched_jmp())
-				lkl_idle_tail_schedule();
-		}
+		thread_sched_jb();
 		return;
 	}
 
@@ -242,10 +236,8 @@ void arch_cpu_idle_prepare(void)
 	 * We hijack the idle loop here so that we can let the idle thread
 	 * jump back to the beginning.
 	 */
-	while (1) {
-		if (!lkl_ops->jmp_buf_set(&cpu.idle_jb))
-			cpu_idle_loop();
-	}
+	while (1)
+		lkl_ops->jmp_buf_set(&cpu.idle_jb, cpu_idle_loop);
 }
 
 void lkl_cpu_wakeup_idle(void)

arch/lkl/kernel/syscalls.c

@@ -14,6 +14,7 @@
 #include <asm/syscalls.h>
 #include <asm/syscalls_32.h>
 #include <asm/cpu.h>
+#include <asm/sched.h>
 
 static asmlinkage long sys_virtio_mmio_device_add(long base, long size,
 						  unsigned int irq);
@@ -114,9 +115,7 @@ long lkl_syscall(long no, long *params)
 	ret = run_syscall(no, params);
 
 	if (no == __NR_reboot) {
-		set_current_state(TASK_UNINTERRUPTIBLE);
-		if (!thread_set_sched_jmp())
-			schedule();
+		thread_sched_jb();
 		return ret;
 	}
 

arch/lkl/kernel/threads.c

@@ -3,6 +3,7 @@
 #include <linux/sched.h>
 #include <asm/host_ops.h>
 #include <asm/cpu.h>
+#include <asm/sched.h>
 
 static volatile int threads_counter;
 
@@ -140,15 +141,7 @@ void switch_to_host_task(struct task_struct *task)
 	task_thread_info(task)->tid = lkl_ops->thread_self();
 
 	wake_up_process(task);
-	if (test_thread_flag(TIF_HOST_THREAD)) {
-		set_current_state(TASK_UNINTERRUPTIBLE);
-		if (!thread_set_sched_jmp())
-			schedule();
-	} else {
-		if (!thread_set_sched_jmp())
-			lkl_idle_tail_schedule();
-	}
-
+	thread_sched_jb();
 	lkl_ops->sem_down(task_thread_info(task)->sched_sem);
 	schedule_tail(abs_prev);
 }

tools/lkl/lib/jmp_buf.c

@@ -1,9 +1,10 @@
 #include <setjmp.h>
 #include <lkl_host.h>
 
-int jmp_buf_set(struct lkl_jmp_buf *jmpb)
+void jmp_buf_set(struct lkl_jmp_buf *jmpb, void (*f)(void))
 {
-	return setjmp(*((jmp_buf *)jmpb->buf));
+	if (!setjmp(*((jmp_buf *)jmpb->buf)))
+		f();
 }
 
 void jmp_buf_longjmp(struct lkl_jmp_buf *jmpb, int val)

tools/lkl/lib/jmp_buf.h

@@ -1,7 +1,7 @@
 #ifndef _LKL_LIB_JMP_BUF_H
 #define _LKL_LIB_JMP_BUF_H
 
-int jmp_buf_set(struct lkl_jmp_buf *jmpb);
+void jmp_buf_set(struct lkl_jmp_buf *jmpb, void (*f)(void));
 void jmp_buf_longjmp(struct lkl_jmp_buf *jmpb, int val);
 
 #endif