diff --git a/src/proto.rs b/src/proto.rs index 587d866..8e3b07a 100644 --- a/src/proto.rs +++ b/src/proto.rs @@ -118,6 +118,10 @@ impl DiscoveryMsg { Ok(()) } + pub fn get_data(&self, key: &str) -> Option<&Vec> { + self.items.iter().find(|item| item.0.as_str() == key).map(|item| &item.1) + } + /// Serializes discovery message to `peer-discovery` binary format. pub fn to_bytes(&self) -> Vec { let mut bytes = Vec::with_capacity(1 + self.id.len() + 1 + self.service_name.len() + 2); @@ -417,6 +421,14 @@ mod tests { assert!(res.is_err()); } + it "adds arbitrary data to the discovery message" { + let mut msg = DiscoveryMsg::new("service1".to_string(), TransportProtocol::Tcp, 3000); + + unwrap!(msg.add_data("file_body".to_string(), vec![1, 2, 3])); + + assert_eq!(msg.get_data("file_body"), Some(&vec![1, 2, 3])) + } + it "returns error when message already has 255 data items" { let mut msg = DiscoveryMsg::new("service1".to_string(), TransportProtocol::Tcp, 3000); for i in 1..=255 { diff --git a/src/udp.rs b/src/udp.rs index 1f9610d..fff5cf7 100644 --- a/src/udp.rs +++ b/src/udp.rs @@ -6,7 +6,7 @@ use futures::channel::mpsc; use futures::{future, FutureExt, TryFutureExt}; use futures_timer::Delay; use log::{error, info}; -use std::net::Ipv4Addr; +use std::net::{IpAddr, Ipv4Addr}; use std::time::Duration; use crate::error::Error; @@ -62,7 +62,13 @@ async fn listen_for_udp( loop { let (bytes_read, from_peer) = socket.recv_from(&mut buf).await?; let msg = match DiscoveryMsg::from_bytes(&buf[0..bytes_read]) { - Ok(msg) => msg, + Ok(mut msg) => { + if let IpAddr::V4(addrv4) = from_peer.ip() { + msg.add_addrv4(addrv4); + } + + msg + } Err(e) => { info!( "Invalid discovery message received: {}, from: {}",