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