icmp_socket/
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    mem::MaybeUninit,
22    time::Duration,
23};
24
25use socket2::{Domain, Protocol, SockAddr, Socket, Type};
26
27use crate::packet::{Icmpv4Packet, Icmpv6Packet};
28
29fn ip_to_socket(ip: &IpAddr) -> SocketAddr {
30    SocketAddr::new(*ip, 0)
31}
32
33/// Trait for an IcmpSocket implemented by Icmpv4Socket and Icmpv6Socket.
34pub trait IcmpSocket {
35    /// The type of address this socket operates on.
36    type AddrType;
37    /// The type of packet this socket handles.
38    type PacketType;
39
40    /// Sets the timeout on the socket for rcv_from. A value of None
41    /// will cause rcv_from to block.
42    fn set_timeout(&mut self, timeout: Option<Duration>);
43
44    /// Sets the ttl for packets sent on this socket. Controls the number of
45    /// hops the packet will be allowed to traverse.
46    fn set_max_hops(&mut self, hops: u32);
47
48    /// Binds this socket to an address.
49    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()>;
50
51    /// Sends the packet to the given destination.
52    fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()>;
53
54    /// Receive a packet on this socket.
55    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)>;
56}
57
58/// Options for this socket.
59struct Opts {
60    hops: u32,
61    timeout: Option<Duration>,
62}
63
64/// An ICMPv4 socket.
65pub struct IcmpSocket4 {
66    bound_to: Option<Ipv4Addr>,
67    buf: Vec<u8>,
68    inner: Socket,
69    opts: Opts,
70}
71
72impl IcmpSocket4 {
73    /// Construct a new socket. The socket must be bound to an address using `bind_to`
74    /// before it can be used to send and receive packets.
75    pub fn new() -> std::io::Result<Self> {
76        let socket = Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::ICMPV4))?;
77        socket.set_recv_buffer_size(512)?;
78        Ok(Self {
79            bound_to: None,
80            inner: socket,
81            buf: vec![0; 512],
82            opts: Opts {
83                hops: 50,
84                timeout: None,
85            },
86        })
87    }
88}
89
90impl IcmpSocket for IcmpSocket4 {
91    type AddrType = Ipv4Addr;
92    type PacketType = Icmpv4Packet;
93
94    fn set_max_hops(&mut self, hops: u32) {
95        self.opts.hops = hops;
96    }
97
98    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
99        let addr = addr.into();
100        self.bound_to = Some(addr.clone());
101        let sock = ip_to_socket(&IpAddr::V4(addr));
102        self.inner.bind(&(sock.into()))?;
103        Ok(())
104    }
105
106    fn send_to(&mut self, dest: Self::AddrType, packet: Self::PacketType) -> std::io::Result<()> {
107        let dest = ip_to_socket(&IpAddr::V4(dest));
108        self.inner.set_ttl(self.opts.hops)?;
109        self.inner
110            .send_to(&packet.with_checksum().get_bytes(true), &(dest.into()))?;
111        Ok(())
112    }
113
114    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
115        self.inner.set_read_timeout(self.opts.timeout)?;
116        // NOTE(jwall): the `recv_from` implementation promises not to write uninitialised
117        // bytes to the `buf`fer, so this casting is safe.
118        // TODO(jwall): change to `Vec::spare_capacity_mut` when it stabilizes.
119        let mut buf =
120            unsafe { &mut *(self.buf.as_mut_slice() as *mut [u8] as *mut [MaybeUninit<u8>]) };
121        let (read_count, addr) = self.inner.recv_from(&mut buf)?;
122        Ok((self.buf[0..read_count].try_into()?, addr))
123    }
124
125    fn set_timeout(&mut self, timeout: Option<Duration>) {
126        self.opts.timeout = timeout;
127    }
128}
129
130/// An Icmpv6 socket.
131pub struct IcmpSocket6 {
132    bound_to: Option<Ipv6Addr>,
133    inner: Socket,
134    buf: Vec<u8>,
135    opts: Opts,
136}
137
138impl IcmpSocket6 {
139    /// Construct a new socket. The socket must be bound to an address using `bind_to`
140    /// before it can be used to send and receive packets.
141    pub fn new() -> std::io::Result<Self> {
142        let socket = Socket::new(Domain::IPV6, Type::RAW, Some(Protocol::ICMPV6))?;
143        socket.set_recv_buffer_size(512)?;
144        Ok(Self {
145            bound_to: None,
146            inner: socket,
147            buf: vec![0; 512],
148            opts: Opts {
149                hops: 50,
150                timeout: None,
151            },
152        })
153    }
154}
155
156impl IcmpSocket for IcmpSocket6 {
157    type AddrType = Ipv6Addr;
158    type PacketType = Icmpv6Packet;
159
160    fn set_max_hops(&mut self, hops: u32) {
161        self.opts.hops = hops;
162    }
163
164    fn bind<A: Into<Self::AddrType>>(&mut self, addr: A) -> std::io::Result<()> {
165        let addr = addr.into();
166        self.bound_to = Some(addr.clone());
167        let sock = ip_to_socket(&IpAddr::V6(addr));
168        self.inner.bind(&(sock.into()))?;
169        Ok(())
170    }
171
172    fn send_to(
173        &mut self,
174        dest: Self::AddrType,
175        mut packet: Self::PacketType,
176    ) -> std::io::Result<()> {
177        let source = match self.bound_to {
178            Some(ref addr) => addr,
179            None => {
180                return Err(std::io::Error::new(
181                    std::io::ErrorKind::Other,
182                    "Socket not bound to an address",
183                ))
184            }
185        };
186        packet = packet.with_checksum(source, &dest);
187        let dest = ip_to_socket(&IpAddr::V6(dest));
188        self.inner.set_unicast_hops_v6(self.opts.hops)?;
189        let pkt = packet.get_bytes(true);
190        self.inner.send_to(&pkt, &(dest.into()))?;
191        Ok(())
192    }
193
194    fn rcv_from(&mut self) -> std::io::Result<(Self::PacketType, SockAddr)> {
195        self.inner.set_read_timeout(self.opts.timeout)?;
196        // NOTE(jwall): the `recv_from` implementation promises not to write uninitialised
197        // bytes to the `buf`fer, so this casting is safe.
198        // TODO(jwall): change to `Vec::spare_capacity_mut` when it stabilizes.
199        let mut buf =
200            unsafe { &mut *(self.buf.as_mut_slice() as *mut [u8] as *mut [MaybeUninit<u8>]) };
201        let (read_count, addr) = self.inner.recv_from(&mut buf)?;
202        Ok((self.buf[0..read_count].try_into()?, addr))
203    }
204
205    fn set_timeout(&mut self, timeout: Option<Duration>) {
206        self.opts.timeout = timeout;
207    }
208}
209
210impl TryFrom<Ipv4Addr> for IcmpSocket4 {
211    type Error = std::io::Error;
212
213    fn try_from(addr: Ipv4Addr) -> Result<Self, Self::Error> {
214        let mut sock = IcmpSocket4::new()?;
215        sock.bind(addr)?;
216        Ok(sock)
217    }
218}
219
220impl TryFrom<Ipv6Addr> for IcmpSocket6 {
221    type Error = std::io::Error;
222
223    fn try_from(addr: Ipv6Addr) -> Result<Self, Self::Error> {
224        let mut sock = IcmpSocket6::new()?;
225        sock.bind(addr)?;
226        Ok(sock)
227    }
228}
icmp_socket/socket.rs

icmp_socket/
socket.rs
