From 45e1eb7ba4469be3b303ffff7ba233f6963610b9 Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Mon, 4 Sep 2023 21:07:30 +0200
Subject: [PATCH] Set IA32_TSC_AUX if present.

---
 src/arch/x86_64/misc.rs  | 14 ++++++++++++--
 src/arch/x86_64/start.rs |  4 ++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/arch/x86_64/misc.rs b/src/arch/x86_64/misc.rs
index 733e42eb78..20fa33e9e0 100644
--- a/src/arch/x86_64/misc.rs
+++ b/src/arch/x86_64/misc.rs
@@ -1,8 +1,14 @@
 use x86::controlregs::Cr4;
 use x86::cpuid::ExtendedFeatures;
 
-pub unsafe fn init() {
-    let has_ext_feat = |feat: fn(ExtendedFeatures) -> bool| crate::cpuid::cpuid_always().get_extended_feature_info().map_or(false, feat);
+use crate::cpuid::cpuid_always;
+
+pub unsafe fn init(cpu_id: usize) {
+    let has_ext_feat = |feat: fn(ExtendedFeatures) -> bool| {
+        cpuid_always()
+            .get_extended_feature_info()
+            .map_or(false, feat)
+    };
 
     if has_ext_feat(|feat| feat.has_umip()) {
         // UMIP (UserMode Instruction Prevention) forbids userspace from calling SGDT, SIDT, SLDT,
@@ -25,4 +31,8 @@ pub unsafe fn init() {
         // Clear CLAC in (the probably unlikely) case the bootloader set it earlier.
         x86::bits64::rflags::clac();
     }
+
+    if let Some(feats) = cpuid_always().get_extended_processor_and_feature_identifiers() && feats.has_rdtscp() {
+        x86::msr::wrmsr(x86::msr::IA32_TSC_AUX, cpu_id as u64);
+    }
 }
diff --git a/src/arch/x86_64/start.rs b/src/arch/x86_64/start.rs
index 9accbc53fc..5f471a7dc8 100644
--- a/src/arch/x86_64/start.rs
+++ b/src/arch/x86_64/start.rs
@@ -155,7 +155,7 @@ pub unsafe extern fn kstart(args_ptr: *const KernelArgs) -> ! {
         log::init();
 
         // Initialize miscellaneous processor features
-        misc::init();
+        misc::init(0);
 
         // Initialize devices
         device::init();
@@ -234,7 +234,7 @@ pub unsafe extern fn kstart_ap(args_ptr: *const KernelArgsAp) -> ! {
         interrupt::syscall::init();
 
         // Initialize miscellaneous processor features
-        misc::init();
+        misc::init(cpu_id);
 
         // Initialize devices (for AP)
         device::init_ap();