From ee180a157615b282e8d13f40511ce58b78104191 Mon Sep 17 00:00:00 2001 From: yuguorui Date: Tue, 30 May 2023 23:55:33 +0800 Subject: [PATCH 1/2] Add recv_with_ancillary_data for udp socket recv_with_ancillary_data allows to cooperate with sockopts such as IP_PKTINFO and IP_RECVORIGDSTADDR to obtain ancillary data. For more information, please refer to the manual page[1]. [1] https://man7.org/linux/man-pages/man7/ip.7.html --- src/net/udp.rs | 7 +++++++ src/sys/unix/udp.rs | 22 ++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/src/net/udp.rs b/src/net/udp.rs index 6129527a7..aec739da3 100644 --- a/src/net/udp.rs +++ b/src/net/udp.rs @@ -303,6 +303,13 @@ impl UdpSocket { self.inner.do_io(|inner| inner.recv(buf)) } + /// Receives in-band and ancillary data from the socket previously bound with connect(). On success, returns + /// the number of bytes read. + #[cfg(unix)] + pub fn recv_with_ancillary_data(&self, buf: &mut [u8], ancillary_data: &mut [u8]) -> io::Result { + self.inner.do_io(|inner| sys::udp::recv_with_ancillary_data(inner, buf, ancillary_data)) + } + /// Receives data from the socket, without removing it from the input queue. /// On success, returns the number of bytes read. /// diff --git a/src/sys/unix/udp.rs b/src/sys/unix/udp.rs index 843ae885c..cd14b2b77 100644 --- a/src/sys/unix/udp.rs +++ b/src/sys/unix/udp.rs @@ -29,3 +29,25 @@ pub(crate) fn only_v6(socket: &net::UdpSocket) -> io::Result { Ok(optval != 0) } + +pub(crate) fn recv_with_ancillary_data( + socket: &net::UdpSocket, + buf: &mut [u8], + ancillary_data: &mut [u8], +) -> io::Result { + let mut buf = libc::iovec { + iov_base: buf.as_mut_ptr().cast(), + iov_len: buf.len(), + }; + let mut msg_hdr = libc::msghdr { + msg_iov: &mut buf, + msg_iovlen: 1, + msg_name: std::ptr::null_mut(), + msg_namelen: 0, + msg_control: ancillary_data.as_mut_ptr().cast(), + msg_controllen: ancillary_data.len(), + msg_flags: 0, + }; + syscall!(recvmsg(socket.as_raw_fd(), &mut msg_hdr, 0)).map(|n| n as usize) +} + From a6953328217cbb836af35374ef7d6ea402c9b997 Mon Sep 17 00:00:00 2001 From: yuguorui Date: Wed, 31 May 2023 00:03:10 +0800 Subject: [PATCH 2/2] Add set_pktinfo for udp socket Add support for set_pktinfo and some basic unittests for set_pktinfo and recv_with_ancillary_data. --- src/net/udp.rs | 7 ++++ src/sys/unix/udp.rs | 11 ++++++ tests/udp_socket.rs | 89 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 107 insertions(+) diff --git a/src/net/udp.rs b/src/net/udp.rs index aec739da3..48e5b979c 100644 --- a/src/net/udp.rs +++ b/src/net/udp.rs @@ -303,6 +303,13 @@ impl UdpSocket { self.inner.do_io(|inner| inner.recv(buf)) } + /// Receives IP_PKTINFO ancillary message that contains a pktinfo structure that supplies some information + /// about the incoming packet. This works only for datagram oriented sockets. + #[cfg(unix)] + pub fn set_pktinfo(&self, enable: bool) -> io::Result<()> { + self.inner.do_io(|inner| sys::udp::set_pktinfo(inner, enable)) + } + /// Receives in-band and ancillary data from the socket previously bound with connect(). On success, returns /// the number of bytes read. #[cfg(unix)] diff --git a/src/sys/unix/udp.rs b/src/sys/unix/udp.rs index cd14b2b77..d15b4a0ed 100644 --- a/src/sys/unix/udp.rs +++ b/src/sys/unix/udp.rs @@ -51,3 +51,14 @@ pub(crate) fn recv_with_ancillary_data( syscall!(recvmsg(socket.as_raw_fd(), &mut msg_hdr, 0)).map(|n| n as usize) } +pub(crate) fn set_pktinfo(socket: &net::UdpSocket, enable: bool) -> io::Result<()> { + let val: libc::c_int = i32::from(enable); + syscall!(setsockopt( + socket.as_raw_fd(), + libc::IPPROTO_IP, + libc::IP_PKTINFO, + &val as *const libc::c_int as *const libc::c_void, + core::mem::size_of::() as libc::socklen_t, + ))?; + Ok(()) +} diff --git a/tests/udp_socket.rs b/tests/udp_socket.rs index 79b5d0312..cf7a77d8c 100644 --- a/tests/udp_socket.rs +++ b/tests/udp_socket.rs @@ -384,6 +384,95 @@ fn smoke_test_connected_udp_socket(mut socket1: UdpSocket, mut socket2: UdpSocke assert!(socket2.take_error().unwrap().is_none()); } +#[test] +fn udp_socket_iptinfo() { + let mut socket1 = UdpSocket::bind(any_local_address()).unwrap(); + let address1 = socket1.local_addr().unwrap(); + + let mut socket2 = UdpSocket::bind(any_local_address()).unwrap(); + let address2 = socket2.local_addr().unwrap(); + + socket1.set_pktinfo(true).unwrap(); + socket2.set_pktinfo(true).unwrap(); + + socket1.connect(address2).unwrap(); + socket2.connect(address1).unwrap(); + + let (mut poll, mut events) = init_with_poll(); + + assert_socket_non_blocking(&socket1); + assert_socket_close_on_exec(&socket1); + assert_socket_non_blocking(&socket2); + assert_socket_close_on_exec(&socket2); + + poll.registry() + .register( + &mut socket1, + ID1, + Interest::READABLE.add(Interest::WRITABLE), + ) + .expect("unable to register UDP socket"); + poll.registry() + .register( + &mut socket2, + ID2, + Interest::READABLE.add(Interest::WRITABLE), + ) + .expect("unable to register UDP socket"); + + expect_events( + &mut poll, + &mut events, + vec![ + ExpectEvent::new(ID1, Interest::WRITABLE), + ExpectEvent::new(ID2, Interest::WRITABLE), + ], + ); + + // ancillary_data is a cmsghdr followed by an in_pktinfo + // struct cmsghdr { + // size_t cmsg_len; /* Data byte count, including header + // (type is socklen_t in POSIX) */ + // int cmsg_level; /* Originating protocol */ + // int cmsg_type; /* Protocol-specific type */ + // /* followed by + // unsigned char cmsg_data[]; */ + // }; + // + // struct in_pktinfo { + // unsigned int ipi_ifindex; /* Interface index */ + // struct in_addr ipi_spec_dst; /* Local address */ + // struct in_addr ipi_addr; /* Header Destination + // address */ + // }; + const PKTINFO_LOCALHOST_V4 :[u8;28] = [28, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 1, 0, 0, 0, 127, 0, 0, 1, 127, 0, 0, 1]; + + let mut buf = [0; 20]; + let mut ancillary_data = [0; 28]; + assert_would_block(socket1.recv_with_ancillary_data(&mut buf, &mut ancillary_data)); + assert_would_block(socket2.recv_with_ancillary_data(&mut buf, &mut ancillary_data)); + + checked_write!(socket1.send(DATA1)); + checked_write!(socket2.send(DATA2)); + + expect_events( + &mut poll, + &mut events, + vec![ + ExpectEvent::new(ID1, Interest::READABLE), + ExpectEvent::new(ID2, Interest::READABLE), + ], + ); + + expect_read!(socket1.recv_with_ancillary_data(&mut buf, &mut ancillary_data), DATA2); + assert_eq!(ancillary_data, PKTINFO_LOCALHOST_V4); + expect_read!(socket2.recv_with_ancillary_data(&mut buf, &mut ancillary_data), DATA1); + assert_eq!(ancillary_data, PKTINFO_LOCALHOST_V4); + + assert!(socket1.take_error().unwrap().is_none()); + assert!(socket2.take_error().unwrap().is_none()); +} + #[test] fn reconnect_udp_socket_sending() { let (mut poll, mut events) = init_with_poll();