aarch64_dit/
lib.rs

1#![no_std]
2#![doc = include_str!("../README.md")]
3#![doc(
4    html_logo_url = "https://raw.githubusercontent.com/RustCrypto/meta/master/logo.svg",
5    html_favicon_url = "https://raw.githubusercontent.com/RustCrypto/meta/master/logo.svg"
6)]
7#![warn(missing_docs, rust_2018_idioms, unused_qualifications)]
8
9//! ## Usage
10//!
11//! ```
12//! use aarch64_dit::Dit;
13//!
14//! let dit = Dit::init();
15//! assert!(!dit.is_enabled());
16//! let _guard = dit.enable();
17//! assert!(dit.is_enabled());
18//! ```
19
20#[cfg(not(target_arch = "aarch64"))]
21compile_error!("This crate only builds on `aarch64` targets");
22
23use core::arch::asm;
24
25cpufeatures::new!(dit_supported, "dit");
26
27/// Data-Independent Timing: support for enabling features of AArch64 CPUs which improve
28/// constant-time operation.
29pub struct Dit {
30    supported: dit_supported::InitToken,
31}
32
33impl Dit {
34    /// Initialize Data-Independent Timing using runtime CPU feature detection.
35    pub fn init() -> Self {
36        Self {
37            supported: dit_supported::init(),
38        }
39    }
40
41    /// Enable Data-Independent Timing (if available).
42    ///
43    /// Returns an RAII guard that will return DIT to its previous state when dropped.
44    #[must_use]
45    pub fn enable(&self) -> Guard<'_> {
46        let was_enabled = if self.is_supported() {
47            unsafe { set_dit_enabled() }
48        } else {
49            false
50        };
51
52        Guard {
53            dit: self,
54            was_enabled,
55        }
56    }
57
58    /// Check if DIT has been enabled.
59    pub fn is_enabled(&self) -> bool {
60        if self.is_supported() {
61            unsafe { get_dit_enabled() }
62        } else {
63            false
64        }
65    }
66
67    /// Check if DIT is supported by this CPU.
68    pub fn is_supported(&self) -> bool {
69        self.supported.get()
70    }
71}
72
73/// RAII guard which returns DIT to its previous state when dropped.
74pub struct Guard<'a> {
75    /// DIT implementation.
76    dit: &'a Dit,
77
78    /// Previous DIT state before it was enabled.
79    was_enabled: bool,
80}
81
82impl Drop for Guard<'_> {
83    fn drop(&mut self) {
84        if self.dit.supported.get() {
85            unsafe { restore_dit(self.was_enabled) }
86        }
87    }
88}
89
90/// Detect if DIT is enabled for the current thread by checking the processor state register.
91#[target_feature(enable = "dit")]
92unsafe fn get_dit_enabled() -> bool {
93    let mut dit: u64;
94    asm!(
95        "mrs {dit}, DIT",
96        dit = out(reg) dit,
97        options(nomem, nostack, preserves_flags)
98    );
99    (dit >> 24) & 1 != 0
100}
101
102/// Enable DIT for the current thread.
103///
104/// Returns the previous DIT state prior to enabling DIT.
105#[target_feature(enable = "dit")]
106unsafe fn set_dit_enabled() -> bool {
107    let was_enabled = get_dit_enabled();
108    asm!("msr DIT, #1", options(nomem, nostack, preserves_flags));
109    was_enabled
110}
111
112/// Restore DIT state depending on the enabled bit.
113#[target_feature(enable = "dit")]
114unsafe fn restore_dit(enabled: bool) {
115    if !enabled {
116        // Disable DIT
117        asm!("msr DIT, #0", options(nomem, nostack, preserves_flags));
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::{get_dit_enabled, restore_dit, set_dit_enabled, Dit};
124    cpufeatures::new!(dit_supported, "dit");
125
126    #[test]
127    fn high_level_api() {
128        let dit = Dit::init();
129        assert!(dit.is_supported());
130
131        {
132            assert!(!dit.is_enabled());
133            let _guard = dit.enable();
134            assert!(dit.is_enabled());
135
136            // Test nested usage
137            {
138                let _guard2 = dit.enable();
139                assert!(dit.is_enabled());
140            }
141
142            assert!(dit.is_enabled());
143        }
144
145        assert!(!dit.is_enabled());
146    }
147
148    #[test]
149    fn asm_wrappers() {
150        let dit_token = dit_supported::init();
151        if !dit_token.get() {
152            panic!("DIT is not available on this CPU");
153        }
154
155        let dit_enabled = unsafe { get_dit_enabled() };
156        assert!(!dit_enabled);
157
158        let was_enabled = unsafe { set_dit_enabled() };
159        assert!(!was_enabled);
160        let dit_enabled = unsafe { get_dit_enabled() };
161        assert!(dit_enabled);
162
163        unsafe { restore_dit(true) };
164        let dit_enabled = unsafe { get_dit_enabled() };
165        assert!(dit_enabled);
166
167        unsafe { restore_dit(false) };
168        let dit_enabled = unsafe { get_dit_enabled() };
169        assert!(!dit_enabled);
170    }
171}
aarch64_dit/lib.rs

aarch64_dit/
lib.rs
