6107939931

Committed 07 Sep 2023 08:58AM UTC coverage: 70.376% (+0.3%) from 70.111%

Build # 6107939931

Build Type

push

github

Committed by

ensc

Commit Message

version 0.3.0

Run Details

1720 of 2444 relevant lines covered (70.38%)

388.0 hits per line

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

83.13

/src/util/udpsocket.rs

use std::io::IoSlice;
use std::net::IpAddr;
use std::os::fd::{OwnedFd, AsFd};
use std::os::unix::prelude::AsRawFd;

use nix::libc;
use nix::sys::socket::{self, SockaddrStorage};
use tokio::io::unix::AsyncFd;

use crate::{ Result, Error };

use super::SocketAddr;

trait OptUnwrap<T: Sized> {
    fn unwrap_opt(self) -> Result<T>;
}

impl <T: Sized> OptUnwrap<T> for Option<T>
{
    fn unwrap_opt(self) -> Result<T> {
        self.ok_or(Error::Internal("Option is None"))
    }
}

#[derive(Default, Debug)]
struct RecvInfoOpt {
    size:        usize,
    if_idx:        Option<libc::c_int>,
    local:        Option<IpAddr>,
    remote:        Option<SocketAddr>,
    spec_dest:        Option<IpAddr>,
}

impl RecvInfoOpt {
    fn convert_v4(ip_raw: nix::libc::in_addr) -> Option<std::net::IpAddr>
    {
        Some(std::net::Ipv4Addr::from(ip_raw.s_addr.to_be()).into())
    }

    fn convert_v6(ip_raw: nix::libc::in6_addr) -> Option<std::net::IpAddr>
    {
        Some(std::net::Ipv6Addr::from(ip_raw.s6_addr).into())
    }

    pub fn set_local_v4(&mut self, ip_raw: nix::libc::in_addr) {
        self.local = Self::convert_v4(ip_raw);
    }

    pub fn set_local_v6(&mut self, ip_raw: nix::libc::in6_addr) {
        self.local = Self::convert_v6(ip_raw);
    }

    pub fn set_spec_dest_v4(&mut self, ip_raw: nix::libc::in_addr) {
        self.spec_dest = Self::convert_v4(ip_raw);
    }
}

pub struct RecvInfo {
    pub size:        usize,
    #[allow(dead_code)]
    if_idx:        libc::c_int,
    pub local:        IpAddr,
    pub remote:        SocketAddr,
}

impl TryFrom<RecvInfoOpt> for RecvInfo {
    type Error = Error;

    fn try_from(v: RecvInfoOpt) -> std::result::Result<Self, Self::Error> {
        Ok(Self {
            size:        v.size,
            if_idx:        v.if_idx.unwrap_opt()?,
            // TODO: prefer spec_dest when set?
            local:        v.local.unwrap_opt()?,
            remote:        v.remote.unwrap_opt()?,
        })
    }
}

pub struct UdpSocket {
    af:                socket::AddressFamily,
    fd:                AsyncFd<OwnedFd>,
}

impl AsRawFd for UdpSocket {
    #[inline]
    fn as_raw_fd(&self) -> std::os::fd::RawFd {
        self.fd.as_raw_fd()
    }
}

impl UdpSocket {
    #[inline]
    fn get_fd(&self) -> &AsyncFd<OwnedFd> {
        &self.fd
    }

    pub async fn sendto(&self, buf: &[u8], addr: &SocketAddr) -> Result<()>
    {
        use socket::MsgFlags as M;
        use nix::Error as E;

        loop {
            let mut async_guard = self.get_fd().writable().await?;

            match self.sendto_sync(buf, addr, M::MSG_NOSIGNAL | M::MSG_DONTWAIT) {
                Ok(_)                        => break Ok(()),
                Err(E::EAGAIN)                => async_guard.clear_ready(),
                Err(e)                        => break Err(e.into())
            };
        }
    }

    fn sendto_sync(&self, buf: &[u8], addr: &SocketAddr,
                   flags: socket::MsgFlags) -> nix::Result<()>
    {
        use nix::Error as E;

        match socket::sendto(self.as_raw_fd(), buf, addr.as_nix(), flags) {
            Ok(sz) if sz == buf.len()        => Ok(()),
            Ok(sz)                        => {
                error!("sent only {} bytes out of {} ones", sz, buf.len());
                Err(E::ENOPKG)
            },
            Err(e)                        => Err(e)
        }
    }

    pub async fn sendmsg(&self, iov: &[IoSlice<'_>], addr: &SocketAddr) -> Result<()>
    {
        use socket::MsgFlags as M;
        use nix::Error as E;

        loop {
            let mut async_guard = self.get_fd().writable().await?;

            match self.sendmsg_sync(iov, addr, M::MSG_NOSIGNAL | M::MSG_DONTWAIT) {
                Ok(_)                        => break Ok(()),
                Err(E::EAGAIN)                => async_guard.clear_ready(),
                Err(e)                        => break Err(e.into())
            }
        }
    }

    fn sendmsg_sync(&self, iov: &[IoSlice<'_>], addr: &SocketAddr,
                    flags: socket::MsgFlags) -> nix::Result<()>
    {
        use nix::Error as E;

        let total_sz: usize = iov.iter().map(|v| v.len()).sum();

        match socket::sendmsg(self.as_raw_fd(), iov, &[], flags, Some(addr.as_nix())) {
            Ok(sz) if sz == total_sz        => Ok(()),
            Ok(sz)                        => {
                error!("sent only {} bytes out of {} ones", sz, total_sz);
                Err(E::ENOPKG)
            },
            Err(e)                        => Err(e),
        }
    }

    pub async fn recvfrom(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)>
    {
        use nix::Error as E;

        loop {
            let mut async_guard = self.get_fd().readable().await?;

            match socket::recvfrom::<SockaddrStorage>(self.as_raw_fd(), buf) {
                Ok((sz, Some(addr)))        => break Ok((sz, addr.try_into()?)),
                Ok((_, None))                => break Err(Error::Internal("no address from recvfrom")),
                Err(E::EAGAIN)                => async_guard.clear_ready(),
                Err(e)                        => break Err(e.into())
            }
        }
    }

    pub async fn recvmsg(&self, buf: &mut [u8]) -> Result<RecvInfo>
    {
        use socket::MsgFlags as M;
        use nix::Error as E;

        loop {
            let mut async_guard = self.get_fd().readable().await?;

            match self.recvmsg_sync(buf, M::MSG_DONTWAIT) {
                Ok(info)                        => break Ok(info),
                Err(Error::Nix(E::EAGAIN))        => async_guard.clear_ready(),
                Err(e)                                => break Err(e),
            }
        }
    }

    fn recvmsg_sync(&self, buf: &mut [u8], flags: socket::MsgFlags) -> Result<RecvInfo>
    {

        let mut iov = [std::io::IoSliceMut::new(buf)];
        let mut cmsg = nix::cmsg_space!(libc::in6_pktinfo,
                                        libc::in_pktinfo);

        let recv = socket::recvmsg::<SockaddrStorage>(self.as_raw_fd(), &mut iov, Some(&mut cmsg), flags)?;

        let mut res = RecvInfoOpt {
            size:        recv.bytes,
            ..Default::default()
        };

        for msg in recv.cmsgs() {
            use socket::ControlMessageOwned as C;

            match msg {
                C::Ipv4PacketInfo(i)        => {
                    res.set_local_v4(i.ipi_addr);
                    res.set_spec_dest_v4(i.ipi_spec_dst);
                    res.if_idx = Some(i.ipi_ifindex);
                },

                C::Ipv6PacketInfo(i)        => {
                    res.set_local_v6(i.ipi6_addr);
                    res.spec_dest = None;
                    res.if_idx = Some(i.ipi6_ifindex as libc::c_int);
                },

                m                        => {
                    debug!("unhandled msg {:?}", m);
                },
            }
        }

        match recv.address {
            Some(addr)        => res.remote = Some(SocketAddr::try_from(addr)?),
            None        => {
                warn!("missing remote address");
                return Err(Error::Internal("missing remote address"));
            },
        };

        res.try_into()
    }

    pub fn bind(addr: &SocketAddr) -> Result<Self> {
        let fd = addr.socket()?;
        let af = addr.get_af();

        match socket::bind(fd.as_raw_fd(), addr.as_nix()) {
            Ok(_)        => Ok(Self {
                fd:                AsyncFd::new(fd)?,
                af:                af,
            }),

            Err(e)        =>
                Err(std::io::Error::from(e).into())
        }
    }

    pub fn from_raw(fd: OwnedFd) -> Result<Self> {
        let addr = SocketAddr::from_fd(fd.as_fd())?;

        Ok(Self {
            fd:                AsyncFd::new(fd)?,
            af:                addr.get_af(),
        })
    }

    pub fn local_addr(&self) -> Result<SocketAddr> {
        let addr: SockaddrStorage = socket::getsockname(self.as_raw_fd())?;

        addr.try_into()
    }

    pub fn set_request_pktinfo(&mut self) -> Result<()> {
        use socket::AddressFamily as AF;
        use nix::sys::socket::sockopt as O;

        match self.af {
            AF::Inet        => socket::setsockopt(&self.fd, O::Ipv4PacketInfo, &true),
            AF::Inet6        => socket::setsockopt(&self.fd, O::Ipv6RecvPacketInfo, &true),
            _                => return Err(Error::Internal("unexpected af")),
        }?;

        Ok(())
    }

    pub fn set_nonblocking(&self) -> Result<()> {
        use libc::{ fcntl, F_GETFL, F_SETFL, O_NONBLOCK };

        let fd = self.as_raw_fd();
        let rc = match unsafe { fcntl(fd, F_GETFL) } {
            e if e < 0                        => e,
            f if f & O_NONBLOCK != 0        => 0,
            f                                => unsafe { fcntl(fd, F_SETFL, f | O_NONBLOCK) }
        };

        if rc < 0 {
            return Err(std::io::Error::last_os_error().into());
        }

        Ok(())
    }
}

1	use std::io::IoSlice;
2	use std::net::IpAddr;
3	use std::os::fd::{OwnedFd, AsFd};
4	use std::os::unix::prelude::AsRawFd;
5
6	use nix::libc;
7	use nix::sys::socket::{self, SockaddrStorage};
8	use tokio::io::unix::AsyncFd;
9
10	use crate::{ Result, Error };
11
12	use super::SocketAddr;
13
14	trait OptUnwrap<T: Sized> {
15	fn unwrap_opt(self) -> Result<T>;
16	}
17
18	impl <T: Sized> OptUnwrap<T> for Option<T>
19	{
20	fn unwrap_opt(self) -> Result<T> {	486✔
21	self.ok_or(Error::Internal("Option is None"))	486✔
22	}	486✔
23	}
24
25	#[derive(Default, Debug)]	324✔
26	struct RecvInfoOpt {
27	size: usize,	162✔
28	if_idx: Option<libc::c_int>,	162✔
29	local: Option<IpAddr>,	162✔
30	remote: Option<SocketAddr>,	162✔
31	spec_dest: Option<IpAddr>,	162✔
32	}
33
34	impl RecvInfoOpt {
35	fn convert_v4(ip_raw: nix::libc::in_addr) -> Option<std::net::IpAddr>	162✔
36	{
37	Some(std::net::Ipv4Addr::from(ip_raw.s_addr.to_be()).into())	162✔
38	}	162✔
39
40	fn convert_v6(ip_raw: nix::libc::in6_addr) -> Option<std::net::IpAddr>	81✔
41	{
42	Some(std::net::Ipv6Addr::from(ip_raw.s6_addr).into())	81✔
43	}	81✔
44
45	pub fn set_local_v4(&mut self, ip_raw: nix::libc::in_addr) {	81✔
46	self.local = Self::convert_v4(ip_raw);	81✔
47	}	81✔
48
49	pub fn set_local_v6(&mut self, ip_raw: nix::libc::in6_addr) {	81✔
50	self.local = Self::convert_v6(ip_raw);	81✔
51	}	81✔
52
53	pub fn set_spec_dest_v4(&mut self, ip_raw: nix::libc::in_addr) {	81✔
54	self.spec_dest = Self::convert_v4(ip_raw);	81✔
55	}	81✔
56	}
57
58	pub struct RecvInfo {
59	pub size: usize,
60	#[allow(dead_code)]
61	if_idx: libc::c_int,
62	pub local: IpAddr,
63	pub remote: SocketAddr,
64	}
65
66	impl TryFrom<RecvInfoOpt> for RecvInfo {
67	type Error = Error;
68
69	fn try_from(v: RecvInfoOpt) -> std::result::Result<Self, Self::Error> {	162✔
70	Ok(Self {	162✔
71	size: v.size,	162✔
72	if_idx: v.if_idx.unwrap_opt()?,	162✔
73	// TODO: prefer spec_dest when set?
74	local: v.local.unwrap_opt()?,	162✔
75	remote: v.remote.unwrap_opt()?,	162✔
76	})
77	}	162✔
78	}
79
80	pub struct UdpSocket {
81	af: socket::AddressFamily,
82	fd: AsyncFd<OwnedFd>,
83	}
84
85	impl AsRawFd for UdpSocket {
86	#[inline]
87	fn as_raw_fd(&self) -> std::os::fd::RawFd {	7,714✔
88	self.fd.as_raw_fd()	7,714✔
89	}	7,714✔
90	}
91
92	impl UdpSocket {
93	#[inline]
94	fn get_fd(&self) -> &AsyncFd<OwnedFd> {	7,392✔
95	&self.fd
96	}	7,392✔
97
98	pub async fn sendto(&self, buf: &[u8], addr: &SocketAddr) -> Result<()>	244✔
99	{	284✔
100	use socket::MsgFlags as M;
101	use nix::Error as E;
102
103	loop {	122✔
104	let mut async_guard = self.get_fd().writable().await?;	162✔
105
106	match self.sendto_sync(buf, addr, M::MSG_NOSIGNAL \| M::MSG_DONTWAIT) {	122✔
107	Ok(_) => break Ok(()),	122✔
108	Err(E::EAGAIN) => async_guard.clear_ready(),	×
109	Err(e) => break Err(e.into())	×
110	};
111	}
112	}	244✔
113
114	fn sendto_sync(&self, buf: &[u8], addr: &SocketAddr,	122✔
115	flags: socket::MsgFlags) -> nix::Result<()>
116	{
117	use nix::Error as E;
118
119	match socket::sendto(self.as_raw_fd(), buf, addr.as_nix(), flags) {	122✔
120	Ok(sz) if sz == buf.len() => Ok(()),	122✔
121	Ok(sz) => {	×
122	error!("sent only {} bytes out of {} ones", sz, buf.len());	×
123	Err(E::ENOPKG)	×
124	},	×
125	Err(e) => Err(e)	×
126	}
127	}	122✔
128
129	pub async fn sendmsg(&self, iov: &[IoSlice<'_>], addr: &SocketAddr) -> Result<()>	4,584✔
130	{	4,596✔
131	use socket::MsgFlags as M;
132	use nix::Error as E;
133
134	loop {	2,292✔
135	let mut async_guard = self.get_fd().writable().await?;	2,304✔
136
137	match self.sendmsg_sync(iov, addr, M::MSG_NOSIGNAL \| M::MSG_DONTWAIT) {	2,292✔
138	Ok(_) => break Ok(()),	2,292✔
139	Err(E::EAGAIN) => async_guard.clear_ready(),	×
140	Err(e) => break Err(e.into())	×
141	}
142	}
143	}	4,584✔
144
145	fn sendmsg_sync(&self, iov: &[IoSlice<'_>], addr: &SocketAddr,	2,292✔
146	flags: socket::MsgFlags) -> nix::Result<()>
147	{
148	use nix::Error as E;
149
150	let total_sz: usize = iov.iter().map(\|v\| v.len()).sum();	6,876✔
151
152	match socket::sendmsg(self.as_raw_fd(), iov, &[], flags, Some(addr.as_nix())) {	2,292✔
153	Ok(sz) if sz == total_sz => Ok(()),	2,292✔
154	Ok(sz) => {	×
155	error!("sent only {} bytes out of {} ones", sz, total_sz);	×
156	Err(E::ENOPKG)	×
157	},	×
158	Err(e) => Err(e),	×
159	}
160	}	2,292✔
161
162	pub async fn recvfrom(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)>	4,828✔
163	{	7,230✔
164	use nix::Error as E;
165
166	loop {	2,414✔
167	let mut async_guard = self.get_fd().readable().await?;	7,058✔
168
169	match socket::recvfrom::<SockaddrStorage>(self.as_raw_fd(), buf) {	4,656✔
170	Ok((sz, Some(addr))) => break Ok((sz, addr.try_into()?)),	2,414✔
171	Ok((_, None)) => break Err(Error::Internal("no address from recvfrom")),	×
172	Err(E::EAGAIN) => async_guard.clear_ready(),	2,242✔
173	Err(e) => break Err(e.into())	×
174	}
175	}
176	}	4,828✔
177
178	pub async fn recvmsg(&self, buf: &mut [u8]) -> Result<RecvInfo>	324✔
179	{	486✔
180	use socket::MsgFlags as M;
181	use nix::Error as E;
182
183	loop {	162✔
184	let mut async_guard = self.get_fd().readable().await?;	484✔
185
186	match self.recvmsg_sync(buf, M::MSG_DONTWAIT) {	322✔
187	Ok(info) => break Ok(info),	162✔
188	Err(Error::Nix(E::EAGAIN)) => async_guard.clear_ready(),	160✔
189	Err(e) => break Err(e),	×
190	}
191	}	160✔
192	}	324✔
193
194	fn recvmsg_sync(&self, buf: &mut [u8], flags: socket::MsgFlags) -> Result<RecvInfo>	322✔
195	{
196
197	let mut iov = [std::io::IoSliceMut::new(buf)];	322✔
198	let mut cmsg = nix::cmsg_space!(libc::in6_pktinfo,	322✔
199	libc::in_pktinfo);
200
201	let recv = socket::recvmsg::<SockaddrStorage>(self.as_raw_fd(), &mut iov, Some(&mut cmsg), flags)?;	322✔
202
203	let mut res = RecvInfoOpt {	162✔
204	size: recv.bytes,	162✔
205	..Default::default()	162✔
206	};
207
208	for msg in recv.cmsgs() {	324✔
209	use socket::ControlMessageOwned as C;
210
211	match msg {	162✔
212	C::Ipv4PacketInfo(i) => {	81✔
213	res.set_local_v4(i.ipi_addr);	81✔
214	res.set_spec_dest_v4(i.ipi_spec_dst);	81✔
215	res.if_idx = Some(i.ipi_ifindex);	81✔
216	},	81✔
217
218	C::Ipv6PacketInfo(i) => {	81✔
219	res.set_local_v6(i.ipi6_addr);	81✔
220	res.spec_dest = None;	81✔
221	res.if_idx = Some(i.ipi6_ifindex as libc::c_int);	81✔
222	},	81✔
223
224	m => {	×
225	debug!("unhandled msg {:?}", m);	×
226	},	×
227	}
228	}	162✔
229
230	match recv.address {	162✔
231	Some(addr) => res.remote = Some(SocketAddr::try_from(addr)?),	162✔
232	None => {
233	warn!("missing remote address");	×
234	return Err(Error::Internal("missing remote address"));	×
235	},
236	};
237
238	res.try_into()	162✔
239	}	322✔
240
241	pub fn bind(addr: &SocketAddr) -> Result<Self> {	160✔
242	let fd = addr.socket()?;	160✔
243	let af = addr.get_af();	160✔
244
245	match socket::bind(fd.as_raw_fd(), addr.as_nix()) {	160✔
246	Ok(_) => Ok(Self {	160✔
247	fd: AsyncFd::new(fd)?,	160✔
248	af: af,
249	}),	160✔
250
251	Err(e) =>	×
252	Err(std::io::Error::from(e).into())	×
253	}
254	}	160✔
255
256	pub fn from_raw(fd: OwnedFd) -> Result<Self> {	2✔
257	let addr = SocketAddr::from_fd(fd.as_fd())?;	2✔
258
259	Ok(Self {	2✔
260	fd: AsyncFd::new(fd)?,	2✔
261	af: addr.get_af(),	2✔
262	})
263	}	2✔
264
265	pub fn local_addr(&self) -> Result<SocketAddr> {	320✔
266	let addr: SockaddrStorage = socket::getsockname(self.as_raw_fd())?;	320✔
267
268	addr.try_into()	320✔
269	}	320✔
270
271	pub fn set_request_pktinfo(&mut self) -> Result<()> {	2✔
272	use socket::AddressFamily as AF;
273	use nix::sys::socket::sockopt as O;
274
275	match self.af {	4✔
276	AF::Inet => socket::setsockopt(&self.fd, O::Ipv4PacketInfo, &true),	1✔
277	AF::Inet6 => socket::setsockopt(&self.fd, O::Ipv6RecvPacketInfo, &true),	1✔
278	_ => return Err(Error::Internal("unexpected af")),	×
279	}?;	×
280
281	Ok(())	2✔
282	}	2✔
283
284	pub fn set_nonblocking(&self) -> Result<()> {	2✔
285	use libc::{ fcntl, F_GETFL, F_SETFL, O_NONBLOCK };
286
287	let fd = self.as_raw_fd();	2✔
288	let rc = match unsafe { fcntl(fd, F_GETFL) } {	2✔
289	e if e < 0 => e,	2✔
290	f if f & O_NONBLOCK != 0 => 0,	2✔
291	f => unsafe { fcntl(fd, F_SETFL, f \| O_NONBLOCK) }	2✔
292	};
293
294	if rc < 0 {	2✔
295	return Err(std::io::Error::last_os_error().into());	×
296	}
297
298	Ok(())	2✔
299	}	2✔
300	}

ensc / r-tftpd / 6107939931

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