From 5c051aa8b0309083e5ebc346201b0a94c6ca7727 Mon Sep 17 00:00:00 2001
From: Povilas Balciunas <balciunas90@gmail.com>
Date: Sun, 3 Jan 2021 17:27:27 +1300
Subject: [PATCH 1/2] Yield address from which the service discovery message
 came from

This actually makes it more ergonomic as in many cases the discovery message sender
will NOT need to specify its interface addresses.
Instead the receiver will see which address the packet is coming from.
---
 src/udp.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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

From 317f67a4770291c63c07802c619c25ddf119d7c0 Mon Sep 17 00:00:00 2001
From: Povilas Balciunas <balciunas90@gmail.com>
Date: Tue, 12 Jan 2021 12:20:29 +1300
Subject: [PATCH 2/2] Implement a way to get arbitrary data from DiscoveryMsg

---
 src/proto.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

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<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 {