lab_proxy_core/
state_bank.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
5//! Software state banks, runtime knobs, and diagnostic triggers for in-situ
6//! driver experimentation.
7//!
8//! [`StateBank`] implements [`MmioBackend`] and [`ResourceBackend`] over a
9//! 32-bit word-addressed virtual bank of 4-byte aligned offsets. A driver can
10//! register:
11//! - **Observable state slots** ([`StateBank::define_state_slot`]): lock-free
12//!   [`AtomicU32`] slots readable by the host (via `read32`, `poll32`, or
13//!   `snapshot`) and updated by the driver via [`StateSlotHandle`].
14//! - **Runtime knobs** ([`StateBank::define_knob`]): writable [`AtomicU32`]
15//!   slots (included in `writable_registers`) for fault/race-window injection
16//!   or runtime fix-toggle experiments.
17//! - **Read probes** ([`StateBank::define_read_probe`]): synchronous callbacks
18//!   evaluated on demand when the host reads the offset.
19//! - **Diagnostic triggers** ([`StateBank::define_trigger`]): writable offsets
20//!   backed by a bounded callback `Fn(u32) -> Result<u32, BackendError>` whose
21//!   returned status word is stored for readback and subsequent polling.
22
23use crate::access_policy::{ResourceCeiling, WritableRegister};
24use crate::hardware_backend::{BackendError, MmioBackend};
25use crate::protocol_resource_adapter::ResourceBackend;
26use std::collections::BTreeMap;
27use std::sync::Arc;
28use std::sync::atomic::{AtomicU32, Ordering, fence};
29
30const WORD_SIZE: u64 = 4;
31const FULL_WRITE_MASK: u32 = 0xFFFF_FFFF;
32
33/// Lock-free handle to a 32-bit state slot or runtime knob shared between an
34/// instrumented driver and a [`StateBank`].
35#[derive(Clone, Debug)]
36pub struct StateSlotHandle {
37    inner: Arc<AtomicU32>,
38}
39
40impl StateSlotHandle {
41    /// Creates a new slot handle initialized to `initial`.
42    pub fn new(initial: u32) -> Self {
43        Self { inner: Arc::new(AtomicU32::new(initial)) }
44    }
45
46    /// Loads the current 32-bit value with sequentially consistent ordering.
47    pub fn load(&self) -> u32 {
48        self.inner.load(Ordering::SeqCst)
49    }
50
51    /// Stores a new 32-bit value with sequentially consistent ordering.
52    pub fn store(&self, value: u32) {
53        self.inner.store(value, Ordering::SeqCst);
54    }
55
56    /// Atomically adds `delta` to the current value, returning the previous value.
57    pub fn fetch_add(&self, delta: u32) -> u32 {
58        self.inner.fetch_add(delta, Ordering::SeqCst)
59    }
60
61    /// Atomically bitwise-ORs `mask` into the current value, returning the previous value.
62    pub fn fetch_or(&self, mask: u32) -> u32 {
63        self.inner.fetch_or(mask, Ordering::SeqCst)
64    }
65
66    /// Atomically bitwise-ANDs `mask` into the current value, returning the previous value.
67    pub fn fetch_and(&self, mask: u32) -> u32 {
68        self.inner.fetch_and(mask, Ordering::SeqCst)
69    }
70
71    /// Returns a reference to the underlying [`AtomicU32`].
72    pub fn as_atomic(&self) -> &AtomicU32 {
73        &self.inner
74    }
75}
76
77impl Default for StateSlotHandle {
78    fn default() -> Self {
79        Self::new(0)
80    }
81}
82
83/// Callback type for on-demand 32-bit state read probes.
84pub type ReadProbeFn = dyn Fn() -> Result<u32, BackendError> + Send + Sync;
85
86/// Callback type for write-triggered diagnostic actions.
87pub type TriggerFn = dyn Fn(u32) -> Result<u32, BackendError> + Send + Sync;
88
89/// Entry bound to a 4-byte aligned offset inside a [`StateBank`].
90#[derive(Clone)]
91pub enum StateEntry {
92    /// Lock-free atomic 32-bit slot (read-only state slot or writable runtime knob).
93    AtomicSlot(Arc<AtomicU32>),
94    /// Synchronous read callback computed on demand during `read32`.
95    ReadProbe(Arc<ReadProbeFn>),
96    /// Write-triggered callback that stores its returned 32-bit status code into
97    /// `last_result` for immediate readback or subsequent `read32`/`poll32`.
98    Trigger {
99        /// Most recent result returned by `callback` (initialized to `0`).
100        last_result: Arc<AtomicU32>,
101        /// Bounded callback executed on `write32`.
102        callback: Arc<TriggerFn>,
103    },
104}
105
106impl std::fmt::Debug for StateEntry {
107    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        match self {
109            Self::AtomicSlot(slot) => {
110                f.debug_tuple("AtomicSlot").field(&slot.load(Ordering::SeqCst)).finish()
111            }
112            Self::ReadProbe(_) => f.debug_struct("ReadProbe").field("callback", &"<fn>").finish(),
113            Self::Trigger { last_result, .. } => f
114                .debug_struct("Trigger")
115                .field("last_result", &last_result.load(Ordering::SeqCst))
116                .finish(),
117        }
118    }
119}
120
121/// Word-addressed software state bank exposing observable driver state slots,
122/// policy-gated runtime knobs, read probes, and diagnostic triggers.
123#[derive(Clone, Debug)]
124pub struct StateBank {
125    logical_size: u64,
126    entries: BTreeMap<u64, StateEntry>,
127    writable_offsets: Vec<u64>,
128}
129
130impl StateBank {
131    /// Creates an empty [`StateBank`] with the given byte `logical_size`.
132    pub fn new(logical_size: u64) -> Self {
133        Self { logical_size, entries: BTreeMap::new(), writable_offsets: Vec::new() }
134    }
135
136    /// Returns the logical byte size of the state bank.
137    pub fn logical_size(&self) -> u64 {
138        self.logical_size
139    }
140
141    /// Returns the 4-byte aligned offsets that are configured as writable knobs
142    /// or triggers.
143    pub fn writable_offsets(&self) -> &[u64] {
144        &self.writable_offsets
145    }
146
147    /// Builds the default [`ResourceCeiling`] for this state bank, allowing
148    /// 32-bit reads and polls within `logical_size` and permitting writes only
149    /// on offsets registered via [`Self::define_knob`] or [`Self::define_trigger`].
150    pub fn ceiling(&self) -> ResourceCeiling {
151        ResourceCeiling {
152            hard_denied: Vec::new(),
153            allow_unknown_reads: true,
154            allow_poll: true,
155            writable_registers: self
156                .writable_offsets
157                .iter()
158                .map(|&offset| WritableRegister {
159                    offset,
160                    width: WORD_SIZE as u32,
161                    allow_mask: FULL_WRITE_MASK,
162                    allow_rmw: true,
163                    require_precondition: false,
164                    precondition_mask: 0,
165                    readback: true,
166                })
167                .collect(),
168            protocol: None,
169            allow_interrupt: false,
170        }
171    }
172
173    /// Defines a read-only 32-bit observable state slot at a 4-byte aligned
174    /// `offset` and returns a [`StateSlotHandle`] for the driver to update.
175    pub fn define_state_slot(&mut self, offset: u64, initial: u32) -> StateSlotHandle {
176        let handle = StateSlotHandle::new(initial);
177        self.writable_offsets.retain(|&existing| existing != offset);
178        self.entries.insert(offset, StateEntry::AtomicSlot(handle.inner.clone()));
179        handle
180    }
181
182    /// Defines a policy-gated writable 32-bit runtime knob at a 4-byte aligned
183    /// `offset` and returns a [`StateSlotHandle`] for the driver to read.
184    pub fn define_knob(&mut self, offset: u64, initial: u32) -> StateSlotHandle {
185        let handle = StateSlotHandle::new(initial);
186        if !self.writable_offsets.contains(&offset) {
187            self.writable_offsets.push(offset);
188        }
189        self.entries.insert(offset, StateEntry::AtomicSlot(handle.inner.clone()));
190        handle
191    }
192
193    /// Defines a read-only on-demand probe callback at a 4-byte aligned `offset`.
194    pub fn define_read_probe(
195        &mut self,
196        offset: u64,
197        probe: impl Fn() -> Result<u32, BackendError> + Send + Sync + 'static,
198    ) {
199        self.writable_offsets.retain(|&existing| existing != offset);
200        self.entries.insert(offset, StateEntry::ReadProbe(Arc::new(probe)));
201    }
202
203    /// Defines a write-triggered diagnostic action at a 4-byte aligned `offset`.
204    ///
205    /// Writing `value` to `offset` invokes `trigger(value)` and stores the
206    /// returned `u32` status into the entry's `last_result` slot so `readback`
207    /// or subsequent `read32`/`poll32` calls observe the trigger outcome.
208    pub fn define_trigger(
209        &mut self,
210        offset: u64,
211        trigger: impl Fn(u32) -> Result<u32, BackendError> + Send + Sync + 'static,
212    ) {
213        if !self.writable_offsets.contains(&offset) {
214            self.writable_offsets.push(offset);
215        }
216        self.entries.insert(
217            offset,
218            StateEntry::Trigger {
219                last_result: Arc::new(AtomicU32::new(0)),
220                callback: Arc::new(trigger),
221            },
222        );
223    }
224
225    fn check_aligned_bounds(&self, offset: u64) -> Result<(), BackendError> {
226        if !offset.is_multiple_of(WORD_SIZE) {
227            return Err(BackendError::Fault);
228        }
229        let end = offset.checked_add(WORD_SIZE).ok_or(BackendError::Fault)?;
230        if end > self.logical_size {
231            return Err(BackendError::Fault);
232        }
233        Ok(())
234    }
235}
236
237impl MmioBackend for StateBank {
238    fn read32(&mut self, offset: u64) -> Result<u32, BackendError> {
239        self.check_aligned_bounds(offset)?;
240        match self.entries.get(&offset) {
241            Some(StateEntry::AtomicSlot(slot)) => Ok(slot.load(Ordering::SeqCst)),
242            Some(StateEntry::ReadProbe(probe)) => probe(),
243            Some(StateEntry::Trigger { last_result, .. }) => Ok(last_result.load(Ordering::SeqCst)),
244            None => Ok(0),
245        }
246    }
247
248    fn write32(&mut self, offset: u64, value: u32) -> Result<(), BackendError> {
249        self.check_aligned_bounds(offset)?;
250        if !self.writable_offsets.contains(&offset) {
251            return Err(BackendError::Fault);
252        }
253        match self.entries.get(&offset) {
254            Some(StateEntry::AtomicSlot(slot)) => {
255                slot.store(value, Ordering::SeqCst);
256                Ok(())
257            }
258            Some(StateEntry::Trigger { last_result, callback }) => {
259                let outcome = callback(value)?;
260                last_result.store(outcome, Ordering::SeqCst);
261                Ok(())
262            }
263            Some(StateEntry::ReadProbe(_)) | None => Err(BackendError::Fault),
264        }
265    }
266
267    fn barrier(&mut self) {
268        fence(Ordering::SeqCst);
269    }
270}
271
272impl ResourceBackend for StateBank {
273    fn read32(&mut self, offset: u64) -> Result<u32, BackendError> {
274        MmioBackend::read32(self, offset)
275    }
276
277    fn write32(&mut self, offset: u64, value: u32) -> Result<(), BackendError> {
278        MmioBackend::write32(self, offset, value)
279    }
280
281    fn barrier(&mut self) {
282        MmioBackend::barrier(self);
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289    use crate::access_policy::{
290        AccessClass, AccessPolicy, AccessRule, Denial, MmioResource, SessionMode,
291    };
292    use crate::audit_ring::AuditRing;
293    use crate::executor::{ExecLimits, Executor};
294    use crate::hardware_backend::FakeClock;
295
296    #[test]
297    fn state_slot_and_knob_roundtrip() {
298        let mut bank = StateBank::new(0x20);
299        let state_slot = bank.define_state_slot(0x00, 0x1111_2222);
300        let knob_slot = bank.define_knob(0x04, 0);
301
302        assert_eq!(bank.logical_size(), 0x20);
303        assert_eq!(bank.writable_offsets(), &[0x04]);
304
305        // Driver updates state slot; host reads it.
306        assert_eq!(MmioBackend::read32(&mut bank, 0x00), Ok(0x1111_2222));
307        state_slot.store(0x3333_4444);
308        assert_eq!(MmioBackend::read32(&mut bank, 0x00), Ok(0x3333_4444));
309        assert_eq!(state_slot.fetch_add(1), 0x3333_4444);
310        assert_eq!(MmioBackend::read32(&mut bank, 0x00), Ok(0x3333_4445));
311
312        // Direct backend write to read-only state slot fails.
313        assert_eq!(MmioBackend::write32(&mut bank, 0x00, 0x9999), Err(BackendError::Fault));
314        assert_eq!(state_slot.load(), 0x3333_4445);
315
316        // Host writes runtime knob; driver observes new knob value.
317        assert_eq!(MmioBackend::write32(&mut bank, 0x04, 250), Ok(()));
318        assert_eq!(knob_slot.load(), 250);
319        assert_eq!(MmioBackend::read32(&mut bank, 0x04), Ok(250));
320    }
321
322    #[test]
323    fn read_probe_and_trigger_execution() {
324        let mut bank = StateBank::new(0x20);
325        let counter = Arc::new(AtomicU32::new(10));
326        let counter_for_probe = counter.clone();
327        bank.define_read_probe(0x08, move || Ok(counter_for_probe.fetch_add(1, Ordering::SeqCst)));
328
329        let trigger_calls = Arc::new(AtomicU32::new(0));
330        let trigger_calls_clone = trigger_calls.clone();
331        bank.define_trigger(0x0C, move |arg| {
332            if arg == 0xBAD {
333                return Err(BackendError::Fault);
334            }
335            trigger_calls_clone.fetch_add(1, Ordering::SeqCst);
336            Ok(arg.wrapping_mul(2))
337        });
338
339        // Read probe computes on demand and rejects writes.
340        assert_eq!(MmioBackend::read32(&mut bank, 0x08), Ok(10));
341        assert_eq!(MmioBackend::read32(&mut bank, 0x08), Ok(11));
342        assert_eq!(MmioBackend::write32(&mut bank, 0x08, 1), Err(BackendError::Fault));
343
344        // Trigger starts with last_result == 0, updates on write32, and propagates faults.
345        assert_eq!(MmioBackend::read32(&mut bank, 0x0C), Ok(0));
346        assert_eq!(MmioBackend::write32(&mut bank, 0x0C, 21), Ok(()));
347        assert_eq!(trigger_calls.load(Ordering::SeqCst), 1);
348        assert_eq!(MmioBackend::read32(&mut bank, 0x0C), Ok(42));
349
350        assert_eq!(MmioBackend::write32(&mut bank, 0x0C, 0xBAD), Err(BackendError::Fault));
351        assert_eq!(MmioBackend::read32(&mut bank, 0x0C), Ok(42));
352    }
353
354    #[test]
355    fn executor_enforces_state_bank_ceiling_and_supports_trigger_readback() {
356        let mut bank = StateBank::new(0x20);
357        let violations = bank.define_state_slot(0x00, 0);
358        let fix_knob = bank.define_knob(0x04, 0);
359        let violations_for_trigger = violations.clone();
360        let fix_for_trigger = fix_knob.clone();
361        bank.define_trigger(0x08, move |_arg| {
362            if fix_for_trigger.load() == 0 {
363                violations_for_trigger.fetch_add(1);
364                Ok(1)
365            } else {
366                Ok(0)
367            }
368        });
369
370        let ceiling = bank.ceiling();
371        let mut resources = BTreeMap::new();
372        resources.insert(0, MmioResource::mmio("state0", 0x20, 0x20));
373        let mut ceilings = BTreeMap::new();
374        ceilings.insert(0, ceiling);
375
376        // Policy rejects allowlist attempting to write read-only state slot 0x00.
377        let bad_rules =
378            vec![AccessRule { resource: 0, offset: 0x00, width: 4, class: AccessClass::Write }];
379        assert!(
380            AccessPolicy::new(
381                SessionMode::Mutating,
382                resources.clone(),
383                ceilings.clone(),
384                bad_rules
385            )
386            .is_err()
387        );
388
389        // Valid mutating session allowlist for reading 0x00 and writing knob 0x04 + trigger 0x08.
390        let rules = vec![
391            AccessRule { resource: 0, offset: 0x00, width: 4, class: AccessClass::ReadOnce },
392            AccessRule { resource: 0, offset: 0x04, width: 4, class: AccessClass::Write },
393            AccessRule { resource: 0, offset: 0x08, width: 4, class: AccessClass::Write },
394        ];
395        let policy = AccessPolicy::new(SessionMode::Mutating, resources, ceilings, rules)
396            .expect("valid policy");
397
398        let mut backends = BTreeMap::new();
399        backends.insert(0, bank);
400        let mut executor = Executor::new(backends, FakeClock::new(100, 10), ExecLimits::default());
401        let mut audit = AuditRing::new(32);
402
403        // Attempting to write read-only slot 0x00 is denied by policy.
404        let denied =
405            executor.write32(&policy, &mut audit, 1, 0, 0x00, 1, FULL_WRITE_MASK, None, true);
406        assert!(matches!(
407            denied,
408            Err(crate::executor::WriteError::Denied { denial: Denial::WriteNotPermitted, .. })
409        ));
410
411        // Trigger with fix_knob == 0 increments violations and returns readback 1.
412        let trig_bug = executor
413            .write32(&policy, &mut audit, 1, 0, 0x08, 1, FULL_WRITE_MASK, None, true)
414            .expect("trigger write");
415        assert_eq!(trig_bug.readback_value, 1);
416        let read_v = executor.read32(&policy, &mut audit, 1, 0, 0x00).expect("read violations");
417        assert_eq!(read_v.value, 1);
418
419        // Enable fix_knob = 1 and re-trigger: readback is 0 and violations remain 1.
420        let knob_write = executor
421            .write32(&policy, &mut audit, 1, 0, 0x04, 1, FULL_WRITE_MASK, None, true)
422            .expect("knob write");
423        assert_eq!(knob_write.readback_value, 1);
424        let trig_fixed = executor
425            .write32(&policy, &mut audit, 1, 0, 0x08, 1, FULL_WRITE_MASK, None, true)
426            .expect("trigger write with fix");
427        assert_eq!(trig_fixed.readback_value, 0);
428        let read_v2 =
429            executor.read32(&policy, &mut audit, 1, 0, 0x00).expect("read violations after fix");
430        assert_eq!(read_v2.value, 1);
431    }
432}
lab_proxy_core/state_bank.rs

lab_proxy_core/
state_bank.rs
