When a time slice extension was granted in the need_resched() check on exit
to user space, the task can still be scheduled out in one of the other
pending work items. When it gets scheduled back in, and need_resched() is
not set, then the stale grant would be preserved, which is just wrong.

RSEQ already keeps track of that and sets TIF_RSEQ, which invokes the
critical section and ID update mechanisms.

Utilize them and clear the user space slice control member of struct rseq
unconditionally within the existing user access sections. That's just an
unconditional store more in that path.

Signed-off-by: Thomas Gleixner <tglx@linutronix.de>
Cc: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
Cc: Peter Zijlstra <peterz@infradead.org>
Cc: "Paul E. McKenney" <paulmck@kernel.org>
Cc: Boqun Feng <boqun.feng@gmail.com>
---
 include/linux/rseq_entry.h |   28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

--- a/include/linux/rseq_entry.h
+++ b/include/linux/rseq_entry.h
@@ -103,9 +103,17 @@ static __always_inline bool rseq_arm_sli
 	return __rseq_arm_slice_extension_timer();
 }
 
+static __always_inline void rseq_slice_clear_grant(struct task_struct *t)
+{
+	if (IS_ENABLED(CONFIG_RSEQ_STATS) && t->rseq.slice.state.granted)
+		rseq_stat_inc(rseq_stats.s_revoked);
+	t->rseq.slice.state.granted = false;
+}
+
 #else /* CONFIG_RSEQ_SLICE_EXTENSION */
 static inline bool rseq_slice_extension_enabled(void) { return false; }
 static inline bool rseq_arm_slice_extension_timer(void) { return false; }
+static inline void rseq_slice_clear_grant(struct task_struct *t) { }
 #endif /* !CONFIG_RSEQ_SLICE_EXTENSION */
 
 bool rseq_debug_update_user_cs(struct task_struct *t, struct pt_regs *regs, unsigned long csaddr);
@@ -404,6 +412,13 @@ bool rseq_set_ids_get_csaddr(struct task
 	unsafe_put_user(ids->mm_cid, &rseq->mm_cid, efault);
 	if (csaddr)
 		unsafe_get_user(*csaddr, &rseq->rseq_cs, efault);
+
+	/* Open coded, so it's in the same user access region */
+	if (rseq_slice_extension_enabled()) {
+		/* Unconditionally clear it, no point in conditionals */
+		unsafe_put_user(0U, &rseq->slice_ctrl, efault);
+		rseq_slice_clear_grant(t);
+	}
 	user_access_end();
 
 	/* Cache the new values */
@@ -518,10 +533,19 @@ static __always_inline bool __rseq_exit_
 		 * If IDs have not changed rseq_event::user_irq must be true
 		 * See rseq_sched_switch_event().
 		 */
+		struct rseq __user *rseq = t->rseq.usrptr;
 		u64 csaddr;
 
-		if (unlikely(get_user_masked_u64(&csaddr, &t->rseq.usrptr->rseq_cs)))
+		if (!user_rw_masked_begin(rseq))
 			goto fail;
+		unsafe_get_user(csaddr, &rseq->rseq_cs, fault);
+		/* Open coded, so it's in the same user access region */
+		if (rseq_slice_extension_enabled()) {
+			/* Unconditionally clear it, no point in conditionals */
+			unsafe_put_user(0U, &rseq->slice_ctrl, fault);
+			rseq_slice_clear_grant(t);
+		}
+		user_access_end();
 
 		if (static_branch_unlikely(&rseq_debug_enabled) || unlikely(csaddr)) {
 			if (unlikely(!rseq_update_user_cs(t, regs, csaddr)))
@@ -545,6 +569,8 @@ static __always_inline bool __rseq_exit_
 	t->rseq.event.events = 0;
 	return false;
 
+fault:
+	user_access_end();
 fail:
 	pagefault_enable();
 	/* Force it into the slow path. Don't clear the state! */