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![0; 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        let mut buf = Vec::spare_capacity_mut(&mut self.buf);
130        let (read_count, addr) = self.inner.recv_from(&mut buf)?;
131        Ok((self.buf[0..read_count].try_into()?, addr))
132    }
133
134    fn set_timeout(&mut self, timeout: Option<Duration>) {
135        self.opts.timeout = timeout;
136    }
137
138    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
139    fn bind_device(&mut self, interface_name: &str) -> std::io::Result<()> {
140        self.inner.bind_device(Some(interface_name.as_bytes()))
141    }
142}
143
144/// An Icmpv6 socket.
145pub struct IcmpSocket6 {
146    bound_to: Option<Ipv6Addr>,
147    inner: Socket,
148    buf: Vec<u8>,
149    opts: Opts,
150}
151
152impl IcmpSocket6 {
153    /// Construct a new raw socket. The socket must be bound to an address using `bind_to`
154    /// before it can be used to send and receive packets.
155    pub fn new() -> std::io::Result<Self> {
156        let socket = Socket::new(Domain::IPV6, Type::RAW, Some(Protocol::ICMPV6))?;
157        Self::new_from_socket(socket)
158    }
159
160    fn new_from_socket(socket: Socket) -> std::io::Result<Self> {
161        socket.set_recv_buffer_size(512)?;
162        Ok(Self {
163            bound_to: None,
164            inner: socket,
165            buf: vec![0; 512],
166            opts: Opts {
167                hops: 50,
168                timeout: None,
169            },
170        })
171    }
172
173    /// Construct a new dgram socket. The socket must be bound to an address using `bind_to`
174    /// before it can be used to send and receive packets.
175    pub fn new_dgram_socket() -> std::io::Result<Self> {
176        let socket = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::ICMPV6))?;
177        Self::new_from_socket(socket)
178    }
179}
180
181impl IcmpSocket for IcmpSocket6 {
182    type AddrType = Ipv6Addr;
183    type PacketType = Icmpv6Packet;
184
185    fn set_max_hops(&mut self, hops: u32) {
186        self.opts.hops = hops;
187    }
188
189    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
190        let addr = addr.into();
191        self.bound_to = Some(addr.clone());
192        let sock = ip_to_socket(&IpAddr::V6(addr));
193        self.inner.bind(&(sock.into()))?;
194        Ok(())
195    }
196
197    fn send_to(
198        &mut self,
199        dest: Self::AddrType,
200        mut packet: Self::PacketType,
201    ) -> std::io::Result<()> {
202        let source = match self.bound_to {
203            Some(ref addr) => addr,
204            None => {
205                return Err(std::io::Error::new(
206                    std::io::ErrorKind::Other,
207                    "Socket not bound to an address",
208                ));
209            }
210        };
211        packet = packet.with_checksum(source, &dest);
212        let dest = ip_to_socket(&IpAddr::V6(dest));
213        self.inner.set_unicast_hops_v6(self.opts.hops)?;
214        let pkt = packet.get_bytes(true);
215        self.inner.send_to(&pkt, &(dest.into()))?;
216        Ok(())
217    }
218
219    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
220        self.inner.set_read_timeout(self.opts.timeout)?;
221        let mut buf = Vec::spare_capacity_mut(&mut self.buf);
222        let (read_count, addr) = self.inner.recv_from(&mut buf)?;
223        Ok((self.buf[0..read_count].try_into()?, addr))
224    }
225
226    fn set_timeout(&mut self, timeout: Option<Duration>) {
227        self.opts.timeout = timeout;
228    }
229
230    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
231    fn bind_device(&mut self, interface_name: &str) -> std::io::Result<()> {
232        self.inner.bind_device(Some(interface_name.as_bytes()))
233    }
234}
235
236impl TryFrom<Ipv4Addr> for IcmpSocket4 {
237    type Error = std::io::Error;
238
239    fn try_from(addr: Ipv4Addr) -> Result<Self, Self::Error> {
240        let mut sock = IcmpSocket4::new()?;
241        sock.bind(addr)?;
242        Ok(sock)
243    }
244}
245
246impl TryFrom<Ipv6Addr> for IcmpSocket6 {
247    type Error = std::io::Error;
248
249    fn try_from(addr: Ipv6Addr) -> Result<Self, Self::Error> {
250        let mut sock = IcmpSocket6::new()?;
251        sock.bind(addr)?;
252        Ok(sock)
253    }
254}