Merge pull request #4 from zaphar/allow_dgram_sockets

Add dgram sockets as an option

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "icmp-socket"
 description = "ICMP sockets for both IPv4 and IPv6"
-version = "0.1.2"
+version = "0.2.0"
 repository = "https://github.com/zaphar/icmp-socket"
 authors = ["Jeremy Wall <jeremy@marzhillstudios.com>"]
 edition = "2018"

README.md

@@ -14,9 +14,8 @@ It can construct and parse the common ICMP packets for both ICMP4 and ICMP6.
 ```rust
 let packet4 = Icmpv4Packet::with_echo_request(42, 1, "payload".to_bytes());
 let packet6 = Icmpv6Packet::with_echo_request(42, 1, "payload".to_bytes());
-)
 ```
 
 # API Documentation
 
-https://docs.rs/icmp-socket/0.1.1
+https://docs.rs/icmp-socket/0.2.0

examples/ping4.rs

@@ -74,7 +74,7 @@ pub fn main() {
         socket4
             .send_to(address.parse::<Ipv4Addr>().unwrap(), packet)
             .unwrap();
-        socket4.set_timeout(Duration::from_secs(1)).unwrap();
+        socket4.set_timeout(Some(Duration::from_secs(1)));
         loop {
             let (resp, sock_addr) = match socket4.rcv_from() {
                 Ok(tpl) => tpl,

examples/ping6.rs

@@ -72,7 +72,7 @@ pub fn main() {
         socket6
             .send_to(address.parse::<Ipv6Addr>().unwrap(), packet)
             .unwrap();
-        socket6.set_timeout(Duration::from_secs(1)).unwrap();
+        socket6.set_timeout(Some(Duration::from_secs(1)));
         loop {
             let (resp, sock_addr) = match socket6.rcv_from() {
                 Ok(tpl) => tpl,

src/socket.rs

@@ -37,9 +37,9 @@ pub trait IcmpSocket {
     /// The type of packet this socket handles.
     type PacketType;
 
-    /// Sets the timeout on the socket for rcv_from. A value of 0 will cause
-    /// rcv_from to block.
-    fn set_timeout(&mut self, timeout: Duration) -> std::io::Result<()>;
+    /// Sets the timeout on the socket for rcv_from. A value of None
+    /// will cause rcv_from to block.
+    fn set_timeout(&mut self, timeout: Option<Duration>);
 
     /// Sets the ttl for packets sent on this socket. Controls the number of
     /// hops the packet will be allowed to traverse.
@@ -56,8 +56,9 @@ pub trait IcmpSocket {
 }
 
 /// Options for this socket.
-pub struct Opts {
+struct Opts {
     hops: u32,
+    timeout: Option<Duration>,
 }
 
 /// An ICMPv4 socket.
@@ -69,18 +70,32 @@ pub struct IcmpSocket4 {
 }
 
 impl IcmpSocket4 {
-    /// Construct a new socket. The socket must be bound to an address using `bind_to`
+    /// Construct a new raw socket. The socket must be bound to an address using `bind_to`
     /// before it can be used to send and receive packets.
     pub fn new() -> std::io::Result<Self> {
         let socket = Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::ICMPV4))?;
+        Self::new_from_socket(socket)
+    }
+
+    fn new_from_socket(socket: Socket) -> std::io::Result<Self> {
         socket.set_recv_buffer_size(512)?;
         Ok(Self {
             bound_to: None,
             inner: socket,
             buf: vec![0; 512],
-            opts: Opts { hops: 50 },
+            opts: Opts {
+                hops: 50,
+                timeout: None,
+            },
         })
     }
+
+    /// Construct a new dgram socket. The socket must be bound to an address using `bind_to`
+    /// before it can be used to send and receive packets.
+    pub fn new_dgram_socket() -> std::io::Result<Self> {
+        let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::ICMPV4))?;
+        Self::new_from_socket(socket)
+    }
 }
 
 impl IcmpSocket for IcmpSocket4 {
@@ -108,7 +123,7 @@ impl IcmpSocket for IcmpSocket4 {
     }
 
     fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
-        self.inner.set_read_timeout(None)?;
+        self.inner.set_read_timeout(self.opts.timeout)?;
         // NOTE(jwall): the `recv_from` implementation promises not to write uninitialised
         // bytes to the `buf`fer, so this casting is safe.
         // TODO(jwall): change to `Vec::spare_capacity_mut` when it stabilizes.
@@ -118,8 +133,8 @@ impl IcmpSocket for IcmpSocket4 {
         Ok((self.buf[0..read_count].try_into()?, addr))
     }
 
-    fn set_timeout(&mut self, timeout: Duration) -> std::io::Result<()> {
-        self.inner.set_read_timeout(Some(timeout))
+    fn set_timeout(&mut self, timeout: Option<Duration>) {
+        self.opts.timeout = timeout;
     }
 }
 
@@ -132,18 +147,32 @@ pub struct IcmpSocket6 {
 }
 
 impl IcmpSocket6 {
-    /// Construct a new socket. The socket must be bound to an address using `bind_to`
+    /// Construct a new raw socket. The socket must be bound to an address using `bind_to`
     /// before it can be used to send and receive packets.
     pub fn new() -> std::io::Result<Self> {
         let socket = Socket::new(Domain::IPV6, Type::RAW, Some(Protocol::ICMPV6))?;
+        Self::new_from_socket(socket)
+    }
+
+    fn new_from_socket(socket: Socket) -> std::io::Result<Self> {
         socket.set_recv_buffer_size(512)?;
         Ok(Self {
             bound_to: None,
             inner: socket,
             buf: vec![0; 512],
-            opts: Opts { hops: 50 },
+            opts: Opts {
+                hops: 50,
+                timeout: None,
+            },
         })
     }
+
+    /// Construct a new dgram socket. The socket must be bound to an address using `bind_to`
+    /// before it can be used to send and receive packets.
+    pub fn new_dgram_socket() -> std::io::Result<Self> {
+        let socket = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::ICMPV6))?;
+        Self::new_from_socket(socket)
+    }
 }
 
 impl IcmpSocket for IcmpSocket6 {
@@ -185,7 +214,7 @@ impl IcmpSocket for IcmpSocket6 {
     }
 
     fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
-        self.inner.set_read_timeout(None)?;
+        self.inner.set_read_timeout(self.opts.timeout)?;
         // NOTE(jwall): the `recv_from` implementation promises not to write uninitialised
         // bytes to the `buf`fer, so this casting is safe.
         // TODO(jwall): change to `Vec::spare_capacity_mut` when it stabilizes.
@@ -195,8 +224,8 @@ impl IcmpSocket for IcmpSocket6 {
         Ok((self.buf[0..read_count].try_into()?, addr))
     }
 
-    fn set_timeout(&mut self, timeout: Duration) -> std::io::Result<()> {
-        self.inner.set_read_timeout(Some(timeout))
+    fn set_timeout(&mut self, timeout: Option<Duration>) {
+        self.opts.timeout = timeout;
     }
 }