[kvm-unit-tests PATCH v4 17/18] x86: cet: Reset IBT tracker state on #CP violations

[Sean Christopherson <seanjc@google.com>]

Reset the IBT tracker state back to IDLE on #CP violations to not
influence follow-up tests with a poisoned starting state.

Opportunistically rename "rvc" to "got_cp" to make it more obvious what
the flag tracks ("rvc" is presumably "raised vector CP"?).

Signed-off-by: Mathias Krause <minipli@grsecurity.net>
[sean: add helper, align indentation, use handler+callback instead of "extra"]
Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 lib/x86/usermode.c | 12 +++++++++---
 lib/x86/usermode.h | 13 ++++++++++---
 x86/cet.c          | 31 +++++++++++++++++++++++++++----
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/lib/x86/usermode.c b/lib/x86/usermode.c
index f896e3bd..b65c5378 100644
--- a/lib/x86/usermode.c
+++ b/lib/x86/usermode.c
@@ -21,12 +21,17 @@ static void restore_exec_to_jmpbuf(void)
 	longjmp(jmpbuf, 1);
 }
 
+static handler ex_callback;
+
 static void restore_exec_to_jmpbuf_exception_handler(struct ex_regs *regs)
 {
 	this_cpu_write_exception_vector(regs->vector);
 	this_cpu_write_exception_rflags_rf((regs->rflags >> 16) & 1);
 	this_cpu_write_exception_error_code(regs->error_code);
 
+	if (ex_callback)
+		ex_callback(regs);
+
 	/* longjmp must happen after iret, so do not do it now.  */
 	regs->rip = (unsigned long)&restore_exec_to_jmpbuf;
 	regs->cs = KERNEL_CS;
@@ -35,9 +40,9 @@ static void restore_exec_to_jmpbuf_exception_handler(struct ex_regs *regs)
 #endif
 }
 
-uint64_t run_in_user(usermode_func func, unsigned int fault_vector,
-		uint64_t arg1, uint64_t arg2, uint64_t arg3,
-		uint64_t arg4, bool *raised_vector)
+uint64_t run_in_user_ex(usermode_func func, unsigned int fault_vector,
+			uint64_t arg1, uint64_t arg2, uint64_t arg3,
+			uint64_t arg4, bool *raised_vector, handler ex_handler)
 {
 	extern char ret_to_kernel;
 	volatile uint64_t rax = 0;
@@ -45,6 +50,7 @@ uint64_t run_in_user(usermode_func func, unsigned int fault_vector,
 	handler old_ex;
 
 	*raised_vector = 0;
+	ex_callback = ex_handler;
 	set_idt_entry(RET_TO_KERNEL_IRQ, &ret_to_kernel, 3);
 	old_ex = handle_exception(fault_vector,
 				  restore_exec_to_jmpbuf_exception_handler);
diff --git a/lib/x86/usermode.h b/lib/x86/usermode.h
index 04e358e2..7eca9079 100644
--- a/lib/x86/usermode.h
+++ b/lib/x86/usermode.h
@@ -20,11 +20,18 @@ typedef uint64_t (*usermode_func)(void);
  * Supports running functions with up to 4 arguments.
  * fault_vector: exception vector that might get thrown during the function.
  * raised_vector: outputs true if exception occurred.
+ * ex_handler: optiona handler to call when handling @fault_vector exceptions
  *
  * returns: return value returned by function, or 0 if an exception occurred.
  */
-uint64_t run_in_user(usermode_func func, unsigned int fault_vector,
-		uint64_t arg1, uint64_t arg2, uint64_t arg3,
-		uint64_t arg4, bool *raised_vector);
+uint64_t run_in_user_ex(usermode_func func, unsigned int fault_vector,
+			uint64_t arg1, uint64_t arg2, uint64_t arg3,
+			uint64_t arg4, bool *raised_vector, handler ex_handler);
 
+static inline uint64_t run_in_user(usermode_func func, unsigned int fault_vector,
+				   uint64_t arg1, uint64_t arg2, uint64_t arg3,
+				   uint64_t arg4, bool *raised_vector)
+{
+	return run_in_user_ex(func, fault_vector, arg1, arg2, arg3, arg4, raised_vector, NULL);
+}
 #endif
diff --git a/x86/cet.c b/x86/cet.c
index 74d3f701..7ffe234b 100644
--- a/x86/cet.c
+++ b/x86/cet.c
@@ -1,4 +1,3 @@
-
 #include "libcflat.h"
 #include "x86/desc.h"
 #include "x86/processor.h"
@@ -85,6 +84,8 @@ static uint64_t cet_ibt_func(void)
 #define CET_ENABLE_SHSTK			BIT(0)
 #define CET_ENABLE_IBT				BIT(2)
 #define CET_ENABLE_NOTRACK			BIT(4)
+#define CET_IBT_SUPPRESS			BIT(10)
+#define CET_IBT_TRACKER_WAIT_FOR_ENDBRANCH	BIT(11)
 
 static void test_shstk(void)
 {
@@ -132,9 +133,31 @@ static void test_shstk(void)
 	report(vector == GP_VECTOR, "MSR_IA32_PL3_SSP alignment test.");
 }
 
+static void ibt_tracker_cp_fixup(struct ex_regs *regs)
+{
+	u64 cet_u = rdmsr(MSR_IA32_U_CET);
+
+	/*
+	 * Switch the IBT tracker state to IDLE to have a clean state for
+	 * following tests.
+	 */
+	if (cet_u & CET_IBT_TRACKER_WAIT_FOR_ENDBRANCH) {
+		cet_u &= ~CET_IBT_TRACKER_WAIT_FOR_ENDBRANCH;
+		printf("CET: suppressing IBT WAIT_FOR_ENDBRANCH state at RIP: %lx\n",
+		       regs->rip);
+		wrmsr(MSR_IA32_U_CET, cet_u);
+	}
+}
+
+static uint64_t ibt_run_in_user(usermode_func func, bool *got_cp)
+{
+	return run_in_user_ex(func, CP_VECTOR, 0, 0, 0, 0, got_cp,
+			      ibt_tracker_cp_fixup);
+}
+
 static void test_ibt(void)
 {
-	bool rvc;
+	bool got_cp;
 
 	if (!this_cpu_has(X86_FEATURE_IBT)) {
 		report_skip("IBT not supported");
@@ -144,8 +167,8 @@ static void test_ibt(void)
 	/* Enable indirect-branch tracking (notrack handling for jump tables) */
 	wrmsr(MSR_IA32_U_CET, CET_ENABLE_IBT | CET_ENABLE_NOTRACK);
 
-	run_in_user(cet_ibt_func, CP_VECTOR, 0, 0, 0, 0, &rvc);
-	report(rvc && exception_error_code() == CP_ERR_ENDBR,
+	ibt_run_in_user(cet_ibt_func, &got_cp);
+	report(got_cp && exception_error_code() == CP_ERR_ENDBR,
 	       "Indirect-branch tracking test");
 }
 
-- 
2.52.0.rc1.455.g30608eb744-goog