From df0cfab3a675bfe382824edbda06a9e4608191f5 Mon Sep 17 00:00:00 2001
From: Andri Saar <andrisaar@google.com>
Date: Tue, 23 Jul 2024 18:47:31 +0000
Subject: [PATCH] stage0 HAL: shuffle code that reads/writes MSRs to the HAL.

This is mostly just copying code around as the original implementation
was fairly generic to begin with.

Bug: 350496083
Change-Id: Ia70d5f10db6aee1cca041482ae081bd0d601b2b8
---
 stage0/src/hal/base/mod.rs |  1 +
 stage0/src/hal/mod.rs      | 41 +++++++++++++++++++++++++++++++++
 stage0/src/hal/sev/mod.rs  |  2 ++
 stage0/src/hal/sev/msr.rs  | 43 ++++++++++++++++++++++++++++++++++
 stage0/src/msr.rs          | 47 ++------------------------------------
 5 files changed, 89 insertions(+), 45 deletions(-)
 create mode 100644 stage0/src/hal/sev/msr.rs

diff --git a/stage0/src/hal/base/mod.rs b/stage0/src/hal/base/mod.rs
index 82469adeee5..2a7795458ce 100644
--- a/stage0/src/hal/base/mod.rs
+++ b/stage0/src/hal/base/mod.rs
@@ -19,3 +19,4 @@ mod mmio;
 
 pub use cpuid::*;
 pub use mmio::*;
+pub use x86_64::registers::model_specific::Msr;
diff --git a/stage0/src/hal/mod.rs b/stage0/src/hal/mod.rs
index d707c5bde74..3c3f420c4d5 100644
--- a/stage0/src/hal/mod.rs
+++ b/stage0/src/hal/mod.rs
@@ -86,3 +86,44 @@ pub fn cpuid(leaf: u32) -> CpuidResult {
     #[cfg(not(feature = "sev"))]
     return base::cpuid(leaf);
 }
+/// Wrapper that can access a MSR either directly or through the GHCB, depending
+/// on the environment.
+pub struct Msr {
+    #[cfg(feature = "sev")]
+    msr_id: u32,
+    msr: base::Msr,
+}
+
+impl Msr {
+    pub const fn new(reg: u32) -> Self {
+        Self {
+            #[cfg(feature = "sev")]
+            msr_id: reg,
+            msr: base::Msr::new(reg),
+        }
+    }
+
+    /// Read the MSR.
+    ///
+    /// ## Safety
+    ///
+    /// The caller must guarantee that the MSR is valid.
+    pub unsafe fn read(&self) -> u64 {
+        #[cfg(feature = "sev")]
+        return sev::read_msr(&self.msr, self.msr_id);
+        #[cfg(not(feature = "sev"))]
+        return self.msr.read();
+    }
+
+    /// Write the MSR.
+    ///
+    /// ## Safety
+    ///
+    /// The caller must guarantee that the MSR is valid.
+    pub unsafe fn write(&mut self, val: u64) {
+        #[cfg(feature = "sev")]
+        return sev::write_msr(&mut self.msr, self.msr_id, val);
+        #[cfg(not(feature = "sev"))]
+        return self.msr.write(val);
+    }
+}
diff --git a/stage0/src/hal/sev/mod.rs b/stage0/src/hal/sev/mod.rs
index 82469adeee5..c7ce09701d6 100644
--- a/stage0/src/hal/sev/mod.rs
+++ b/stage0/src/hal/sev/mod.rs
@@ -16,6 +16,8 @@
 
 mod cpuid;
 mod mmio;
+mod msr;
 
 pub use cpuid::*;
 pub use mmio::*;
+pub use msr::*;
diff --git a/stage0/src/hal/sev/msr.rs b/stage0/src/hal/sev/msr.rs
new file mode 100644
index 00000000000..14e7647b6a1
--- /dev/null
+++ b/stage0/src/hal/sev/msr.rs
@@ -0,0 +1,43 @@
+//
+// Copyright 2024 The Project Oak Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+use crate::sev::GHCB_WRAPPER;
+
+/// Read a MSR.
+///
+/// ## Safety
+///
+/// The caller must guarantee that the MSR is valid.
+pub unsafe fn read_msr(msr: &crate::hal::base::Msr, msr_id: u32) -> u64 {
+    if let Some(ghcb) = GHCB_WRAPPER.get() {
+        ghcb.lock().msr_read(msr_id).expect("couldn't read the MSR using the GHCB protocol")
+    } else {
+        msr.read()
+    }
+}
+
+/// Write to a MSR.
+///
+/// ## Safety
+///
+/// The caller must guarantee that the MSR is valid.
+pub unsafe fn write_msr(msr: &mut crate::hal::base::Msr, msr_id: u32, val: u64) {
+    if let Some(ghcb) = GHCB_WRAPPER.get() {
+        ghcb.lock().msr_write(msr_id, val).expect("couldn't write the MSR using the GHCB protocol")
+    } else {
+        msr.write(val)
+    }
+}
diff --git a/stage0/src/msr.rs b/stage0/src/msr.rs
index ac49010e353..c54ff8fb712 100644
--- a/stage0/src/msr.rs
+++ b/stage0/src/msr.rs
@@ -16,52 +16,9 @@
 
 use bitflags::bitflags;
 use strum::FromRepr;
-use x86_64::{registers::model_specific::Msr as DirectMsr, PhysAddr};
+use x86_64::PhysAddr;
 
-use crate::sev::GHCB_WRAPPER;
-
-/// Wrapper that can access a MSR either directly or through the GHCB, depending
-/// on the environment.
-pub struct Msr {
-    msr_id: u32,
-    msr: DirectMsr,
-}
-
-impl Msr {
-    pub const fn new(reg: u32) -> Self {
-        Self { msr_id: reg, msr: DirectMsr::new(reg) }
-    }
-
-    /// Read the MSR.
-    ///
-    /// ## Safety
-    ///
-    /// The caller must guarantee that the MSR is valid.
-    pub unsafe fn read(&self) -> u64 {
-        if let Some(ghcb) = GHCB_WRAPPER.get() {
-            ghcb.lock()
-                .msr_read(self.msr_id)
-                .expect("couldn't read the MSR using the GHCB protocol")
-        } else {
-            self.msr.read()
-        }
-    }
-
-    /// Write the MSR.
-    ///
-    /// ## Safety
-    ///
-    /// The caller must guarantee that the MSR is valid.
-    pub unsafe fn write(&mut self, val: u64) {
-        if let Some(ghcb) = GHCB_WRAPPER.get() {
-            ghcb.lock()
-                .msr_write(self.msr_id, val)
-                .expect("couldn't write the MSR using the GHCB protocol")
-        } else {
-            self.msr.write(val)
-        }
-    }
-}
+use crate::hal::Msr;
 
 bitflags! {
     /// Flags in the APIC Base Address Register (MSR 0x1B)