umsh_ulcp_runtime/
node_counters.rs

1//! The device node's persisted frame counters.
2//!
3//! A RAM image of the [`counter_map`](crate::counter_map) plus the
4//! journal handle behind it. `store` upserts the map; a dirty `flush`
5//! writes the whole map as one record in the counter journal.
6//!
7//! This is what stops a power cycle from reusing MAC frame-counter
8//! space: the device identity's TX reservation boundary and each peer's
9//! RX replay boundary both survive here. The MAC batches its calls (one
10//! flush per persist block of secured frames), so each flush costs one
11//! journal record write.
12
13use embassy_sync::blocking_mutex::raw::RawMutex;
14use embassy_sync::mutex::Mutex;
15
16use crate::counter_map::{CounterMap, ENCODED_MAX};
17use crate::journal::{JournalFlash, ProtoStore, SharedFlash};
18use crate::log::debug_log;
19
20/// RAM image + journal handle behind [`NodeCounterStore`].
21///
22/// Lives in a board's `StaticCell` rather than a plain static because
23/// the journal handle carries the board's flash reference, which is
24/// deliberately only ever shared through the cell pattern.
25pub struct NodeCounters<MF: RawMutex + 'static, F: JournalFlash + 'static> {
26    map: CounterMap,
27    dirty: bool,
28    /// Mounted counter journal; `None` between [`NodeCounters::new`] and
29    /// [`mount`], where flushes stay RAM-only.
30    journal: Option<ProtoStore<MF, F>>,
31}
32
33impl<MF: RawMutex + 'static, F: JournalFlash + 'static> NodeCounters<MF, F> {
34    /// The still-journal-less counter state. Call once, early in boot;
35    /// the board attaches the journal with [`mount`] before the device
36    /// node comes up.
37    pub const fn new() -> Self {
38        Self {
39            map: CounterMap::new(),
40            dirty: false,
41            journal: None,
42        }
43    }
44}
45
46impl<MF: RawMutex + 'static, F: JournalFlash + 'static> Default for NodeCounters<MF, F> {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52/// The mutex a board wraps its [`NodeCounters`] in.
53///
54/// The counter mutex (`MC`) and the flash mutex (`MF`) are separate
55/// parameters because a board can legitimately want different kinds: the
56/// ESP32 image guards its uncontended flash with a `NoopRawMutex` while
57/// the counter state, reached from the MAC pump, takes a critical
58/// section.
59pub type NodeCountersMutex<MC, MF, F> = Mutex<MC, NodeCounters<MF, F>>;
60
61/// Mount the counter journal at `page0` and load the persisted map.
62pub async fn mount<MC: RawMutex + 'static, MF: RawMutex + 'static, F: JournalFlash + 'static>(
63    counters: &'static NodeCountersMutex<MC, MF, F>,
64    flash: &'static SharedFlash<MF, F>,
65    page0: u32,
66) {
67    let (journal, payload) = ProtoStore::mount(flash, page0).await;
68    let map = payload
69        .as_deref()
70        .and_then(CounterMap::decode)
71        .unwrap_or_default();
72    debug_log(format_args!("counter journal: {} entries", map.len()));
73    let mut counters = counters.lock().await;
74    counters.map = map;
75    counters.journal = Some(journal);
76}
77
78/// Drop a previous identity's persisted TX boundary (its context is the
79/// raw 32-byte public key; per-peer RX boundaries are keyed by the
80/// *peer* key and stay meaningful across identity replacement). The next
81/// dirty flush persists the pruned map.
82pub async fn prune_stale_tx<
83    MC: RawMutex + 'static,
84    MF: RawMutex + 'static,
85    F: JournalFlash + 'static,
86>(
87    counters: &'static NodeCountersMutex<MC, MF, F>,
88    public_key: &[u8; 32],
89) {
90    let mut counters = counters.lock().await;
91    if counters.map.prune_tx_except(public_key) {
92        counters.dirty = true;
93    }
94}
95
96/// Drop all persisted device-node counters (factory clear). The RAM map
97/// clears unconditionally; a failed tombstone write self-heals because
98/// the map is left dirty and the next flush rewrites the (now empty)
99/// state.
100pub async fn clear<MC: RawMutex + 'static, MF: RawMutex + 'static, F: JournalFlash + 'static>(
101    counters: &'static NodeCountersMutex<MC, MF, F>,
102) {
103    let mut counters = counters.lock().await;
104    counters.map.clear();
105    counters.dirty = match counters.journal.as_mut() {
106        Some(journal) => journal.clear().await.is_err(),
107        None => false,
108    };
109}
110
111/// The device node's [`umsh_hal::CounterStore`].
112pub struct NodeCounterStore<
113    MC: RawMutex + 'static,
114    MF: RawMutex + 'static,
115    F: JournalFlash + 'static,
116> {
117    counters: &'static NodeCountersMutex<MC, MF, F>,
118}
119
120impl<MC: RawMutex + 'static, MF: RawMutex + 'static, F: JournalFlash + 'static>
121    NodeCounterStore<MC, MF, F>
122{
123    pub fn new(counters: &'static NodeCountersMutex<MC, MF, F>) -> Self {
124        Self { counters }
125    }
126}
127
128impl<MC: RawMutex + 'static, MF: RawMutex + 'static, F: JournalFlash + 'static>
129    umsh_hal::CounterStore for NodeCounterStore<MC, MF, F>
130{
131    type Error = ();
132
133    async fn load(&self, context: &[u8]) -> Result<u32, Self::Error> {
134        // Missing entries read as 0, the MAC's "no boundary persisted
135        // yet" sentinel.
136        Ok(self.counters.lock().await.map.get(context).unwrap_or(0))
137    }
138
139    async fn store(&self, context: &[u8], value: u32) -> Result<(), Self::Error> {
140        let mut counters = self.counters.lock().await;
141        let changed = counters.map.set(context, value).map_err(|_| ())?;
142        counters.dirty |= changed;
143        Ok(())
144    }
145
146    async fn flush(&self) -> Result<(), Self::Error> {
147        let mut counters = self.counters.lock().await;
148        if !counters.dirty {
149            return Ok(());
150        }
151        let mut payload = [0u8; ENCODED_MAX];
152        let len = counters.map.encode(&mut payload).ok_or(())?;
153        match counters.journal.as_mut() {
154            Some(journal) => journal.persist(&payload[..len]).await?,
155            // No journal mounted (a board built without persistence, or
156            // before the boot-time mount): RAM only. Report success so
157            // the MAC marks the boundary instead of re-flushing every
158            // cycle.
159            None => {}
160        }
161        counters.dirty = false;
162        Ok(())
163    }
164}