Refactor API for more ergonomic usage

Cargo.toml

@@ -8,5 +8,4 @@ edition = "2018"
 
 [dependencies]
 socket2 = "0.3.19"
-packet = "0.1.4"
 byteorder = "1.3.4"

examples/ping6.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 use std::net::Ipv6Addr;
 
+use icmp_socket::socket::IcmpSocket;
 use icmp_socket::*;
 
 pub fn main() {

src/echo.rs

@@ -14,15 +14,10 @@
 use std::convert::{From, TryFrom, TryInto};
 use std::net::{Ipv4Addr, Ipv6Addr};
 
-use packet::icmp::echo::Packet;
-use packet::{Builder, Packet as P};
-
-use crate::packet::{Icmpv6Message::EchoReply, Icmpv6Packet};
-
-// TODO(jwall): It turns out that the ICMPv6 packets are sufficiently
-// different from the ICMPv4 packets. In order to handle them appropriately
-// It is going to take some consideration.
-use crate::{IcmpSocket4, IcmpSocket6};
+use crate::{
+    packet::{Icmpv4Message, Icmpv4Packet, Icmpv6Message, Icmpv6Packet},
+    socket::{IcmpSocket, IcmpSocket4, IcmpSocket6},
+};
 
 #[derive(Debug)]
 pub struct EchoResponse {
@@ -35,7 +30,7 @@ impl TryFrom<Icmpv6Packet> for EchoResponse {
     type Error = std::io::Error;
 
     fn try_from(pkt: Icmpv6Packet) -> Result<Self, Self::Error> {
-        if let EchoReply {
+        if let Icmpv6Message::EchoReply {
             identifier,
             sequence,
             payload,
@@ -58,18 +53,44 @@ impl TryFrom<Icmpv6Packet> for EchoResponse {
     }
 }
 
+impl TryFrom<Icmpv4Packet> for EchoResponse {
+    type Error = std::io::Error;
+
+    fn try_from(pkt: Icmpv4Packet) -> Result<Self, Self::Error> {
+        if let Icmpv4Message::EchoReply {
+            identifier,
+            sequence,
+            payload,
+        } = pkt.message
+        {
+            Ok(EchoResponse {
+                identifier,
+                sequence,
+                payload,
+            })
+        } else {
+            Err(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                format!(
+                    "Incorrect icmpv4 message: {:?}, code: {}",
+                    pkt.message, pkt.code
+                ),
+            ))
+        }
+    }
+}
+
 pub struct EchoSocket4 {
     sequence: u16,
-    buf: Vec<u8>,
     inner: IcmpSocket4,
 }
 
+// TODO(jwall): Make this a trait
 impl EchoSocket4 {
     pub fn new(sock: IcmpSocket4) -> Self {
         EchoSocket4 {
             inner: sock,
             sequence: 0,
-            buf: Vec::with_capacity(512),
         }
     }
 
@@ -83,41 +104,16 @@ impl EchoSocket4 {
         identifier: u16,
         payload: &[u8],
     ) -> std::io::Result<()> {
-        let packet = packet::icmp::Builder::default()
-            .echo()
-            .unwrap()
-            .request()
-            .unwrap()
-            .identifier(identifier)
-            .unwrap()
-            .sequence(self.sequence)
-            .unwrap()
-            .payload(payload)
-            .unwrap()
-            .build()
-            .unwrap();
+        let packet =
+            Icmpv4Packet::with_echo_request(identifier, self.sequence, payload.to_owned())?;
         self.sequence += 1;
-        self.inner.send_to(dest, &packet)?;
+        self.inner.send_to(dest, packet)?;
         Ok(())
     }
 
     pub fn recv_ping(&mut self) -> std::io::Result<EchoResponse> {
-        let bytes_read = self.inner.rcv_from(&mut self.buf)?;
-        match Packet::new(&self.buf[0..bytes_read]) {
-            Ok(p) => {
-                return Ok(EchoResponse {
-                    sequence: p.sequence(),
-                    identifier: p.identifier(),
-                    payload: p.payload().to_owned(),
-                })
-            }
-            Err(e) => {
-                return Err(std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    format!("Malformed ICMP Response: {:?}", e),
-                ))
-            }
-        };
+        let packet = self.inner.rcv_from()?;
+        Ok(packet.try_into()?)
     }
 }
 

src/lib.rs

@@ -16,5 +16,6 @@ pub mod echo;
 pub mod packet;
 pub mod socket;
 
-pub use crate::packet::{Icmpv6Message, Icmpv6Packet};
-pub use socket::{IcmpSocket4, IcmpSocket6};
+pub use echo::{EchoSocket4, EchoSocket6};
+pub use packet::{Icmpv4Message, Icmpv4Packet, Icmpv6Message, Icmpv6Packet};
+pub use socket::{IcmpSocket, IcmpSocket4, IcmpSocket6};

src/packet.rs

@@ -406,6 +406,7 @@ impl From<PacketParseError> for std::io::Error {
     }
 }
 
+#[derive(Debug)]
 pub enum Icmpv4Message {
     Unreachable {
         // type 3
@@ -585,6 +586,7 @@ impl Icmpv4Message {
     }
 }
 
+#[derive(Debug)]
 pub struct Icmpv4Packet {
     pub typ: u8,
     pub code: u8,

src/socket.rs

@@ -12,17 +12,30 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::convert::{Into, TryFrom};
+use std::convert::{Into, TryFrom, TryInto};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 
 use socket2::{Domain, Protocol, Socket, Type};
 
-use crate::packet::Icmpv6Packet;
+use crate::packet::{Icmpv4Packet, Icmpv6Packet};
 
 fn ip_to_socket(ip: &IpAddr) -> SocketAddr {
     SocketAddr::new(*ip, 0)
 }
 
+pub trait IcmpSocket {
+    type AddrType;
+    type PacketType;
+
+    fn set_max_hops(&mut self, hops: u32);
+
+    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()>;
+
+    fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()>;
+
+    fn rcv_from(&self) -> std::io::Result<Self::PacketType>;
+}
+
 pub struct Opts {
     hops: u32,
 }
@@ -41,12 +54,17 @@ impl IcmpSocket4 {
             opts: Opts { hops: 50 },
         })
     }
+}
 
-    pub fn set_max_hops(&mut self, hops: u32) {
+impl IcmpSocket for IcmpSocket4 {
+    type AddrType = Ipv4Addr;
+    type PacketType = Icmpv4Packet;
+
+    fn set_max_hops(&mut self, hops: u32) {
         self.opts.hops = hops;
     }
 
-    pub fn bind<A: Into<Ipv4Addr>>(&mut self, addr: A) -> std::io::Result<()> {
+    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
         let addr = addr.into();
         self.bound_to = Some(addr.clone());
         let sock = ip_to_socket(&IpAddr::V4(addr));
@@ -54,17 +72,18 @@ impl IcmpSocket4 {
         Ok(())
     }
 
-    // TODO(jwall): This should take an actual packet not the payload.
-    pub fn send_to(&mut self, dest: Ipv4Addr, payload: &[u8]) -> std::io::Result<()> {
+    fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()> {
         let dest = ip_to_socket(&IpAddr::V4(dest));
         self.inner.set_ttl(self.opts.hops)?;
-        self.inner.send_to(payload, &(dest.into()))?;
+        self.inner
+            .send_to(&packet.get_bytes(true), &(dest.into()))?;
         Ok(())
     }
 
-    pub fn rcv_from(&self, buf: &mut [u8]) -> std::io::Result<usize> {
-        let (read_count, _addr) = self.inner.recv_from(buf)?;
-        Ok(read_count)
+    fn rcv_from(&self) -> std::io::Result<Self::PacketType> {
+        let mut buf = vec![0; 512];
+        let (read_count, _addr) = self.inner.recv_from(&mut buf)?;
+        Ok(buf[0..read_count].try_into()?)
     }
 }
 
@@ -82,12 +101,17 @@ impl IcmpSocket6 {
             opts: Opts { hops: 50 },
         })
     }
+}
 
-    pub fn set_max_hops(&mut self, hops: u32) {
+impl IcmpSocket for IcmpSocket6 {
+    type AddrType = Ipv6Addr;
+    type PacketType = Icmpv6Packet;
+
+    fn set_max_hops(&mut self, hops: u32) {
         self.opts.hops = hops;
     }
 
-    pub fn bind<A: Into<Ipv6Addr>>(&mut self, addr: A) -> std::io::Result<()> {
+    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
         let addr = addr.into();
         self.bound_to = Some(addr.clone());
         let sock = ip_to_socket(&IpAddr::V6(addr));
@@ -95,7 +119,11 @@ impl IcmpSocket6 {
         Ok(())
     }
 
-    pub fn send_to(&mut self, dest: Ipv6Addr, mut packet: Icmpv6Packet) -> std::io::Result<()> {
+    fn send_to(
+        &mut self,
+        dest: Self::AddrType,
+        mut packet: Self::PacketType,
+    ) -> std::io::Result<()> {
         let source = match self.bound_to {
             Some(ref addr) => addr,
             None => {
@@ -113,11 +141,10 @@ impl IcmpSocket6 {
         Ok(())
     }
 
-    pub fn rcv_from(&self) -> std::io::Result<Icmpv6Packet> {
+    fn rcv_from(&self) -> std::io::Result<Self::PacketType> {
         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)
+        Ok(buf[0..read_count].try_into()?)
     }
 }