1use 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#[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#[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 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
103pub use rand::rngs::ThreadRng;
105
106#[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#[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#[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 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 t_frame_ms: 800,
160 rssi: -40,
161 snr: 10,
162 loopback: true,
163 }
164 }
165}
166
167pub 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 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 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 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#[derive(Clone, Debug)]
314pub struct TokioFileCounterStore {
315 root: PathBuf,
316}
317
318impl TokioFileCounterStore {
319 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#[derive(Clone, Debug)]
356pub struct TokioFileKeyValueStore {
357 root: PathBuf,
358}
359
360impl TokioFileKeyValueStore {
361 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#[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#[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#[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}