From 37f3651995fc608cfe40bde9e89649ac5ea72df6 Mon Sep 17 00:00:00 2001
From: Jeremy Wall <jeremy@marzhillstudios.com>
Date: Sat, 30 Jan 2021 12:50:57 -0500
Subject: [PATCH] We need the sockaddr from rcv_from

---
 examples/ping4.rs |  3 ++-
 examples/ping6.rs |  2 +-
 src/echo.rs       |  8 +++++---
 src/socket.rs     | 16 ++++++++--------
 4 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/examples/ping4.rs b/examples/ping4.rs
index a3108b3..de612be 100644
--- a/examples/ping4.rs
+++ b/examples/ping4.rs
@@ -35,7 +35,8 @@ pub fn main() {
             ],
         )
         .unwrap();
-    let resp = echo_socket.recv_ping().unwrap();
+    let _ = echo_socket.recv_ping();
+    let (resp, _addr) = echo_socket.recv_ping().unwrap();
     println!(
         "seq: {}, identifier: {} payload: {}",
         resp.sequence,
diff --git a/examples/ping6.rs b/examples/ping6.rs
index 3163651..2c2e8e4 100644
--- a/examples/ping6.rs
+++ b/examples/ping6.rs
@@ -36,7 +36,7 @@ pub fn main() {
     // TODO(jwall): The first packet we recieve will be the one we sent.
     // We need to implement packet filtering for the socket.
     let _ = echo_socket.recv_ping();
-    let resp = echo_socket.recv_ping().unwrap();
+    let (resp, _addr) = echo_socket.recv_ping().unwrap();
     println!(
         "seq: {}, identifier: {} payload: {}",
         resp.sequence,
diff --git a/src/echo.rs b/src/echo.rs
index 0f41ed7..fb579bb 100644
--- a/src/echo.rs
+++ b/src/echo.rs
@@ -13,6 +13,8 @@
 // limitations under the License.
 use std::convert::{From, TryFrom, TryInto};
 
+use socket2::SockAddr;
+
 use crate::{
     packet::{Icmpv4Message, Icmpv4Packet, Icmpv6Message, Icmpv6Packet, WithEchoRequest},
     socket::IcmpSocket,
@@ -116,9 +118,9 @@ where
         Ok(())
     }
 
-    pub fn recv_ping(&mut self) -> std::io::Result<EchoResponse> {
-        let packet = self.inner.rcv_from()?;
-        Ok(packet.try_into()?)
+    pub fn recv_ping(&mut self) -> std::io::Result<(EchoResponse, SockAddr)> {
+        let (packet, addr) = self.inner.rcv_from()?;
+        Ok((packet.try_into()?, addr))
     }
 }
 
diff --git a/src/socket.rs b/src/socket.rs
index 8e2e82f..339ffe9 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -15,7 +15,7 @@
 use std::convert::{Into, TryFrom, TryInto};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 
-use socket2::{Domain, Protocol, Socket, Type};
+use socket2::{Domain, Protocol, SockAddr, Socket, Type};
 
 use crate::packet::{Icmpv4Packet, Icmpv6Packet};
 
@@ -33,7 +33,7 @@ pub trait IcmpSocket {
 
     fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()>;
 
-    fn rcv_from(&mut self) -> std::io::Result<Self::PacketType>;
+    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)>;
 }
 
 pub struct Opts {
@@ -84,9 +84,9 @@ impl IcmpSocket for IcmpSocket4 {
         Ok(())
     }
 
-    fn rcv_from(&mut self) -> std::io::Result<Self::PacketType> {
-        let (read_count, _addr) = dbg!(self.inner.recv_from(&mut self.buf)?);
-        Ok(self.buf[0..read_count].try_into()?)
+    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
+        let (read_count, addr) = self.inner.recv_from(&mut self.buf)?;
+        Ok((self.buf[0..read_count].try_into()?, addr))
     }
 }
 
@@ -148,9 +148,9 @@ impl IcmpSocket for IcmpSocket6 {
         Ok(())
     }
 
-    fn rcv_from(&mut self) -> std::io::Result<Self::PacketType> {
-        let (read_count, _addr) = self.inner.recv_from(&mut self.buf)?;
-        Ok(self.buf[0..read_count].try_into()?)
+    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
+        let (read_count, addr) = self.inner.recv_from(&mut self.buf)?;
+        Ok((self.buf[0..read_count].try_into()?, addr))
     }
 }