diff --git a/src/transport.rs b/src/transport.rs index 6fe63fa..173de39 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -1,5 +1,5 @@ use alloc::boxed::Box; -use alloc::string::String; +use alloc::string::{String, ToString}; use alloc::vec::Vec; use alloc::vec; use crate::interface::{Interface, TargetingData}; @@ -98,12 +98,35 @@ impl Transport { } } } + + /// Poll all the interfaces and receive a message + /// + /// Returns a result with an option of `(message, peer_id)` + pub fn receive(&mut self) -> Option<(MessageBytes, u64 /* peer id*/)> { + if let Some((interface_id, (msg, peer_data))) = self.interfaces + .iter_mut() + // For each interface return (interface id, message result) + .map(|interface| (interface.id().to_string(), interface.receive())) + // If there was an error, print it + .map(|res| match res { + (id, Err(e)) => { + #[cfg(std)] + println!("An error occurred while receiving: {:?}", e); + (id, Err(e)) + } + (id, Ok(r)) => (id, Ok(r)) + }) + // Find a result where there is a message + .find(|r| matches!(r, (_, Ok(Some(_))))) + // Safely unwrap this result (we already matched `Ok(Some(_))`) + .map(|(id, r)| (id, r.unwrap().unwrap())) { + Some((msg, self.find_or_add_peer(interface_id, peer_data))) + } else { None } + } } #[cfg(test)] use crate::interface::test_interface::TestInterface; -#[cfg(test)] -use alloc::string::ToString; #[test] fn test_adding_peer_to_transport() { @@ -131,3 +154,15 @@ fn test_transport_sending() { transport.send_message(vec![239, 123], Some(peer_id)).unwrap(); assert_eq!(transport.interfaces[0].receive(), IFResult::Ok(Some((vec![239, 123], interface_targeting_data)))); } + +#[test] +fn test_transport_receiving() { + let mut transport = Transport::new(vec![Box::new(TestInterface::default())]); + let (interface_id, interface_targeting_data) = ("test_interface".to_string(), "hi".to_string()); + let peer_id = transport.find_or_add_peer(interface_id.clone(), interface_targeting_data.clone()); + transport.send_message(vec![239, 123], None).unwrap(); + assert_eq!(transport.receive(), Some((vec![239u8, 123], transport.get_peer_by_parameters(interface_id.as_str(), "").unwrap().peer_id))); + assert!(transport.receive().is_none()); + transport.send_message(vec![239, 123], Some(peer_id)).unwrap(); + assert_eq!(transport.receive(), Some((vec![239, 123], peer_id))); +}