Yield address from which the service discovery message came from

src/proto.rs

@@ -118,6 +118,10 @@ impl DiscoveryMsg {
         Ok(())
     }
 
+    pub fn get_data(&self, key: &str) -> Option<&Vec<u8>> {
+        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<u8> {
         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 {

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: {}",