3861426833

Build Type

github

Committed by GitHub

Commit Message

<a href="https://github.com/rng-dynamics/ping-fox/commit/<a class=hub.com/rng-dynamics/ping-fox/commit/08f4965244a8761edc80a7a0a76e4c49a3dc0b54">08f496524<a href="https://github.com/rng-dynamics/ping-fox/commit/08f4965244a8761edc80a7a0a76e4c49a3dc0b54">">Merge </a><a class="double-link" href="https://github.com/rng-dynamics/ping-fox/commit/<a class="double-link" href="https://github.com/rng-dynamics/ping-fox/commit/f3ad7b25b8822e68d3d10a5c961c065152050cbb">f3ad7b25b</a>">f3ad7b25b</a><a href="https://github.com/rng-dynamics/ping-fox/commit/08f4965244a8761edc80a7a0a76e4c49a3dc0b54"> into e61b23b64">e61b23b64</a>

Pull Request Pull Request #19: feat: configurable socket types

Run Details

78 of 78 new or added lines in 3 files covered. (100.0%)

606 of 678 relevant lines covered (89.38%)

3.35 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

73.6

/src/socket.rs

use std::io;
use std::time::Duration;

use pnet_packet::{ipv4::Ipv4Packet, Packet};

use socket2::{Domain, Protocol, Type};

pub(crate) trait Socket: Send + Sync {
    fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize>;

    fn recv_from(
        &self,
        buf: &mut [std::mem::MaybeUninit<u8>],
    ) -> io::Result<(usize, socket2::SockAddr)>;
}

struct DgramSocket {
    socket: socket2::Socket,
}

struct RawSocket {
    socket: socket2::Socket,
}

impl Socket for DgramSocket {
    fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize> {
        self.socket.send_to(buf, addr)
    }

    fn recv_from(
        &self,
        buf: &mut [std::mem::MaybeUninit<u8>],
    ) -> io::Result<(usize, socket2::SockAddr)> {
        socket2::Socket::recv_from(&self.socket, buf)
    }
}

impl Socket for RawSocket {
    fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize> {
        self.socket.send_to(buf, addr)
    }

    fn recv_from(
        &self,
        buf: &mut [std::mem::MaybeUninit<u8>],
    ) -> io::Result<(usize, socket2::SockAddr)> {
        // On a RAW socket we get an IP packet.
        let (n, addr) = socket2::Socket::recv_from(&self.socket, buf)?;
        // Unwrape the IP packet
        let buf2: Vec<u8> = buf
            .iter()
            .take(n)
            .map(|&b| unsafe { b.assume_init() })
            .collect();
        let ipv4_packet = Ipv4Packet::new(&buf2).expect("could not initialize IPv4 package");
        let ip_payload = ipv4_packet.payload();
        // Return only the ICMP content
        for (idx, bval) in ip_payload.iter().enumerate() {
            buf[idx] = std::mem::MaybeUninit::new(*bval);
        }
        Ok((ip_payload.len(), addr))
    }
}

pub(crate) fn create_socket2_dgram_socket(timeout: Duration) -> Result<impl Socket, io::Error> {
    let socket = socket2::Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::ICMPV4))?;
    socket
        .set_read_timeout(Some(timeout))
        .expect("could not set socket timeout");
    Ok(DgramSocket { socket })
}

pub(crate) fn create_socket2_raw_socket(timeout: Duration) -> Result<impl Socket, io::Error> {
    let socket = socket2::Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::ICMPV4))?;
    socket
        .set_read_timeout(Some(timeout))
        .expect("could not set socket timeout");
    Ok(RawSocket { socket })
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    use std::net::IpAddr;
    use std::net::SocketAddr;
    use std::sync::Mutex;

    use pnet_packet::icmp::checksum;
    use pnet_packet::icmp::echo_reply::EchoReplyPacket;
    use pnet_packet::icmp::echo_reply::MutableEchoReplyPacket;
    use pnet_packet::icmp::IcmpCode;
    use pnet_packet::icmp::IcmpPacket;
    use pnet_packet::icmp::IcmpType;
    use pnet_packet::Packet;
    use pnet_packet::PacketSize;

    #[derive(PartialEq, Eq)]
    pub(crate) enum OnSend {
        ReturnErr,
        ReturnDefault,
    }

    #[derive(PartialEq, Eq, Clone, Copy)]
    pub(crate) enum OnReceive {
        ReturnWouldBlock,
        ReturnDefault(usize),
    }

    pub(crate) struct SocketMock {
        on_send: OnSend,
        on_receive: Mutex<OnReceive>,
        sent: Mutex<Vec<(Vec<u8>, IpAddr)>>,
        received_cnt: Mutex<usize>,
    }

    impl SocketMock {
        pub(crate) fn new(on_send: OnSend, on_receive: OnReceive) -> Self {
            Self {
                on_send,
                on_receive: Mutex::new(on_receive),
                sent: Mutex::new(vec![]),
                received_cnt: Mutex::new(0),
            }
        }

        pub(crate) fn should_send_number_of_messages(&self, n: usize) -> &Self {
            assert!(n == self.sent.lock().unwrap().len());
            self
        }

        pub(crate) fn should_send_to_address(&self, addr: &IpAddr) -> &Self {
            assert!(self.sent.lock().unwrap().iter().any(|e| *addr == e.1));
            self
        }

        pub(crate) fn should_receive_number_of_messages(&self, n: usize) -> &Self {
            assert!(n == *self.received_cnt.lock().unwrap());
            self
        }
    }

    impl crate::Socket for SocketMock {
        fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize> {
            if self.on_send == OnSend::ReturnErr {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "simulating error in mock",
                ));
            }
            self.sent.lock().unwrap().push((
                buf.to_vec(),
                addr.as_socket()
                    .ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::Other,
                            "error in extracting IP address from SockAddr",
                        )
                    })?
                    .ip(),
            ));
            Ok(buf.len())
        }

        fn recv_from(
            &self,
            buf: &mut [std::mem::MaybeUninit<u8>],
        ) -> io::Result<(usize, socket2::SockAddr)> {
            let on_receive: OnReceive = *self.on_receive.lock().unwrap();
            match on_receive {
                OnReceive::ReturnWouldBlock => {
                    return Err(io::Error::new(
                        io::ErrorKind::WouldBlock,
                        "simulating would-block in mock",
                    ));
                }
                OnReceive::ReturnDefault(cnt) => {
                    *self.on_receive.lock().unwrap() = if cnt <= 1 {
                        OnReceive::ReturnWouldBlock
                    } else {
                        OnReceive::ReturnDefault(cnt - 1)
                    };
                }
            };

            let payload: Vec<u8> = vec![0xFF, 0xFF, 0xFF, 0xFF];
            if buf.len() < EchoReplyPacket::minimum_packet_size() + payload.len() {
                return Err(io::Error::new(io::ErrorKind::Other, "buffer too small"));
            }

            let mut received_cnt = self.received_cnt.lock().unwrap();
            *received_cnt += 1;

            let buf2 = vec![0u8; EchoReplyPacket::minimum_packet_size() + payload.len()];
            let mut package: MutableEchoReplyPacket<'_> =
                MutableEchoReplyPacket::owned(buf2).unwrap();
            package.set_icmp_type(IcmpType::new(0)); // echo reply
            package.set_icmp_code(IcmpCode::new(0)); // echo reply
            package.set_identifier(0xABCD_u16);
            package.set_sequence_number(0);
            package.set_payload(&payload);
            package.set_checksum(0_u16);
            package.set_checksum(checksum(&IcmpPacket::new(package.packet()).unwrap()));
            for (i, b) in package.packet().iter().enumerate() {
                buf[i].write(*b);
            }

            Ok((
                package.packet_size(),
                "127.0.0.1:12345".parse::<SocketAddr>().unwrap().into(),
            ))
        }
    }
}

1	use std::io;
2	use std::time::Duration;
3
4	use pnet_packet::{ipv4::Ipv4Packet, Packet};
5
6	use socket2::{Domain, Protocol, Type};
7
8	pub(crate) trait Socket: Send + Sync {
9	fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize>;
10
11	fn recv_from(
12	&self,
13	buf: &mut [std::mem::MaybeUninit<u8>],
14	) -> io::Result<(usize, socket2::SockAddr)>;
15	}
16
17	struct DgramSocket {
18	socket: socket2::Socket,
19	}
20
21	struct RawSocket {
22	socket: socket2::Socket,
23	}
24
25	impl Socket for DgramSocket {
26	fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize> {	2✔
27	self.socket.send_to(buf, addr)	2✔
28	}	2✔
29
30	fn recv_from(	2✔
31	&self,	2✔
32	buf: &mut [std::mem::MaybeUninit<u8>],	2✔
33	) -> io::Result<(usize, socket2::SockAddr)> {	2✔
34	socket2::Socket::recv_from(&self.socket, buf)	2✔
35	}	2✔
36	}
37
38	impl Socket for RawSocket {
39	fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize> {	×
40	self.socket.send_to(buf, addr)	×
41	}	×
42
43	fn recv_from(	×
44	&self,	×
45	buf: &mut [std::mem::MaybeUninit<u8>],	×
46	) -> io::Result<(usize, socket2::SockAddr)> {	×
47	// On a RAW socket we get an IP packet.
48	let (n, addr) = socket2::Socket::recv_from(&self.socket, buf)?;	×
49	// Unwrape the IP packet
50	let buf2: Vec<u8> = buf	×
51	.iter()	×
52	.take(n)	×
53	.map(\|&b\| unsafe { b.assume_init() })	×
54	.collect();	×
55	let ipv4_packet = Ipv4Packet::new(&buf2).expect("could not initialize IPv4 package");	×
56	let ip_payload = ipv4_packet.payload();	×
57	// Return only the ICMP content
58	for (idx, bval) in ip_payload.iter().enumerate() {	×
59	buf[idx] = std::mem::MaybeUninit::new(*bval);	×
60	}	×
61	Ok((ip_payload.len(), addr))	×
62	}	×
63	}
64
65	pub(crate) fn create_socket2_dgram_socket(timeout: Duration) -> Result<impl Socket, io::Error> {	2✔
66	let socket = socket2::Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::ICMPV4))?;	2✔
67	socket	2✔
68	.set_read_timeout(Some(timeout))	2✔
69	.expect("could not set socket timeout");	2✔
70	Ok(DgramSocket { socket })	2✔
71	}	2✔
72
73	pub(crate) fn create_socket2_raw_socket(timeout: Duration) -> Result<impl Socket, io::Error> {	×
74	let socket = socket2::Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::ICMPV4))?;	×
75	socket	×
76	.set_read_timeout(Some(timeout))	×
77	.expect("could not set socket timeout");	×
78	Ok(RawSocket { socket })	×
79	}	×
80
81	#[cfg(test)]
82	pub(crate) mod tests {
83	use super::*;
84
85	use std::net::IpAddr;
86	use std::net::SocketAddr;
87	use std::sync::Mutex;
88
89	use pnet_packet::icmp::checksum;
90	use pnet_packet::icmp::echo_reply::EchoReplyPacket;
91	use pnet_packet::icmp::echo_reply::MutableEchoReplyPacket;
92	use pnet_packet::icmp::IcmpCode;
93	use pnet_packet::icmp::IcmpPacket;
94	use pnet_packet::icmp::IcmpType;
95	use pnet_packet::Packet;
96	use pnet_packet::PacketSize;
97
98	#[derive(PartialEq, Eq)]	8✔
99	pub(crate) enum OnSend {
100	ReturnErr,
101	ReturnDefault,
102	}
103
104	#[derive(PartialEq, Eq, Clone, Copy)]	×
105	pub(crate) enum OnReceive {
106	ReturnWouldBlock,
107	ReturnDefault(usize),
108	}
109
110	pub(crate) struct SocketMock {
111	on_send: OnSend,
112	on_receive: Mutex<OnReceive>,
113	sent: Mutex<Vec<(Vec<u8>, IpAddr)>>,
114	received_cnt: Mutex<usize>,
115	}
116
117	impl SocketMock {
118	pub(crate) fn new(on_send: OnSend, on_receive: OnReceive) -> Self {	7✔
119	Self {	7✔
120	on_send,	7✔
121	on_receive: Mutex::new(on_receive),	7✔
122	sent: Mutex::new(vec![]),	7✔
123	received_cnt: Mutex::new(0),	7✔
124	}	7✔
125	}	7✔
126
127	pub(crate) fn should_send_number_of_messages(&self, n: usize) -> &Self {	1✔
128	assert!(n == self.sent.lock().unwrap().len());	1✔
129	self	1✔
130	}	1✔
131
132	pub(crate) fn should_send_to_address(&self, addr: &IpAddr) -> &Self {	1✔
133	assert!(self.sent.lock().unwrap().iter().any(\|e\| *addr == e.1));	1✔
134	self	1✔
135	}	1✔
136
137	pub(crate) fn should_receive_number_of_messages(&self, n: usize) -> &Self {	1✔
138	assert!(n == *self.received_cnt.lock().unwrap());	1✔
139	self	1✔
140	}	1✔
141	}
142
143	impl crate::Socket for SocketMock {
144	fn send_to(&self, buf: &[u8], addr: &socket2::SockAddr) -> io::Result<usize> {	8✔
145	if self.on_send == OnSend::ReturnErr {	8✔
146	return Err(io::Error::new(	1✔
147	io::ErrorKind::Other,	1✔
148	"simulating error in mock",	1✔
149	));	1✔
150	}	7✔
151	self.sent.lock().unwrap().push((	7✔
152	buf.to_vec(),	7✔
153	addr.as_socket()	7✔
154	.ok_or_else(\|\| {	7✔
155	io::Error::new(	×
156	io::ErrorKind::Other,	×
157	"error in extracting IP address from SockAddr",	×
158	)	×
159	})?	7✔
160	.ip(),	7✔
161	));	7✔
162	Ok(buf.len())	7✔
163	}	8✔
164
165	fn recv_from(	7✔
166	&self,	7✔
167	buf: &mut [std::mem::MaybeUninit<u8>],	7✔
168	) -> io::Result<(usize, socket2::SockAddr)> {	7✔
169	let on_receive: OnReceive = *self.on_receive.lock().unwrap();	7✔
170	match on_receive {	7✔
171	OnReceive::ReturnWouldBlock => {
172	return Err(io::Error::new(	2✔
173	io::ErrorKind::WouldBlock,	2✔
174	"simulating would-block in mock",	2✔
175	));	2✔
176	}
177	OnReceive::ReturnDefault(cnt) => {	5✔
178	*self.on_receive.lock().unwrap() = if cnt <= 1 {	5✔
179	OnReceive::ReturnWouldBlock	3✔
180	} else {
181	OnReceive::ReturnDefault(cnt - 1)	2✔
182	};
183	}
184	};
185
186	let payload: Vec<u8> = vec![0xFF, 0xFF, 0xFF, 0xFF];	5✔
187	if buf.len() < EchoReplyPacket::minimum_packet_size() + payload.len() {	5✔
188	return Err(io::Error::new(io::ErrorKind::Other, "buffer too small"));	×
189	}	5✔
190		5✔
191	let mut received_cnt = self.received_cnt.lock().unwrap();	5✔
192	*received_cnt += 1;	5✔
193		5✔
194	let buf2 = vec![0u8; EchoReplyPacket::minimum_packet_size() + payload.len()];	5✔
195	let mut package: MutableEchoReplyPacket<'_> =	5✔
196	MutableEchoReplyPacket::owned(buf2).unwrap();	5✔
197	package.set_icmp_type(IcmpType::new(0)); // echo reply	5✔
198	package.set_icmp_code(IcmpCode::new(0)); // echo reply	5✔
199	package.set_identifier(0xABCD_u16);	5✔
200	package.set_sequence_number(0);	5✔
201	package.set_payload(&payload);	5✔
202	package.set_checksum(0_u16);	5✔
203	package.set_checksum(checksum(&IcmpPacket::new(package.packet()).unwrap()));	5✔
204	for (i, b) in package.packet().iter().enumerate() {	60✔
205	buf[i].write(*b);	60✔
206	}	60✔
207
208	Ok((	5✔
209	package.packet_size(),	5✔
210	"127.0.0.1:12345".parse::<SocketAddr>().unwrap().into(),	5✔
211	))	5✔
212	}	7✔
213	}
214	}

rng-dynamics / ping-fox / 3861426833

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous