From a333a508a20df6d9c94927a0176d40c896dd52de Mon Sep 17 00:00:00 2001
From: Peter Grehan <grehan@FreeBSD.org>
Date: Mon, 24 Aug 2020 11:49:49 +0000
Subject: [PATCH] cpu_auxmsr: assert caller is preventing CPU migration.

Submitted by:	Adam Fenn (adam at fenn dot io)
Requested by:	kib
Reviewed by:	kib, grehan
Approved by:	kib
MFC after:	3 weeks
Differential Revision:	https://reviews.freebsd.org/D26166
---
 sys/amd64/amd64/initcpu.c | 6 +++++-
 sys/i386/i386/initcpu.c   | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/sys/amd64/amd64/initcpu.c b/sys/amd64/amd64/initcpu.c
index eab1d49f2b73..a4e436420875 100644
--- a/sys/amd64/amd64/initcpu.c
+++ b/sys/amd64/amd64/initcpu.c
@@ -42,6 +42,7 @@ __FBSDID("$FreeBSD$");
 
 #include <machine/cputypes.h>
 #include <machine/md_var.h>
+#include <machine/psl.h>
 #include <machine/specialreg.h>
 
 #include <vm/vm.h>
@@ -218,11 +219,14 @@ init_via(void)
 }
 
 /*
- * The value for the TSC_AUX MSR and rdtscp/rdpid.
+ * The value for the TSC_AUX MSR and rdtscp/rdpid on the invoking CPU.
+ *
+ * Caller should prevent CPU migration.
  */
 u_int
 cpu_auxmsr(void)
 {
+	KASSERT((read_rflags() & PSL_I) == 0, ("context switch possible"));
 	return (PCPU_GET(cpuid));
 }
 
diff --git a/sys/i386/i386/initcpu.c b/sys/i386/i386/initcpu.c
index ba86f30cf31c..04bffe568281 100644
--- a/sys/i386/i386/initcpu.c
+++ b/sys/i386/i386/initcpu.c
@@ -41,6 +41,7 @@ __FBSDID("$FreeBSD$");
 
 #include <machine/cputypes.h>
 #include <machine/md_var.h>
+#include <machine/psl.h>
 #include <machine/specialreg.h>
 
 #include <vm/vm.h>
@@ -628,11 +629,14 @@ init_transmeta(void)
 #endif
 
 /*
- * The value for the TSC_AUX MSR and rdtscp/rdpid.
+ * The value for the TSC_AUX MSR and rdtscp/rdpid on the invoking CPU.
+ *
+ * Caller should prevent CPU migration.
  */
 u_int
 cpu_auxmsr(void)
 {
+	KASSERT((read_eflags() & PSL_I) == 0, ("context switch possible"));
 	return (PCPU_GET(cpuid));
 }