From e850bf02360b9fbd352676b9cc6426a2ab2a9e35 Mon Sep 17 00:00:00 2001
From: Jeremy Wall <jeremy@marzhillstudios.com>
Date: Mon, 25 Jan 2021 19:00:09 -0500
Subject: [PATCH] rcv_from returns a packet instead of filling a buffer

---
 src/echo.rs   | 11 +++--------
 src/packet.rs | 15 +++++++++++++++
 src/socket.rs |  9 +++++----
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/echo.rs b/src/echo.rs
index 2dec637..1a8b5c5 100644
--- a/src/echo.rs
+++ b/src/echo.rs
@@ -98,7 +98,6 @@ impl From<IcmpSocket4> for EchoSocket4 {
 
 pub struct EchoSocket6 {
     sequence: u16,
-    buf: Vec<u8>,
     inner: IcmpSocket6,
 }
 
@@ -106,7 +105,7 @@ impl EchoSocket6 {
 
     pub fn new(sock: IcmpSocket6) -> Self {
         // TODO(jwall): How to set ICMPv6 filters.
-        EchoSocket6{inner:sock, sequence: 0, buf: Vec::with_capacity(512)}
+        EchoSocket6{inner:sock, sequence: 0}
     }
 
     pub fn set_max_hops(&mut self, hops: u32) {
@@ -121,12 +120,8 @@ impl EchoSocket6 {
     }
 
     pub fn recv_ping(&mut self) -> std::io::Result<EchoResponse> {
-        self.buf.resize(512, 0);
-        let bytes_read = self.inner.rcv_from(&mut self.buf)?;
-        match Icmpv6Packet::parse(&self.buf[0..bytes_read]) {
-            Ok(p) => return Ok(p.try_into()?),
-            Err(e) => return Err(std::io::Error::new(std::io::ErrorKind::Other, format!("Malformed ICMP Response: {:?}", e))),
-        };
+        let pkt = self.inner.rcv_from()?;
+        Ok(pkt.try_into()?)
     }
 }
 
diff --git a/src/packet.rs b/src/packet.rs
index 743f5b2..e75bcf3 100644
--- a/src/packet.rs
+++ b/src/packet.rs
@@ -340,12 +340,27 @@ impl std::fmt::Display for Icmpv6PacketBuildError {
     }
 }
 
+impl std::fmt::Display for PacketParseError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", match self {
+            PacketParseError::PacketTooSmall(c) => format!("Packet Too Small size: {}", c),
+            PacketParseError::UnrecognizedICMPType => "UnrecognizedIcmpType".to_owned(),
+        })
+    }
+}
+
 impl From<Icmpv6PacketBuildError> for std::io::Error {
     fn from(err: Icmpv6PacketBuildError) -> Self {
         std::io::Error::new(std::io::ErrorKind::Other, format!("{}", err))
     }
 }
 
+impl From<PacketParseError> for std::io::Error {
+    fn from(err: PacketParseError) -> Self {
+        std::io::Error::new(std::io::ErrorKind::Other, format!("{}", err))
+    }
+}
+
 #[cfg(test)]
 mod checksum_tests {
     use super::*;
diff --git a/src/socket.rs b/src/socket.rs
index ce4eb59..00359b0 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -108,10 +108,11 @@ impl IcmpSocket6 {
         Ok(())
     }
    
-    // TODO(jwall): This should return a packet not bytes.
-    pub fn rcv_from(&self, buf: &mut [u8]) -> std::io::Result<usize> {
-        let (read_count, _addr) = self.inner.recv_from(buf)?;
-        Ok(read_count)
+    pub fn rcv_from(&self) -> std::io::Result<Icmpv6Packet> {
+        let mut buf = vec![0; 512];
+        let (read_count, _addr) = self.inner.recv_from(&mut buf)?;
+        let pkt = Icmpv6Packet::parse(&buf[0..read_count])?;
+        Ok(pkt)
     }
 }