From 8bdc7cbde9406b70bcd9bb385c4b0f3e87ff9fda Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Fri, 6 Sep 2024 07:25:47 -0600
Subject: [PATCH] aarch64-dit: safe, high-level API (#1108)

Removes the previous low-level unsafe API and wraps it up instead in a
safe API which automatically performs runtime CPU feature detection and
RAII guards for enabling DIT and restoring its previous state when done.
---
 aarch64-dit/Cargo.toml |   2 +-
 aarch64-dit/src/lib.rs | 110 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 106 insertions(+), 6 deletions(-)

diff --git a/aarch64-dit/Cargo.toml b/aarch64-dit/Cargo.toml
index d3ad30f2..bba53ad9 100644
--- a/aarch64-dit/Cargo.toml
+++ b/aarch64-dit/Cargo.toml
@@ -14,7 +14,7 @@ readme = "README.md"
 edition = "2021"
 rust-version = "1.61"
 
-[dev-dependencies]
+[dependencies]
 cpufeatures = { version = "0.2.14", path = "../cpufeatures" }
 
 [package.metadata.docs.rs]
diff --git a/aarch64-dit/src/lib.rs b/aarch64-dit/src/lib.rs
index 2239537e..8b5a4b66 100644
--- a/aarch64-dit/src/lib.rs
+++ b/aarch64-dit/src/lib.rs
@@ -6,14 +6,92 @@
 )]
 #![warn(missing_docs, rust_2018_idioms, unused_qualifications)]
 
+//! ## Usage
+//!
+//! ```
+//! use aarch64_dit::Dit;
+//!
+//! let dit = Dit::init();
+//! assert!(!dit.is_enabled());
+//! let _guard = dit.enable();
+//! assert!(dit.is_enabled());
+//! ```
+
 #[cfg(not(target_arch = "aarch64"))]
 compile_error!("This crate only builds on `aarch64` targets");
 
 use core::arch::asm;
 
+cpufeatures::new!(dit_supported, "dit");
+
+/// Data-Independent Timing: support for enabling features of AArch64 CPUs which improve
+/// constant-time operation.
+pub struct Dit {
+    supported: dit_supported::InitToken,
+}
+
+impl Dit {
+    /// Initialize Data-Independent Timing using runtime CPU feature detection.
+    pub fn init() -> Self {
+        Self {
+            supported: dit_supported::init(),
+        }
+    }
+
+    /// Enable Data-Independent Timing (if available).
+    ///
+    /// Returns an RAII guard that will return DIT to its previous state when dropped.
+    #[must_use]
+    pub fn enable(&self) -> Guard<'_> {
+        let was_enabled = if self.is_supported() {
+            let was_enabled = unsafe { get_dit_enabled() };
+            unsafe { set_dit_enabled() };
+            was_enabled
+        } else {
+            false
+        };
+
+        Guard {
+            dit: self,
+            was_enabled,
+        }
+    }
+
+    /// Check if DIT has been enabled.
+    pub fn is_enabled(&self) -> bool {
+        if self.is_supported() {
+            unsafe { get_dit_enabled() }
+        } else {
+            false
+        }
+    }
+
+    /// Check if DIT is supported by this CPU.
+    pub fn is_supported(&self) -> bool {
+        self.supported.get()
+    }
+}
+
+/// RAII guard which returns DIT to its previous state when dropped.
+pub struct Guard<'a> {
+    /// DIT implementation.
+    dit: &'a Dit,
+
+    /// Previous DIT state before it was enabled.
+    was_enabled: bool,
+}
+
+impl Drop for Guard<'_> {
+    fn drop(&mut self) {
+        if self.dit.supported.get() {
+            unsafe { restore_dit(self.was_enabled) }
+        }
+    }
+}
+
 /// Detect if DIT is enabled for the current thread by checking the processor state register.
 #[target_feature(enable = "dit")]
-pub unsafe fn get_dit_enabled() -> bool {
+unsafe fn get_dit_enabled() -> bool {
     let mut dit: u64;
     asm!(
         "mrs {dit}, DIT",
@@ -25,13 +103,13 @@ pub unsafe fn get_dit_enabled() -> bool {
 
 /// Enable DIT for the current thread.
 #[target_feature(enable = "dit")]
-pub unsafe fn set_dit_enabled() {
+unsafe fn set_dit_enabled() {
     asm!("msr DIT, #1", options(nomem, nostack, preserves_flags));
 }
 
 /// Restore DIT state depending on the enabled bit.
 #[target_feature(enable = "dit")]
-pub unsafe fn restore_dit(enabled: bool) {
+unsafe fn restore_dit(enabled: bool) {
     if !enabled {
         // Disable DIT
         asm!("msr DIT, #0", options(nomem, nostack, preserves_flags));
@@ -40,11 +118,33 @@ pub unsafe fn restore_dit(enabled: bool) {
 
 #[cfg(test)]
 mod tests {
-    use super::{get_dit_enabled, restore_dit, set_dit_enabled};
+    use super::{get_dit_enabled, restore_dit, set_dit_enabled, Dit};
     cpufeatures::new!(dit_supported, "dit");
 
     #[test]
-    fn get() {
+    fn high_level_api() {
+        let dit = Dit::init();
+        assert!(dit.is_supported());
+
+        {
+            assert!(!dit.is_enabled());
+            let _guard = dit.enable();
+            assert!(dit.is_enabled());
+
+            // Test nested usage
+            {
+                let _guard2 = dit.enable();
+                assert!(dit.is_enabled());
+            }
+
+            assert!(dit.is_enabled());
+        }
+
+        assert!(!dit.is_enabled());
+    }
+
+    #[test]
+    fn asm_wrappers() {
         let dit_token = dit_supported::init();
         if !dit_token.get() {
             panic!("DIT is not available on this CPU");