Skip to main content

starnix_registers/
lib.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#[cfg(target_arch = "aarch64")]
6mod arm64;
7
8#[cfg(target_arch = "aarch64")]
9pub use arm64::*;
10
11#[cfg(target_arch = "x86_64")]
12mod x64;
13
14#[cfg(target_arch = "x86_64")]
15pub use x64::*;
16
17#[cfg(target_arch = "riscv64")]
18mod riscv64;
19
20#[cfg(target_arch = "riscv64")]
21pub use riscv64::*;
22
23use starnix_logging::{CATEGORY_STARNIX, NAME_MAP_RESTRICTED_STATE};
24use starnix_uapi::__static_assertions::assert_not_impl_any;
25use std::ops::Deref;
26use std::ptr::NonNull;
27use std::sync::LazyLock;
28
29/// Size of `RESTRICTED_STATE_VMAR`.
30///
31/// With single-page state VMOs this fits 65536 concurrently bound restricted states, which covers
32/// the default Linux `pid_max` (32768). Threads beyond that capacity fall back to mapping their
33/// state into the root VMAR.
34const RESTRICTED_STATE_VMAR_SIZE: usize = 256 * 1024 * 1024;
35
36/// A dedicated `ZX_VM_COMPACT` sub-VMAR for restricted state mappings.
37///
38/// Mapping the per-thread state VMOs into a compact sub-VMAR, rather than randomly placing each
39/// one in the root VMAR, keeps `zx_vmar_map` cheap and keeps the mappings densely packed in
40/// shared page tables. The sub-VMAR itself is randomly placed, but the mappings within it are
41/// packed next to each other without gaps.
42///
43/// This is `None` if the sub-VMAR could not be allocated, in which case all mappings go to the
44/// root VMAR.
45static RESTRICTED_STATE_VMAR: LazyLock<Option<zx::Vmar>> = LazyLock::new(|| {
46    fuchsia_runtime::vmar_root_self()
47        .allocate(
48            0,
49            RESTRICTED_STATE_VMAR_SIZE,
50            zx::VmarFlags::CAN_MAP_READ | zx::VmarFlags::CAN_MAP_WRITE | zx::VmarFlags::COMPACT,
51        )
52        .map(|(vmar, _base)| vmar)
53        .ok()
54});
55
56/// Maps `state_vmo` for read/write access into `RESTRICTED_STATE_VMAR`, falling back to the root
57/// VMAR.
58///
59/// Returns the mapped address and whether the mapping is in `RESTRICTED_STATE_VMAR` (as opposed
60/// to the root VMAR).
61fn map_state_vmo(state_vmo: &zx::Vmo, size: usize) -> Result<(usize, bool), zx::Status> {
62    map_state_vmo_in(RESTRICTED_STATE_VMAR.as_ref(), state_vmo, size)
63}
64
65/// Maps `state_vmo` for read/write access into `compact_vmar`, or into the root VMAR if
66/// `compact_vmar` is `None` or full.
67///
68/// Returns the mapped address and whether the mapping is in `compact_vmar`.
69fn map_state_vmo_in(
70    compact_vmar: Option<&zx::Vmar>,
71    state_vmo: &zx::Vmo,
72    size: usize,
73) -> Result<(usize, bool), zx::Status> {
74    // `zx_restricted_bind_state` commits and pins the state VMO's pages, so `MAP_RANGE` populates
75    // the page table entries during the map and the first access below does not page fault.
76    let flags = zx::VmarFlags::PERM_READ | zx::VmarFlags::PERM_WRITE | zx::VmarFlags::MAP_RANGE;
77    if let Some(vmar) = compact_vmar {
78        match vmar.map(0, state_vmo, 0, size, flags) {
79            Ok(address) => return Ok((address, true)),
80            // The compact VMAR is full; fall back to the root VMAR.
81            Err(zx::Status::NO_RESOURCES) => {}
82            Err(status) => return Err(status),
83        }
84    }
85    let address = fuchsia_runtime::vmar_root_self().map(0, state_vmo, 0, size, flags)?;
86    Ok((address, false))
87}
88
89/// Unmaps a mapping created by `map_state_vmo_in` with the same `compact_vmar`.
90///
91/// # Safety
92///
93/// The mapping at `address` must not be accessed after this call.
94unsafe fn unmap_state_vmo_in(
95    compact_vmar: Option<&zx::Vmar>,
96    in_compact_vmar: bool,
97    address: usize,
98    size: usize,
99) -> Result<(), zx::Status> {
100    if in_compact_vmar {
101        let vmar = compact_vmar.expect("the compact VMAR must exist if it holds a mapping");
102        // SAFETY: The caller guarantees that the mapping is no longer accessed.
103        unsafe { vmar.unmap(address, size) }
104    } else {
105        // SAFETY: The caller guarantees that the mapping is no longer accessed.
106        unsafe { fuchsia_runtime::vmar_root_self().unmap(address, size) }
107    }
108}
109
110/// `RestrictedState` manages accesses into the restricted state VMO.
111///
112/// See `zx_restricted_bind_state`.
113pub struct RestrictedState {
114    pub bound_state: NonNull<zx::sys::zx_restricted_state_t>,
115    state_size: usize,
116    /// Whether `bound_state` is mapped in `RESTRICTED_STATE_VMAR` rather than the root VMAR.
117    in_compact_vmar: bool,
118}
119
120impl RestrictedState {
121    /// Allocates a VMO for the restricted state and maps it into the Starnix kernel.
122    ///
123    /// The VMO is created using `zx_restricted_bind_state` and then mapped into the
124    /// kernel's address space. This allows Starnix to inspect and modify the
125    /// restricted state (registers) of the task.
126    pub fn bind_and_map(
127        register_state: &mut RegisterState<RegisterStorageEnum>,
128        exception_report: &mut zx::sys::zx_exception_report_t,
129    ) -> Result<Self, zx::Status> {
130        fuchsia_trace::duration!(CATEGORY_STARNIX, NAME_MAP_RESTRICTED_STATE);
131        let mut out_vmo_handle = 0;
132        // SAFETY: `out_vmo_handle` is a valid pointer to a handle on the stack.
133        let status = zx::Status::ok(unsafe {
134            zx::sys::zx_restricted_bind_state(
135                0,
136                &mut out_vmo_handle,
137                std::ptr::from_mut(exception_report),
138            )
139        });
140        match status {
141            Ok(()) => {
142                // We've successfully attached the VMO to the current thread. This VMO will be
143                // mapped and used for the kernel to store restricted mode register state as it
144                // enters and exits restricted mode.
145            }
146            Err(status) => panic!("zx_restricted_bind_state failed with {status}!"),
147        }
148        // SAFETY: `out_vmo_handle` is a valid handle as `zx_restricted_bind_state` returned OK.
149        let state_vmo = unsafe { zx::Vmo::from(zx::NullableHandle::from_raw(out_vmo_handle)) };
150
151        // Name the VMO so external tools (like the CPU profiler) can capture
152        // this specific VMO.
153        let name = format!(
154            "restricted_state_vmo:{}",
155            fuchsia_runtime::with_thread_self(|t| t.koid())?.raw_koid()
156        );
157        let name = zx::Name::new_lossy(&name);
158        state_vmo.set_name(&name)?;
159
160        let state_size = state_vmo.get_size()? as usize;
161        if state_size < std::mem::size_of::<zx::sys::zx_restricted_exception_t>() {
162            return Err(zx::Status::INVALID_ARGS);
163        }
164
165        let (state_address, in_compact_vmar) = map_state_vmo(&state_vmo, state_size)?;
166
167        // This memory is not managed by Rust's stack, heap, etc. so treat it as "foreign" memory
168        // with no provenance.
169        let state_address: *mut zx::sys::zx_restricted_state_t =
170            std::ptr::without_provenance_mut(state_address);
171        assert!(state_address.is_aligned(), "Zircon must map restricted-state-aligned memory");
172        let bound_state =
173            NonNull::new(state_address).expect("Zircon must map non-null restricted-state");
174
175        // Copy the initial register state into the mapped VMO and link the VMO to the register
176        // state of the current thread.
177        // SAFETY: `bound_state` is valid to read/write to as long as `RestrictedState` is live.
178        unsafe {
179            let vmo_ptr = bound_state.as_ptr();
180            vmo_ptr.write(**register_state);
181            register_state.real_registers =
182                RegisterStorageEnum::Vmo(MappedVmoRegs(NonNull::new_unchecked(vmo_ptr)));
183        }
184
185        Ok(Self { state_size, bound_state, in_compact_vmar })
186    }
187}
188
189impl std::ops::Drop for RestrictedState {
190    fn drop(&mut self) {
191        let mapping_addr = self.bound_state.as_ptr() as usize;
192        // Safety: We are un-mapping the state VMO. This is safe because we route all access
193        // into this memory region though this struct so it is safe to unmap on Drop.
194        unsafe {
195            unmap_state_vmo_in(
196                RESTRICTED_STATE_VMAR.as_ref(),
197                self.in_compact_vmar,
198                mapping_addr,
199                self.state_size,
200            )
201            .expect("Failed to unmap");
202            zx::sys::zx_restricted_unbind_state(0);
203        }
204    }
205}
206
207pub trait RegisterStorage:
208    std::ops::Deref<Target = zx::sys::zx_restricted_state_t>
209    + std::ops::DerefMut
210    + Eq
211    + PartialEq
212    + std::fmt::Debug
213    + Clone
214{
215}
216
217#[derive(Eq, PartialEq, Clone)]
218struct MappedVmoRegs(NonNull<zx::sys::zx_restricted_state_t>);
219// MappedVmoRegs should be tied to the CurrentTask, so it is not Send or Sync.
220assert_not_impl_any!(MappedVmoRegs: Send, Sync);
221
222impl std::fmt::Debug for MappedVmoRegs {
223    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224        f.debug_tuple("MappedVmoRegs").field(&format_args!("{:?}", self.deref())).finish()
225    }
226}
227
228impl std::ops::Deref for MappedVmoRegs {
229    type Target = zx::sys::zx_restricted_state_t;
230
231    fn deref(&self) -> &Self::Target {
232        // SAFETY: The pointer is valid and points to a valid `zx_restricted_state_t`.
233        unsafe { self.0.as_ref() }
234    }
235}
236
237impl std::ops::DerefMut for MappedVmoRegs {
238    fn deref_mut(&mut self) -> &mut Self::Target {
239        // SAFETY: The pointer is valid and points to a valid `zx_restricted_state_t`.
240        unsafe { self.0.as_mut() }
241    }
242}
243
244impl RegisterStorage for MappedVmoRegs {}
245
246#[derive(Eq, PartialEq, Debug, Clone, Default)]
247pub struct HeapRegs(Box<zx::sys::zx_restricted_state_t>);
248
249impl std::ops::Deref for HeapRegs {
250    type Target = zx::sys::zx_restricted_state_t;
251
252    fn deref(&self) -> &Self::Target {
253        &self.0
254    }
255}
256
257impl std::ops::DerefMut for HeapRegs {
258    fn deref_mut(&mut self) -> &mut Self::Target {
259        &mut self.0
260    }
261}
262
263impl RegisterStorage for HeapRegs {}
264
265impl From<RegisterStorageEnum> for HeapRegs {
266    fn from(regs: RegisterStorageEnum) -> Self {
267        match regs {
268            RegisterStorageEnum::Vmo(vmo) => HeapRegs(Box::new(*vmo)),
269            RegisterStorageEnum::Heap(heap) => heap,
270        }
271    }
272}
273
274/// An enum to hold the registers in either a heap or a vmo.
275///
276/// This is introduced to allow `CurrentTask` to store registers in heap during initialization and
277/// link to the vmo after the `RestrictedState` is created.
278#[derive(Eq, PartialEq, Debug)]
279pub enum RegisterStorageEnum {
280    // Keep it private to prevent from using it directly.
281    #[allow(private_interfaces)]
282    Vmo(MappedVmoRegs),
283    Heap(HeapRegs),
284}
285assert_not_impl_any!(RegisterStorageEnum: Send, Sync);
286
287impl std::ops::Deref for RegisterStorageEnum {
288    type Target = zx::sys::zx_restricted_state_t;
289
290    fn deref(&self) -> &Self::Target {
291        match self {
292            RegisterStorageEnum::Vmo(vmo) => vmo.deref(),
293            RegisterStorageEnum::Heap(heap) => heap.deref(),
294        }
295    }
296}
297
298impl std::ops::DerefMut for RegisterStorageEnum {
299    fn deref_mut(&mut self) -> &mut Self::Target {
300        match self {
301            RegisterStorageEnum::Vmo(vmo) => vmo.deref_mut(),
302            RegisterStorageEnum::Heap(heap) => heap.deref_mut(),
303        }
304    }
305}
306
307impl Clone for RegisterStorageEnum {
308    fn clone(&self) -> RegisterStorageEnum {
309        RegisterStorageEnum::Heap(HeapRegs(Box::new(**self)))
310    }
311}
312
313impl RegisterStorage for RegisterStorageEnum {}
314
315impl From<MappedVmoRegs> for RegisterStorageEnum {
316    fn from(regs: MappedVmoRegs) -> Self {
317        RegisterStorageEnum::Vmo(regs)
318    }
319}
320
321impl From<HeapRegs> for RegisterStorageEnum {
322    fn from(regs: HeapRegs) -> Self {
323        RegisterStorageEnum::Heap(regs)
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330
331    fn page_size() -> usize {
332        zx::system_get_page_size() as usize
333    }
334
335    /// Creates a one-page VMO with its page committed, like the VMO returned by
336    /// `zx_restricted_bind_state`.
337    fn create_state_vmo() -> zx::Vmo {
338        let vmo = zx::Vmo::create(page_size() as u64).expect("failed to create VMO");
339        vmo.op_range(zx::VmoOp::COMMIT, 0, page_size() as u64).expect("failed to commit VMO");
340        vmo
341    }
342
343    /// Allocates a compact sub-VMAR with room for exactly one state mapping.
344    fn create_one_page_compact_vmar() -> zx::Vmar {
345        let (vmar, _base) = fuchsia_runtime::vmar_root_self()
346            .allocate(
347                0,
348                page_size(),
349                zx::VmarFlags::CAN_MAP_READ | zx::VmarFlags::CAN_MAP_WRITE | zx::VmarFlags::COMPACT,
350            )
351            .expect("failed to allocate compact VMAR");
352        vmar
353    }
354
355    fn vmar_contains(vmar: &zx::Vmar, address: usize) -> bool {
356        let info = vmar.info().expect("failed to get VMAR info");
357        (info.base..info.base + info.len).contains(&address)
358    }
359
360    /// Checks that the mapping at `address` is backed by `vmo` and writable.
361    fn check_mapping(vmo: &zx::Vmo, address: usize) {
362        let ptr = std::ptr::with_exposed_provenance_mut::<u64>(address);
363        // SAFETY: `address` is the start of a live, writable mapping of at least one page.
364        unsafe { ptr.write_volatile(0x0123_4567_89ab_cdef) };
365        let mut value = [0u8; 8];
366        vmo.read(&mut value, 0).expect("failed to read VMO");
367        assert_eq!(u64::from_ne_bytes(value), 0x0123_4567_89ab_cdef);
368    }
369
370    #[test]
371    fn map_state_vmo_uses_restricted_state_vmar() {
372        assert!(RESTRICTED_STATE_VMAR.is_some(), "failed to allocate RESTRICTED_STATE_VMAR");
373        let vmo = create_state_vmo();
374        let (address, in_compact_vmar) =
375            map_state_vmo(&vmo, page_size()).expect("failed to map state VMO");
376        assert!(in_compact_vmar);
377        assert!(vmar_contains(RESTRICTED_STATE_VMAR.as_ref().unwrap(), address));
378        check_mapping(&vmo, address);
379        // SAFETY: The mapping is not accessed after this.
380        unsafe {
381            unmap_state_vmo_in(
382                RESTRICTED_STATE_VMAR.as_ref(),
383                in_compact_vmar,
384                address,
385                page_size(),
386            )
387        }
388        .expect("failed to unmap state VMO");
389    }
390
391    #[test]
392    fn map_state_vmo_in_compact_vmar_and_unmap() {
393        let compact_vmar = create_one_page_compact_vmar();
394        let vmo = create_state_vmo();
395        let (address, in_compact_vmar) =
396            map_state_vmo_in(Some(&compact_vmar), &vmo, page_size()).expect("failed to map");
397        assert!(in_compact_vmar);
398        assert!(vmar_contains(&compact_vmar, address));
399        check_mapping(&vmo, address);
400        // SAFETY: The mapping is not accessed after this.
401        unsafe { unmap_state_vmo_in(Some(&compact_vmar), in_compact_vmar, address, page_size()) }
402            .expect("failed to unmap");
403
404        // Unmapping freed the only slot in the compact VMAR, so mapping again lands there.
405        let (address, in_compact_vmar) =
406            map_state_vmo_in(Some(&compact_vmar), &vmo, page_size()).expect("failed to remap");
407        assert!(in_compact_vmar);
408        // SAFETY: The mapping is not accessed after this.
409        unsafe { unmap_state_vmo_in(Some(&compact_vmar), in_compact_vmar, address, page_size()) }
410            .expect("failed to unmap");
411    }
412
413    #[test]
414    fn map_state_vmo_in_falls_back_to_root_vmar_when_compact_vmar_is_full() {
415        let compact_vmar = create_one_page_compact_vmar();
416        let first_vmo = create_state_vmo();
417        let second_vmo = create_state_vmo();
418
419        let (first_address, first_in_compact_vmar) =
420            map_state_vmo_in(Some(&compact_vmar), &first_vmo, page_size()).expect("failed to map");
421        assert!(first_in_compact_vmar);
422
423        // The compact VMAR is full, so `map` fails with `ZX_ERR_NO_RESOURCES` and the second
424        // mapping falls back to the root VMAR.
425        let (second_address, second_in_compact_vmar) =
426            map_state_vmo_in(Some(&compact_vmar), &second_vmo, page_size()).expect("failed to map");
427        assert!(!second_in_compact_vmar);
428        assert!(!vmar_contains(&compact_vmar, second_address));
429        check_mapping(&second_vmo, second_address);
430
431        // SAFETY: Neither mapping is accessed after this.
432        unsafe {
433            unmap_state_vmo_in(
434                Some(&compact_vmar),
435                second_in_compact_vmar,
436                second_address,
437                page_size(),
438            )
439            .expect("failed to unmap from the root VMAR");
440            unmap_state_vmo_in(
441                Some(&compact_vmar),
442                first_in_compact_vmar,
443                first_address,
444                page_size(),
445            )
446            .expect("failed to unmap from the compact VMAR");
447        }
448    }
449
450    #[test]
451    fn map_state_vmo_in_without_compact_vmar_uses_root_vmar() {
452        let vmo = create_state_vmo();
453        let (address, in_compact_vmar) =
454            map_state_vmo_in(None, &vmo, page_size()).expect("failed to map");
455        assert!(!in_compact_vmar);
456        check_mapping(&vmo, address);
457        // SAFETY: The mapping is not accessed after this.
458        unsafe { unmap_state_vmo_in(None, in_compact_vmar, address, page_size()) }
459            .expect("failed to unmap");
460    }
461}