Skip to content

Commit

Permalink
Add preallocated buffer version of sendmsg
Browse files Browse the repository at this point in the history
  • Loading branch information
howjmay committed Sep 22, 2024
1 parent e5ac667 commit 1eed9f4
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 8 deletions.
46 changes: 39 additions & 7 deletions src/sys/socket/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,20 @@ pub use self::addr::{SockaddrLike, SockaddrStorage};
pub use self::addr::{AddressFamily, UnixAddr};
#[cfg(not(solarish))]
pub use self::addr::{AddressFamily, UnixAddr};
#[cfg(not(any(solarish, target_os = "haiku", target_os = "hurd", target_os = "redox")))]
#[cfg(not(any(
solarish,
target_os = "haiku",
target_os = "hurd",
target_os = "redox"
)))]
#[cfg(feature = "net")]
pub use self::addr::{LinkAddr, SockaddrIn, SockaddrIn6};
#[cfg(any(solarish, target_os = "haiku", target_os = "hurd", target_os = "redox"))]
#[cfg(any(
solarish,
target_os = "haiku",
target_os = "hurd",
target_os = "redox"
))]
#[cfg(feature = "net")]
pub use self::addr::{SockaddrIn, SockaddrIn6};

Expand Down Expand Up @@ -794,17 +804,17 @@ pub enum ControlMessageOwned {
#[cfg_attr(docsrs, doc(cfg(feature = "net")))]
Ipv6HopLimit(i32),

/// Retrieve the DSCP (ToS) header field of the incoming IPv4 packet.
/// Retrieve the DSCP (ToS) header field of the incoming IPv4 packet.
#[cfg(any(linux_android, target_os = "freebsd"))]
#[cfg(feature = "net")]
#[cfg_attr(docsrs, doc(cfg(feature = "net")))]
Ipv4Tos(u8),

/// Retrieve the DSCP (Traffic Class) header field of the incoming IPv6 packet.
/// Retrieve the DSCP (Traffic Class) header field of the incoming IPv6 packet.
#[cfg(any(linux_android, target_os = "freebsd"))]
#[cfg(feature = "net")]
#[cfg_attr(docsrs, doc(cfg(feature = "net")))]
Ipv6TClass(i32),
Ipv6TClass(i32),

/// UDP Generic Receive Offload (GRO) allows receiving multiple UDP
/// packets from a single sender.
Expand Down Expand Up @@ -1577,7 +1587,7 @@ impl<'a> ControlMessage<'a> {
/// by ancillary data. Optionally direct the message at the given address,
/// as with sendto.
///
/// Allocates if cmsgs is nonempty.
/// Allocates if cmsgs is nonempty, use [`sendmsg_prealloc()`] if you want to use a pre-allocated buffer.
///
/// # Examples
/// When not directing to any specific address, use `()` for the generic type
Expand Down Expand Up @@ -1631,6 +1641,29 @@ pub fn sendmsg<S>(fd: RawFd, iov: &[IoSlice<'_>], cmsgs: &[ControlMessage],
}


/// `sendmsg_prealloc()` is the same as [`sendmsg()`] but it accepts a preallocated
/// `cmsg` buffer vector.
///
/// Send data in scatter-gather vectors to a socket, possibly accompanied
/// by ancillary data. Optionally direct the message at the given address,
/// as with sendto.
pub fn sendmsg_prealloc<S>(fd: RawFd, iov: &[IoSlice<'_>], cmsgs: &[ControlMessage],
flags: MsgFlags, addr: Option<&S>, cmsg_buffer: &mut Vec<u8>) -> Result<usize>
where S: SockaddrLike
{

if cmsg_buffer.len() < cmsgs.len() {
return Err(Errno::ENOBUFS);
}

let mhdr = pack_mhdr_to_send(&mut cmsg_buffer[..], iov, cmsgs, addr);

let ret = unsafe { libc::sendmsg(fd, &mhdr, flags.bits()) };

Errno::result(ret).map(|r| r as usize)
}


/// An extension of `sendmsg` that allows the caller to transmit multiple
/// messages on a socket using a single system call. This has performance
/// benefits for some applications.
Expand Down Expand Up @@ -2456,4 +2489,3 @@ pub fn shutdown(df: RawFd, how: Shutdown) -> Result<()> {
Errno::result(shutdown(df, how)).map(drop)
}
}

54 changes: 53 additions & 1 deletion test/sys/test_socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1204,7 +1204,6 @@ pub fn test_sendmsg_ipv4packetinfo() {
}

let cmsg = [ControlMessage::Ipv4PacketInfo(&pi)];

sendmsg(
sock.as_raw_fd(),
&iov,
Expand All @@ -1215,6 +1214,59 @@ pub fn test_sendmsg_ipv4packetinfo() {
.expect("sendmsg");
}

#[cfg(any(target_os = "linux", apple_targets, target_os = "netbsd"))]
#[test]
pub fn test_sendmsg_prealloc_ipv4packetinfo() {
use cfg_if::cfg_if;
use nix::sys::socket::{
bind, sendmsg_prealloc, socket, AddressFamily, ControlMessage,
MsgFlags, SockFlag, SockType, SockaddrIn,
};
use std::io::IoSlice;

let sock = socket(
AddressFamily::Inet,
SockType::Datagram,
SockFlag::empty(),
None,
)
.expect("socket failed");

let sock_addr = SockaddrIn::new(127, 0, 0, 1, 4000);

bind(sock.as_raw_fd(), &sock_addr).expect("bind failed");

let slice = [1u8, 2, 3, 4, 5, 6, 7, 8];
let iov = [IoSlice::new(&slice)];

cfg_if! {
if #[cfg(target_os = "netbsd")] {
let pi = libc::in_pktinfo {
ipi_ifindex: 0, /* Unspecified interface */
ipi_addr: libc::in_addr { s_addr: 0 },
};
} else {
let pi = libc::in_pktinfo {
ipi_ifindex: 0, /* Unspecified interface */
ipi_addr: libc::in_addr { s_addr: 0 },
ipi_spec_dst: sock_addr.as_ref().sin_addr,
};
}
}

let cmsg = [ControlMessage::Ipv4PacketInfo(&pi)];
let mut cmsg_buffer = vec![0 as u8; 32]; // same size as sendmsg_ipv4packetinfo
sendmsg_prealloc(
sock.as_raw_fd(),
&iov,
&cmsg,
MsgFlags::empty(),
Some(&sock_addr),
&mut cmsg_buffer,
)
.expect("sendmsg_prealloc");
}

// Verify `ControlMessage::Ipv6PacketInfo` for `sendmsg`.
// This creates a (udp) socket bound to ip6-localhost, then sends a message to
// itself but uses Ipv6PacketInfo to force the source address to be
Expand Down

0 comments on commit 1eed9f4

Please sign in to comment.