fake_mmio/inject/
registry.rs

1// Copyright 2026 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
5use std::collections::BTreeMap;
6use std::sync::{Arc, LazyLock, Mutex};
7
8use crate::inject::passthrough::Passthrough;
9use crate::inject::{GenericVmoMemoryHandler, VmoMemoryHandler, VmoMemoryWrapper, VmoOpHelper};
10use crate::operand::MmioOperand;
11
12/// A trait marking a [`VmoMemoryHandler`] as suitable for use within a
13/// [`Registry`].
14pub trait RegistryHandler: VmoMemoryHandler + Send + Sync + 'static {}
15
16impl<H: VmoMemoryHandler + Send + Sync + 'static> RegistryHandler for H {}
17
18static REGISTRY: LazyLock<Mutex<BTreeMap<zx::Koid, Arc<dyn RegistryHandler>>>> =
19    LazyLock::new(|| Mutex::new(BTreeMap::new()));
20
21/// A VMO handler registry that retrieves registered handlers by VMO KOID.
22///
23/// If `STRICT` is true, creating a handler for an unregistered VMO panics.
24/// Otherwise, it falls back to [`Passthrough`].
25///
26/// Use [`BaseRegistry::register_handler`] or
27/// [`ScopedRegistry::register_handler`] to register [`RegistryHandler`]
28/// implementations that can be retrieved at-a-distance based on the VMO KOID.
29///
30/// Use this handler implementation in testing situations where the MMIO VMO
31/// might come over FIDL from an injected dependency. Register the VMO in a
32/// `Registry` before returning the VMO in an injected dependency fake.
33pub struct BaseRegistry<const STRICT: bool>(Arc<dyn RegistryHandler>);
34
35/// A registry that falls back to passthrough behavior for unregistered VMOs.
36pub type Registry = BaseRegistry<false>;
37
38/// A registry that panics when encountering an unregistered VMO.
39pub type StrictRegistry = BaseRegistry<true>;
40
41impl<const STRICT: bool> BaseRegistry<STRICT> {
42    /// Registers `handler` for `vmo`.
43    ///
44    /// The handler is installed in the global registry and never removed. See
45    /// [`ScopedRegistry`] for alternatives.
46    pub fn register_handler<H: RegistryHandler>(vmo: &zx::Vmo, handler: H) {
47        Self::register_shared_handler(vmo, Arc::new(handler))
48    }
49
50    /// Registers a shared reference-counted `handler` for `vmo`.
51    ///
52    /// The handler is installed in the global registry and never removed. See
53    /// [`ScopedRegistry`] for alternatives.
54    pub fn register_shared_handler(vmo: &zx::Vmo, handler: Arc<dyn RegistryHandler>) {
55        let _: zx::Koid = register_global_handler(vmo, handler);
56    }
57
58    /// Clears all the registered VMOs.
59    pub fn clear() {
60        REGISTRY.lock().unwrap().clear();
61    }
62}
63
64impl<const STRICT: bool> VmoMemoryWrapper for BaseRegistry<STRICT> {
65    fn new(vmo: &zx::Vmo) -> Self {
66        let koid = vmo.koid().expect("getting koid");
67        let handler = REGISTRY.lock().unwrap().get(&koid).cloned();
68        let handler = handler.unwrap_or_else(|| {
69            if STRICT {
70                panic!("no handler for vmo {koid:?}")
71            }
72            Arc::new(Passthrough)
73        });
74        Self(handler)
75    }
76}
77
78impl<const STRICT: bool> GenericVmoMemoryHandler for BaseRegistry<STRICT> {
79    fn load<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>) -> T {
80        T::load(&*self.0, op)
81    }
82
83    fn store<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>, value: T) {
84        T::store(&*self.0, op, value)
85    }
86}
87
88fn register_global_handler(vmo: &zx::Vmo, handler: Arc<dyn RegistryHandler>) -> zx::Koid {
89    let koid = vmo.koid().expect("getting koid");
90    assert!(REGISTRY.lock().unwrap().insert(koid, handler).is_none(), "already registered");
91    koid
92}
93
94/// A helper to add VMOs to a global registry that removes all installed VMOs on
95/// drop.
96#[derive(Default, Clone, Debug)]
97pub struct ScopedRegistry(Vec<zx::Koid>);
98
99impl ScopedRegistry {
100    /// Creates a new registry for registering VMOs.
101    pub fn new() -> Self {
102        Self::default()
103    }
104
105    /// Registers `handler` for `vmo`.
106    pub fn register_handler<H: RegistryHandler>(&mut self, vmo: &zx::Vmo, handler: H) {
107        self.register_shared_handler(vmo, Arc::new(handler))
108    }
109
110    /// Registers a shared reference-counted `handler` for `vmo`.
111    pub fn register_shared_handler(&mut self, vmo: &zx::Vmo, handler: Arc<dyn RegistryHandler>) {
112        self.0.push(register_global_handler(vmo, handler));
113    }
114
115    /// Drops this `ScopedRegistry` without unregistering the VMOs.
116    pub fn detach(mut self) {
117        self.0.clear();
118    }
119}
120
121impl Drop for ScopedRegistry {
122    fn drop(&mut self) {
123        let mut registry = REGISTRY.lock().unwrap();
124        for koid in std::mem::take(&mut self.0) {
125            let _ = registry.remove(&koid);
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use crate::inject::{MockVmoMemoryHandler, VmoMemory};
134    use mmio::MmioExt as _;
135
136    const VMO_SIZE: usize = 4096;
137
138    #[test]
139    fn test_registry_registered_vmo() {
140        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
141        let expected_val = 0x1234_5678_u32;
142        let mut mock = MockVmoMemoryHandler::new();
143        let _ = mock
144            .expect_load32()
145            .once()
146            .withf(|op| op.offset() == 0 && op.range() == (0..4))
147            .return_const(expected_val);
148        // Scoped registry ensures we're going to run drop on `mock`.
149        let mut registry = ScopedRegistry::new();
150        registry.register_handler(&vmo, mock);
151
152        let mmio = VmoMemory::<Registry>::map(0, VMO_SIZE, vmo).expect("Failed to map VMO.");
153        assert_eq!(mmio.load::<u32>(0), expected_val);
154    }
155
156    #[test]
157    fn test_registry_unregistered_vmo_fallback() {
158        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
159
160        let mut mmio = VmoMemory::<Registry>::map(0, VMO_SIZE, vmo).expect("Failed to map VMO.");
161        let val = 0x8765_4321_u32;
162        mmio.store::<u32>(0, val);
163        assert_eq!(mmio.load::<u32>(0), val);
164    }
165
166    #[test]
167    fn test_strict_registry_registered_vmo() {
168        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
169        let expected_val = 0xabcd_ef01_u32;
170        let mut mock = MockVmoMemoryHandler::new();
171        let _ = mock
172            .expect_load32()
173            .once()
174            .withf(|op| op.offset() == 0 && op.range() == (0..4))
175            .return_const(expected_val);
176        // Scoped registry ensures we're going to run drop on `mock`.
177        let mut registry = ScopedRegistry::new();
178        registry.register_handler(&vmo, mock);
179
180        let mmio = VmoMemory::<StrictRegistry>::map(0, VMO_SIZE, vmo).expect("Failed to map VMO.");
181        assert_eq!(mmio.load::<u32>(0), expected_val);
182    }
183
184    #[fuchsia::test(logging = false)]
185    #[should_panic(expected = "already registered")]
186    fn test_registry_already_registered_panics() {
187        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
188        Registry::register_handler(&vmo, MockVmoMemoryHandler::new());
189        Registry::register_handler(&vmo, MockVmoMemoryHandler::new());
190    }
191
192    #[fuchsia::test(logging = false)]
193    #[should_panic(expected = "no handler for vmo")]
194    fn test_strict_registry_unregistered_vmo_panics() {
195        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
196        let _ = VmoMemory::<StrictRegistry>::map(0, VMO_SIZE, vmo);
197    }
198
199    #[test]
200    fn test_scoped_registry_drop() {
201        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
202
203        {
204            let expected_val = 0xabcd_ef01_u32;
205            let mut mock = MockVmoMemoryHandler::new();
206            let _ = mock
207                .expect_load32()
208                .once()
209                .withf(|op| op.offset() == 0 && op.range() == (0..4))
210                .return_const(expected_val);
211            let mut scoped = ScopedRegistry::new();
212            scoped.register_handler(&vmo, mock);
213
214            let mmio = VmoMemory::<StrictRegistry>::map(
215                0,
216                VMO_SIZE,
217                vmo.duplicate_handle(zx::Rights::SAME_RIGHTS).unwrap(),
218            )
219            .expect("Failed to map VMO.");
220            assert_eq!(mmio.load::<u32>(0), expected_val);
221        }
222
223        let mmio = VmoMemory::<Registry>::map(0, VMO_SIZE, vmo).expect("Failed to map VMO.");
224        assert_eq!(mmio.load::<u32>(0), 0);
225    }
226
227    #[test]
228    fn test_scoped_registry_detach() {
229        let vmo = zx::Vmo::create(VMO_SIZE.try_into().unwrap()).expect("Failed to create VMO.");
230        let expected_val = 0xabcd_ef01_u32;
231
232        let mut mock = MockVmoMemoryHandler::new();
233        let _ = mock
234            .expect_load32()
235            .times(2)
236            .withf(|op| op.offset() == 0 && op.range() == (0..4))
237            .return_const(expected_val);
238        let mut scoped = ScopedRegistry::new();
239        scoped.register_handler(&vmo, mock);
240
241        let mmio = VmoMemory::<StrictRegistry>::map(
242            0,
243            VMO_SIZE,
244            vmo.duplicate_handle(zx::Rights::SAME_RIGHTS).unwrap(),
245        )
246        .expect("Failed to map VMO.");
247        assert_eq!(mmio.load::<u32>(0), expected_val);
248
249        scoped.detach();
250
251        let mmio = VmoMemory::<StrictRegistry>::map(0, VMO_SIZE, vmo).expect("Failed to map VMO.");
252        assert_eq!(mmio.load::<u32>(0), expected_val);
253
254        // Clear the registry to run drop on mock.
255        Registry::clear();
256    }
257}
fake_mmio/inject/registry.rs

fake_mmio/inject/
registry.rs
