Skip to main content

icmp_socket2/
socket.rs

1// Copyright 2021 Jeremy Wall
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! ICMP Socket implementations for both ICMP4 and ICMP6 protocols.
15//!
16//! There is a common IcmpSocket trait implemented for both the v4 and v6 protocols.
17//! The socket is associated to both an address type and packet type.
18use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
19use std::{
20    convert::{Into, TryFrom, TryInto},
21    time::Duration,
22};
23
24use socket2::{Domain, Protocol, SockAddr, Socket, Type};
25
26use crate::packet::{Icmpv4Packet, Icmpv6Packet};
27
28fn ip_to_socket(ip: &IpAddr) -> SocketAddr {
29    SocketAddr::new(*ip, 0)
30}
31
32/// Trait for an IcmpSocket implemented by Icmpv4Socket and Icmpv6Socket.
33pub trait IcmpSocket {
34    /// The type of address this socket operates on.
35    type AddrType;
36    /// The type of packet this socket handles.
37    type PacketType;
38
39    /// Sets the timeout on the socket for rcv_from. A value of None
40    /// will cause rcv_from to block.
41    fn set_timeout(&mut self, timeout: Option<Duration>);
42
43    /// Sets the ttl for packets sent on this socket. Controls the number of
44    /// hops the packet will be allowed to traverse.
45    fn set_max_hops(&mut self, hops: u32);
46
47    /// Binds this socket to an address.
48    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()>;
49
50    /// Sends the packet to the given destination.
51    fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()>;
52
53    /// Receive a packet on this socket.
54    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)>;
55
56    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
57    fn bind_device(&mut self, interface_name: &str) -> std::io::Result<()>;
58}
59
60/// Options for this socket.
61struct Opts {
62    hops: u32,
63    timeout: Option<Duration>,
64}
65
66/// An ICMPv4 socket.
67pub struct IcmpSocket4 {
68    bound_to: Option<Ipv4Addr>,
69    buf: Vec<u8>,
70    inner: Socket,
71    opts: Opts,
72}
73
74impl IcmpSocket4 {
75    /// Construct a new raw socket. The socket must be bound to an address using `bind_to`
76    /// before it can be used to send and receive packets.
77    pub fn new() -> std::io::Result<Self> {
78        let socket = Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::ICMPV4))?;
79        Self::new_from_socket(socket)
80    }
81
82    fn new_from_socket(socket: Socket) -> std::io::Result<Self> {
83        socket.set_recv_buffer_size(512)?;
84        Ok(Self {
85            bound_to: None,
86            inner: socket,
87            buf: Vec::with_capacity(512),
88            opts: Opts {
89                hops: 50,
90                timeout: None,
91            },
92        })
93    }
94
95    /// Construct a new dgram socket. The socket must be bound to an address using `bind_to`
96    /// before it can be used to send and receive packets.
97    pub fn new_dgram_socket() -> std::io::Result<Self> {
98        let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::ICMPV4))?;
99        Self::new_from_socket(socket)
100    }
101}
102
103impl IcmpSocket for IcmpSocket4 {
104    type AddrType = Ipv4Addr;
105    type PacketType = Icmpv4Packet;
106
107    fn set_max_hops(&mut self, hops: u32) {
108        self.opts.hops = hops;
109    }
110
111    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
112        let addr = addr.into();
113        self.bound_to = Some(addr.clone());
114        let sock = ip_to_socket(&IpAddr::V4(addr));
115        self.inner.bind(&(sock.into()))?;
116        Ok(())
117    }
118
119    fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()> {
120        let dest = ip_to_socket(&IpAddr::V4(dest));
121        self.inner.set_ttl_v4(self.opts.hops)?;
122        self.inner
123            .send_to(&packet.with_checksum().get_bytes(true), &(dest.into()))?;
124        Ok(())
125    }
126
127    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
128        self.inner.set_read_timeout(self.opts.timeout)?;
129        self.buf.clear();
130        let (read_count, addr) = self.inner.recv_from(self.buf.spare_capacity_mut())?;
131        unsafe {
132            self.buf.set_len(read_count);
133        }
134        Ok((self.buf[..].try_into()?, addr))
135    }
136
137    fn set_timeout(&mut self, timeout: Option<Duration>) {
138        self.opts.timeout = timeout;
139    }
140
141    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
142    fn bind_device(&mut self, interface_name: &str) -> std::io::Result<()> {
143        self.inner.bind_device(Some(interface_name.as_bytes()))
144    }
145}
146
147/// An Icmpv6 socket.
148pub struct IcmpSocket6 {
149    bound_to: Option<Ipv6Addr>,
150    inner: Socket,
151    buf: Vec<u8>,
152    opts: Opts,
153}
154
155impl IcmpSocket6 {
156    /// Construct a new raw socket. The socket must be bound to an address using `bind_to`
157    /// before it can be used to send and receive packets.
158    pub fn new() -> std::io::Result<Self> {
159        let socket = Socket::new(Domain::IPV6, Type::RAW, Some(Protocol::ICMPV6))?;
160        Self::new_from_socket(socket)
161    }
162
163    fn new_from_socket(socket: Socket) -> std::io::Result<Self> {
164        socket.set_recv_buffer_size(512)?;
165        Ok(Self {
166            bound_to: None,
167            inner: socket,
168            buf: Vec::with_capacity(512),
169            opts: Opts {
170                hops: 50,
171                timeout: None,
172            },
173        })
174    }
175
176    /// Construct a new dgram socket. The socket must be bound to an address using `bind_to`
177    /// before it can be used to send and receive packets.
178    pub fn new_dgram_socket() -> std::io::Result<Self> {
179        let socket = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::ICMPV6))?;
180        Self::new_from_socket(socket)
181    }
182}
183
184impl IcmpSocket for IcmpSocket6 {
185    type AddrType = Ipv6Addr;
186    type PacketType = Icmpv6Packet;
187
188    fn set_max_hops(&mut self, hops: u32) {
189        self.opts.hops = hops;
190    }
191
192    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
193        let addr = addr.into();
194        self.bound_to = Some(addr.clone());
195        let sock = ip_to_socket(&IpAddr::V6(addr));
196        self.inner.bind(&(sock.into()))?;
197        Ok(())
198    }
199
200    fn send_to(
201        &mut self,
202        dest: Self::AddrType,
203        mut packet: Self::PacketType,
204    ) -> std::io::Result<()> {
205        let source = match self.bound_to {
206            Some(ref addr) => addr,
207            None => {
208                return Err(std::io::Error::new(
209                    std::io::ErrorKind::Other,
210                    "Socket not bound to an address",
211                ));
212            }
213        };
214        packet = packet.with_checksum(source, &dest);
215        let dest = ip_to_socket(&IpAddr::V6(dest));
216        self.inner.set_unicast_hops_v6(self.opts.hops)?;
217        let pkt = packet.get_bytes(true);
218        self.inner.send_to(&pkt, &(dest.into()))?;
219        Ok(())
220    }
221
222    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
223        self.inner.set_read_timeout(self.opts.timeout)?;
224        self.buf.clear();
225        let (read_count, addr) = self.inner.recv_from(self.buf.spare_capacity_mut())?;
226        unsafe {
227            self.buf.set_len(read_count);
228        }
229        Ok((self.buf[..].try_into()?, addr))
230    }
231
232    fn set_timeout(&mut self, timeout: Option<Duration>) {
233        self.opts.timeout = timeout;
234    }
235
236    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
237    fn bind_device(&mut self, interface_name: &str) -> std::io::Result<()> {
238        self.inner.bind_device(Some(interface_name.as_bytes()))
239    }
240}
241
242impl TryFrom<Ipv4Addr> for IcmpSocket4 {
243    type Error = std::io::Error;
244
245    fn try_from(addr: Ipv4Addr) -> Result<Self, Self::Error> {
246        let mut sock = IcmpSocket4::new()?;
247        sock.bind(addr)?;
248        Ok(sock)
249    }
250}
251
252impl TryFrom<Ipv6Addr> for IcmpSocket6 {
253    type Error = std::io::Error;
254
255    fn try_from(addr: Ipv6Addr) -> Result<Self, Self::Error> {
256        let mut sock = IcmpSocket6::new()?;
257        sock.bind(addr)?;
258        Ok(sock)
259    }
260}