umsh/
tokio_support.rs

1//! Tokio-friendly runtime adapters and simple std-backed stores.
2
3use core::{
4    marker::PhantomData,
5    pin::Pin,
6    task::{Context, Poll},
7};
8use std::{
9    cell::RefCell,
10    collections::BTreeMap,
11    fs, io,
12    net::{Ipv4Addr, SocketAddrV4},
13    path::PathBuf,
14    sync::{Arc, Mutex, MutexGuard},
15    time::{Duration, Instant},
16};
17
18use embedded_hal_async::delay::DelayNs;
19use socket2::{Domain, Protocol, Socket, Type};
20use tokio::io::ReadBuf;
21use tokio::net::UdpSocket;
22use umsh_hal::{
23    Clock, CounterStore, KeyValueStore, Radio, RxInfo, RxOrigin, Snr, TxError, TxOptions,
24};
25
26#[cfg(feature = "software-crypto")]
27use crate::{
28    Platform,
29    crypto::software::{SoftwareAes, SoftwareIdentity, SoftwareSha256},
30};
31
32/// [`DelayNs`] adapter backed by `tokio::time::sleep`.
33#[derive(Clone, Copy, Debug, Default)]
34pub struct TokioDelay;
35
36impl DelayNs for TokioDelay {
37    async fn delay_ns(&mut self, ns: u32) {
38        tokio::time::sleep(Duration::from_nanos(u64::from(ns))).await;
39    }
40}
41
42/// Monotonic clock backed by `std::time::Instant`.
43///
44/// Implements [`Clock::poll_delay_until`] using `tokio::time::Sleep` so that
45/// the MAC coordinator can efficiently await timer deadlines.
46#[derive(Debug)]
47pub struct StdClock {
48    origin: Instant,
49    pending_sleep: RefCell<Option<Pin<Box<tokio::time::Sleep>>>>,
50}
51
52impl Clone for StdClock {
53    fn clone(&self) -> Self {
54        Self {
55            origin: self.origin,
56            pending_sleep: RefCell::new(None),
57        }
58    }
59}
60
61impl Default for StdClock {
62    fn default() -> Self {
63        Self {
64            origin: Instant::now(),
65            pending_sleep: RefCell::new(None),
66        }
67    }
68}
69
70impl StdClock {
71    /// Create a clock whose epoch starts at construction time.
72    ///
73    /// [`Clock::now_ms`] is specified as milliseconds since boot, so a
74    /// host that answers `PROP_UPTIME` must construct this when the thing
75    /// it stands for comes up rather than lazily on first use. Cloning
76    /// preserves the origin, so handing copies around is safe.
77    pub fn new() -> Self {
78        Self::default()
79    }
80}
81
82impl Clock for StdClock {
83    fn now_ms(&self) -> u64 {
84        self.origin.elapsed().as_millis() as u64
85    }
86
87    fn poll_delay_until(&self, cx: &mut Context<'_>, deadline_ms: u64) -> Poll<()> {
88        let now = self.now_ms();
89        if now >= deadline_ms {
90            return Poll::Ready(());
91        }
92
93        let remaining = Duration::from_millis(deadline_ms - now);
94        let target = tokio::time::Instant::now() + remaining;
95
96        let mut cell = self.pending_sleep.borrow_mut();
97        let sleep = cell.get_or_insert_with(|| Box::pin(tokio::time::sleep_until(target)));
98        sleep.as_mut().reset(target);
99        sleep.as_mut().poll(cx)
100    }
101}
102
103/// Thread-local cryptographic RNG seeded from the operating system.
104pub use rand::rngs::ThreadRng;
105
106/// Errors returned by the std-backed file and memory stores.
107#[derive(Debug)]
108pub enum FileStoreError {
109    Io(io::Error),
110    BufferTooSmall,
111    Poisoned,
112}
113
114impl From<io::Error> for FileStoreError {
115    fn from(error: io::Error) -> Self {
116        Self::Io(error)
117    }
118}
119
120/// Errors returned by the UDP multicast radio simulator.
121#[derive(Debug)]
122pub enum UdpMulticastRadioError {
123    Io(io::Error),
124    InvalidConfig(&'static str),
125    FrameTooLarge(usize),
126}
127
128impl From<io::Error> for UdpMulticastRadioError {
129    fn from(error: io::Error) -> Self {
130        Self::Io(error)
131    }
132}
133
134/// Configuration for [`UdpMulticastRadio`].
135#[derive(Clone, Copy, Debug)]
136pub struct UdpMulticastRadioConfig {
137    pub bind_addr: Ipv4Addr,
138    pub group_addr: Ipv4Addr,
139    pub interface_addr: Ipv4Addr,
140    pub port: u16,
141    pub max_frame_size: usize,
142    pub t_frame_ms: u32,
143    pub rssi: i16,
144    pub snr: i8,
145    pub loopback: bool,
146}
147
148impl UdpMulticastRadioConfig {
149    /// Create a simple host-local multicast configuration.
150    pub fn localhost(group_addr: Ipv4Addr, port: u16) -> Self {
151        Self {
152            bind_addr: Ipv4Addr::UNSPECIFIED,
153            group_addr,
154            interface_addr: Ipv4Addr::LOCALHOST,
155            port,
156            max_frame_size: 256,
157            // Use a LoRa-like frame duration by default so MAC retry/backoff
158            // behavior on the desktop resembles real deployments.
159            t_frame_ms: 800,
160            rssi: -40,
161            snr: 10,
162            loopback: true,
163        }
164    }
165}
166
167/// Host-side radio simulator backed by UDP multicast.
168///
169/// Raw UMSH frames are sent and received over IPv4 multicast with no additional
170/// framing. Multicast loopback is enabled by default so that multiple processes
171/// on the same host can exchange frames via the loopback interface.
172///
173/// Because multicast loopback echoes every sent frame back to the sender,
174/// [`UdpMulticastRadio`] stores the last transmitted frame and silently drops
175/// any received frame that is an exact byte-for-byte match.  A legitimately
176/// forwarded copy of the frame is never an exact match: repeaters always modify
177/// the packet (decrementing the hop count or removing a source-route entry), so
178/// only the loopback echo is affected.
179pub struct UdpMulticastRadio {
180    socket: UdpSocket,
181    group_addr: SocketAddrV4,
182    max_frame_size: usize,
183    t_frame_ms: u32,
184    rssi: i16,
185    snr: i8,
186    recv_buf: Vec<u8>,
187    last_sent: Vec<u8>,
188}
189
190impl UdpMulticastRadio {
191    /// Bind a UDP multicast simulator socket using a simple localhost-oriented configuration.
192    pub async fn bind_v4(group_addr: Ipv4Addr, port: u16) -> Result<Self, UdpMulticastRadioError> {
193        Self::bind_with_config(UdpMulticastRadioConfig::localhost(group_addr, port)).await
194    }
195
196    /// Bind a UDP multicast simulator socket using the provided configuration.
197    pub async fn bind_with_config(
198        config: UdpMulticastRadioConfig,
199    ) -> Result<Self, UdpMulticastRadioError> {
200        if !config.group_addr.is_multicast() {
201            return Err(UdpMulticastRadioError::InvalidConfig(
202                "group_addr must be an IPv4 multicast address",
203            ));
204        }
205        if config.max_frame_size == 0 {
206            return Err(UdpMulticastRadioError::InvalidConfig(
207                "max_frame_size must be non-zero",
208            ));
209        }
210
211        let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?;
212        socket.set_reuse_address(true)?;
213        #[cfg(unix)]
214        socket.set_reuse_port(true)?;
215        socket.set_nonblocking(true)?;
216        socket.bind(&SocketAddrV4::new(config.bind_addr, config.port).into())?;
217        socket.join_multicast_v4(&config.group_addr, &config.interface_addr)?;
218        socket.set_multicast_if_v4(&config.interface_addr)?;
219        socket.set_multicast_loop_v4(config.loopback)?;
220
221        let group = SocketAddrV4::new(config.group_addr, config.port);
222        let std_socket: std::net::UdpSocket = socket.into();
223        let socket = UdpSocket::from_std(std_socket)?;
224
225        Ok(Self {
226            socket,
227            group_addr: group,
228            max_frame_size: config.max_frame_size,
229            t_frame_ms: config.t_frame_ms,
230            rssi: config.rssi,
231            snr: config.snr,
232            recv_buf: vec![0u8; config.max_frame_size],
233            last_sent: Vec::new(),
234        })
235    }
236}
237
238impl Radio for UdpMulticastRadio {
239    type Error = UdpMulticastRadioError;
240
241    async fn transmit(
242        &mut self,
243        data: &[u8],
244        _options: TxOptions,
245    ) -> Result<(), TxError<Self::Error>> {
246        if data.len() > self.max_frame_size {
247            return Err(TxError::Io(UdpMulticastRadioError::FrameTooLarge(
248                data.len(),
249            )));
250        }
251
252        self.last_sent.clear();
253        self.last_sent.extend_from_slice(data);
254
255        let sent = self
256            .socket
257            .send_to(data, self.group_addr)
258            .await
259            .map_err(UdpMulticastRadioError::Io)
260            .map_err(TxError::Io)?;
261        eprintln!("[udp-radio] TX {} bytes to {}", sent, self.group_addr);
262        Ok(())
263    }
264
265    fn poll_receive(
266        &mut self,
267        cx: &mut core::task::Context<'_>,
268        buf: &mut [u8],
269    ) -> core::task::Poll<Result<RxInfo, Self::Error>> {
270        loop {
271            let mut read_buf = ReadBuf::new(&mut self.recv_buf);
272            let len = match self.socket.poll_recv(cx, &mut read_buf) {
273                core::task::Poll::Ready(Ok(())) => read_buf.filled().len(),
274                core::task::Poll::Ready(Err(error)) => {
275                    return core::task::Poll::Ready(Err(UdpMulticastRadioError::Io(error)));
276                }
277                core::task::Poll::Pending => return core::task::Poll::Pending,
278            };
279            if len == 0 {
280                return core::task::Poll::Pending;
281            }
282
283            // Drop exact loopback echoes of our own last transmission.  A frame
284            // forwarded by a repeater is never an exact match because repeaters
285            // always modify at least the hop-count or source-route field.
286            if self.recv_buf[..len] == self.last_sent[..] {
287                continue;
288            }
289
290            let copy_len = len.min(buf.len());
291            buf[..copy_len].copy_from_slice(&self.recv_buf[..copy_len]);
292            eprintln!("[udp-radio] RX {} bytes", copy_len);
293            return core::task::Poll::Ready(Ok(RxInfo {
294                len: copy_len,
295                rssi: self.rssi,
296                snr: Snr::from_decibels(self.snr),
297                lqi: None,
298                origin: RxOrigin::Air,
299            }));
300        }
301    }
302
303    fn max_frame_size(&self) -> usize {
304        self.max_frame_size
305    }
306
307    fn t_frame_ms(&self) -> u32 {
308        self.t_frame_ms
309    }
310}
311
312/// Counter store that persists one file per context key.
313#[derive(Clone, Debug)]
314pub struct TokioFileCounterStore {
315    root: PathBuf,
316}
317
318impl TokioFileCounterStore {
319    /// Create the store rooted at `root`, creating the directory if needed.
320    pub fn new(root: impl Into<PathBuf>) -> Result<Self, io::Error> {
321        let root = root.into();
322        fs::create_dir_all(&root)?;
323        Ok(Self { root })
324    }
325
326    fn path_for(&self, context: &[u8]) -> PathBuf {
327        self.root.join(format!("{}.ctr", hex_encode(context)))
328    }
329}
330
331impl CounterStore for TokioFileCounterStore {
332    type Error = FileStoreError;
333
334    async fn load(&self, context: &[u8]) -> Result<u32, Self::Error> {
335        match fs::read(self.path_for(context)) {
336            Ok(bytes) if bytes.len() == 4 => Ok(u32::from_be_bytes(
337                bytes.try_into().expect("fixed counter bytes"),
338            )),
339            Ok(_) => Ok(0),
340            Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(0),
341            Err(error) => Err(FileStoreError::Io(error)),
342        }
343    }
344
345    async fn store(&self, context: &[u8], value: u32) -> Result<(), Self::Error> {
346        fs::write(self.path_for(context), value.to_be_bytes()).map_err(FileStoreError::Io)
347    }
348
349    async fn flush(&self) -> Result<(), Self::Error> {
350        Ok(())
351    }
352}
353
354/// Key-value store that persists one file per key.
355#[derive(Clone, Debug)]
356pub struct TokioFileKeyValueStore {
357    root: PathBuf,
358}
359
360impl TokioFileKeyValueStore {
361    /// Create the store rooted at `root`, creating the directory if needed.
362    pub fn new(root: impl Into<PathBuf>) -> Result<Self, io::Error> {
363        let root = root.into();
364        fs::create_dir_all(&root)?;
365        Ok(Self { root })
366    }
367
368    fn path_for(&self, key: &[u8]) -> PathBuf {
369        self.root.join(format!("{}.bin", hex_encode(key)))
370    }
371}
372
373impl KeyValueStore for TokioFileKeyValueStore {
374    type Error = FileStoreError;
375
376    async fn load(&self, key: &[u8], buf: &mut [u8]) -> Result<Option<usize>, Self::Error> {
377        match fs::read(self.path_for(key)) {
378            Ok(value) => {
379                if value.len() > buf.len() {
380                    return Err(FileStoreError::BufferTooSmall);
381                }
382                buf[..value.len()].copy_from_slice(&value);
383                Ok(Some(value.len()))
384            }
385            Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
386            Err(error) => Err(FileStoreError::Io(error)),
387        }
388    }
389
390    async fn store(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
391        fs::write(self.path_for(key), value).map_err(FileStoreError::Io)
392    }
393
394    async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
395        match fs::remove_file(self.path_for(key)) {
396            Ok(()) => Ok(()),
397            Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
398            Err(error) => Err(FileStoreError::Io(error)),
399        }
400    }
401}
402
403/// In-memory counter store convenient for host-side tests.
404#[derive(Clone, Debug, Default)]
405pub struct MemoryCounterStore {
406    entries: Arc<Mutex<BTreeMap<Vec<u8>, u32>>>,
407}
408
409impl CounterStore for MemoryCounterStore {
410    type Error = FileStoreError;
411
412    async fn load(&self, context: &[u8]) -> Result<u32, Self::Error> {
413        Ok(*lock_entries(&self.entries)?.get(context).unwrap_or(&0))
414    }
415
416    async fn store(&self, context: &[u8], value: u32) -> Result<(), Self::Error> {
417        lock_entries(&self.entries)?.insert(context.to_vec(), value);
418        Ok(())
419    }
420
421    async fn flush(&self) -> Result<(), Self::Error> {
422        Ok(())
423    }
424}
425
426/// In-memory key-value store convenient for host-side tests.
427#[derive(Clone, Debug, Default)]
428pub struct MemoryKeyValueStore {
429    entries: Arc<Mutex<BTreeMap<Vec<u8>, Vec<u8>>>>,
430}
431
432impl KeyValueStore for MemoryKeyValueStore {
433    type Error = FileStoreError;
434
435    async fn load(&self, key: &[u8], buf: &mut [u8]) -> Result<Option<usize>, Self::Error> {
436        let entries = lock_entries(&self.entries)?;
437        let Some(value) = entries.get(key) else {
438            return Ok(None);
439        };
440        if value.len() > buf.len() {
441            return Err(FileStoreError::BufferTooSmall);
442        }
443        buf[..value.len()].copy_from_slice(value);
444        Ok(Some(value.len()))
445    }
446
447    async fn store(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
448        lock_entries(&self.entries)?.insert(key.to_vec(), value.to_vec());
449        Ok(())
450    }
451
452    async fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
453        lock_entries(&self.entries)?.remove(key);
454        Ok(())
455    }
456}
457
458/// Convenience [`crate::Platform`] implementation for Tokio-based hosts.
459#[cfg(feature = "software-crypto")]
460pub struct TokioPlatform<R, CS = TokioFileCounterStore, KV = TokioFileKeyValueStore>(
461    PhantomData<(R, CS, KV)>,
462);
463
464#[cfg(feature = "software-crypto")]
465impl<R, CS, KV> Platform for TokioPlatform<R, CS, KV>
466where
467    R: umsh_hal::Radio,
468    CS: CounterStore,
469    KV: KeyValueStore,
470{
471    type Identity = SoftwareIdentity;
472    type Aes = SoftwareAes;
473    type Sha = SoftwareSha256;
474    type Radio = R;
475    type Delay = TokioDelay;
476    type Clock = StdClock;
477    type Rng = ThreadRng;
478    type CounterStore = CS;
479    type KeyValueStore = KV;
480}
481
482fn hex_encode(bytes: &[u8]) -> String {
483    const HEX: &[u8; 16] = b"0123456789abcdef";
484    let mut out = String::with_capacity(bytes.len() * 2);
485    for byte in bytes {
486        out.push(HEX[(byte >> 4) as usize] as char);
487        out.push(HEX[(byte & 0x0f) as usize] as char);
488    }
489    out
490}
491
492fn lock_entries<T>(mutex: &Mutex<T>) -> Result<MutexGuard<'_, T>, FileStoreError> {
493    mutex.lock().map_err(|_| FileStoreError::Poisoned)
494}
495
496#[cfg(test)]
497mod tests {
498    use super::*;
499    use core::future::poll_fn;
500    use std::time::{SystemTime, UNIX_EPOCH};
501
502    fn temp_dir(name: &str) -> PathBuf {
503        let unique = SystemTime::now()
504            .duration_since(UNIX_EPOCH)
505            .expect("system time before unix epoch")
506            .as_nanos();
507        std::env::temp_dir().join(format!("umsh-{name}-{unique}"))
508    }
509
510    #[tokio::test]
511    async fn file_counter_store_round_trips_values() {
512        let root = temp_dir("counter-store");
513        let store = TokioFileCounterStore::new(&root).unwrap();
514        assert_eq!(store.load(b"peer").await.unwrap(), 0);
515        store.store(b"peer", 42).await.unwrap();
516        assert_eq!(store.load(b"peer").await.unwrap(), 42);
517        let _ = fs::remove_dir_all(root);
518    }
519
520    #[tokio::test]
521    async fn udp_multicast_radio_exchanges_frames_between_instances() {
522        let port = 40_000
523            + (SystemTime::now()
524                .duration_since(UNIX_EPOCH)
525                .expect("system time before unix epoch")
526                .subsec_nanos()
527                % 10_000) as u16;
528        let group = Ipv4Addr::new(239, 255, 42, 42);
529
530        let mut left = UdpMulticastRadio::bind_v4(group, port).await.unwrap();
531        let mut right = UdpMulticastRadio::bind_v4(group, port).await.unwrap();
532
533        left.transmit(b"ping", TxOptions::default()).await.unwrap();
534
535        let mut buf = [0u8; 16];
536        let rx = tokio::time::timeout(
537            Duration::from_secs(1),
538            poll_fn(|cx| right.poll_receive(cx, &mut buf)),
539        )
540        .await
541        .expect("udp multicast receive should complete")
542        .unwrap();
543
544        assert_eq!(rx.len, 4);
545        assert_eq!(&buf[..rx.len], b"ping");
546    }
547
548    #[tokio::test]
549    async fn file_key_value_store_round_trips_values() {
550        let root = temp_dir("kv-store");
551        let store = TokioFileKeyValueStore::new(&root).unwrap();
552        store.store(b"node", b"value").await.unwrap();
553        let mut buf = [0u8; 16];
554        let len = store.load(b"node", &mut buf).await.unwrap().unwrap();
555        assert_eq!(&buf[..len], b"value");
556        store.delete(b"node").await.unwrap();
557        assert_eq!(store.load(b"node", &mut buf).await.unwrap(), None);
558        let _ = fs::remove_dir_all(root);
559    }
560}