Skip to content

Commit 70b188a

Browse files
SUPERCILEXSUPERCILEX
committed
Add uninit buffer ancillary APIs
Signed-off-by: Alex Saveau <[email protected]>
1 parent 51a88c0 commit 70b188a

File tree

3 files changed

+73
-39
lines changed

3 files changed

+73
-39
lines changed

‎src/net/send_recv/msg.rs

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::net::UCred;
1010

1111
use core::iter::FusedIterator;
1212
use core::marker::PhantomData;
13-
use core::mem::{align_of, size_of, size_of_val, take};
13+
use core::mem::{align_of, size_of, size_of_val, take, MaybeUninit};
1414
#[cfg(linux_kernel)]
1515
use core::ptr::addr_of;
1616
use core::{ptr, slice};
@@ -24,25 +24,31 @@ use super::{RecvFlags, SendFlags, SocketAddrAny, SocketAddrV4, SocketAddrV6};
2424
///
2525
/// Allocate a buffer for a single file descriptor:
2626
/// ```
27+
/// # use core::mem::MaybeUninit;
2728
/// # use rustix::cmsg_space;
28-
/// let mut space = [0; rustix::cmsg_space!(ScmRights(1))];
29+
/// let mut space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(1))];
30+
/// # let _: &[MaybeUninit<u8>] = space.as_slice();
2931
/// ```
3032
///
3133
/// Allocate a buffer for credentials:
3234
/// ```
3335
/// # #[cfg(linux_kernel)]
3436
/// # {
37+
/// # use core::mem::MaybeUninit;
3538
/// # use rustix::cmsg_space;
36-
/// let mut space = [0; rustix::cmsg_space!(ScmCredentials(1))];
39+
/// let mut space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmCredentials(1))];
40+
/// # let _: &[MaybeUninit<u8>] = space.as_slice();
3741
/// # }
3842
/// ```
3943
///
4044
/// Allocate a buffer for two file descriptors and credentials:
4145
/// ```
4246
/// # #[cfg(linux_kernel)]
4347
/// # {
48+
/// # use core::mem::MaybeUninit;
4449
/// # use rustix::cmsg_space;
45-
/// let mut space = [0; rustix::cmsg_space!(ScmRights(2), ScmCredentials(1))];
50+
/// let mut space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(2), ScmCredentials(1))];
51+
/// # let _: &[MaybeUninit<u8>] = space.as_slice();
4652
/// # }
4753
/// ```
4854
#[macro_export]
@@ -164,7 +170,7 @@ pub enum RecvAncillaryMessage<'a> {
164170
/// [`push`]: SendAncillaryBuffer::push
165171
pub struct SendAncillaryBuffer<'buf, 'slice, 'fd> {
166172
/// Raw byte buffer for messages.
167-
buffer: &'buf mut [u8],
173+
buffer: &'buf mut [MaybeUninit<u8>],
168174

169175
/// The amount of the buffer that is used.
170176
length: usize,
@@ -179,6 +185,12 @@ impl<'buf> From<&'buf mut [u8]> for SendAncillaryBuffer<'buf, '_, '_> {
179185
}
180186
}
181187

188+
impl<'buf> From<&'buf mut [MaybeUninit<u8>]> for SendAncillaryBuffer<'buf, '_, '_> {
189+
fn from(buffer: &'buf mut [MaybeUninit<u8>]) -> Self {
190+
Self::new_(buffer)
191+
}
192+
}
193+
182194
impl Default for SendAncillaryBuffer<'_, '_, '_> {
183195
fn default() -> Self {
184196
Self {
@@ -231,6 +243,10 @@ impl<'buf, 'slice, 'fd> SendAncillaryBuffer<'buf, 'slice, 'fd> {
231243
/// [`send`]: crate::net::send
232244
#[inline]
233245
pub fn new(buffer: &'buf mut [u8]) -> Self {
246+
Self::new_(unsafe { core::mem::transmute::<&mut [u8], &mut [MaybeUninit<u8>]>(buffer) })
247+
}
248+
249+
fn new_(buffer: &'buf mut [MaybeUninit<u8>]) -> Self {
234250
Self {
235251
buffer: align_for_cmsghdr(buffer),
236252
length: 0,
@@ -248,7 +264,7 @@ impl<'buf, 'slice, 'fd> SendAncillaryBuffer<'buf, 'slice, 'fd> {
248264
return core::ptr::null_mut();
249265
}
250266

251-
self.buffer.as_mut_ptr()
267+
self.buffer.as_mut_ptr().cast()
252268
}
253269

254270
/// Returns the length of the message data.
@@ -301,7 +317,7 @@ impl<'buf, 'slice, 'fd> SendAncillaryBuffer<'buf, 'slice, 'fd> {
301317
let buffer = leap!(self.buffer.get_mut(..new_length));
302318

303319
// Fill the new part of the buffer with zeroes.
304-
buffer[self.length..new_length].fill(0);
320+
buffer[self.length..new_length].fill(MaybeUninit::new(0));
305321
self.length = new_length;
306322

307323
// Get the last header in the buffer.
@@ -339,7 +355,7 @@ impl<'slice, 'fd> Extend<SendAncillaryMessage<'slice, 'fd>>
339355
#[derive(Default)]
340356
pub struct RecvAncillaryBuffer<'buf> {
341357
/// Raw byte buffer for messages.
342-
buffer: &'buf mut [u8],
358+
buffer: &'buf mut [MaybeUninit<u8>],
343359

344360
/// The portion of the buffer we've read from already.
345361
read: usize,
@@ -354,6 +370,12 @@ impl<'buf> From<&'buf mut [u8]> for RecvAncillaryBuffer<'buf> {
354370
}
355371
}
356372

373+
impl<'buf> From<&'buf mut [MaybeUninit<u8>]> for RecvAncillaryBuffer<'buf> {
374+
fn from(buffer: &'buf mut [MaybeUninit<u8>]) -> Self {
375+
Self::new_(buffer)
376+
}
377+
}
378+
357379
impl<'buf> RecvAncillaryBuffer<'buf> {
358380
/// Create a new, empty `RecvAncillaryBuffer` from a raw byte buffer.
359381
///
@@ -396,6 +418,10 @@ impl<'buf> RecvAncillaryBuffer<'buf> {
396418
/// [`recv`]: crate::net::recv
397419
#[inline]
398420
pub fn new(buffer: &'buf mut [u8]) -> Self {
421+
Self::new_(unsafe { core::mem::transmute::<&mut [u8], &mut [MaybeUninit<u8>]>(buffer) })
422+
}
423+
424+
fn new_(buffer: &'buf mut [MaybeUninit<u8>]) -> Self {
399425
Self {
400426
buffer: align_for_cmsghdr(buffer),
401427
read: 0,
@@ -413,7 +439,7 @@ impl<'buf> RecvAncillaryBuffer<'buf> {
413439
return core::ptr::null_mut();
414440
}
415441

416-
self.buffer.as_mut_ptr()
442+
self.buffer.as_mut_ptr().cast()
417443
}
418444

419445
/// Returns the length of the message data.
@@ -454,7 +480,7 @@ impl Drop for RecvAncillaryBuffer<'_> {
454480
/// Return a slice of `buffer` starting at the first `cmsghdr` alignment
455481
/// boundary.
456482
#[inline]
457-
fn align_for_cmsghdr(buffer: &mut [u8]) -> &mut [u8] {
483+
fn align_for_cmsghdr(buffer: &mut [MaybeUninit<u8>]) -> &mut [MaybeUninit<u8>] {
458484
// If the buffer is empty, we won't be writing anything into it, so it
459485
// doesn't need to be aligned.
460486
if buffer.is_empty() {
@@ -486,7 +512,9 @@ impl<'buf> AncillaryDrain<'buf> {
486512
/// The buffer must contain valid message data (or be empty).
487513
pub unsafe fn parse(buffer: &'buf mut [u8]) -> Self {
488514
Self {
489-
messages: messages::Messages::new(buffer),
515+
messages: messages::Messages::new(unsafe {
516+
core::mem::transmute::<&mut [u8], &mut [MaybeUninit<u8>]>(buffer)
517+
}),
490518
read_and_length: None,
491519
}
492520
}
@@ -920,6 +948,7 @@ mod messages {
920948
use crate::backend::net::msghdr;
921949
use core::iter::FusedIterator;
922950
use core::marker::PhantomData;
951+
use core::mem::MaybeUninit;
923952
use core::ptr::NonNull;
924953

925954
/// An iterator over the messages in an ancillary buffer.
@@ -933,12 +962,12 @@ mod messages {
933962
header: Option<NonNull<c::cmsghdr>>,
934963

935964
/// Capture the original lifetime of the buffer.
936-
_buffer: PhantomData<&'buf mut [u8]>,
965+
_buffer: PhantomData<&'buf mut [MaybeUninit<u8>]>,
937966
}
938967

939968
impl<'buf> Messages<'buf> {
940969
/// Create a new iterator over messages from a byte buffer.
941-
pub(super) fn new(buf: &'buf mut [u8]) -> Self {
970+
pub(super) fn new(buf: &'buf mut [MaybeUninit<u8>]) -> Self {
942971
let msghdr = {
943972
let mut h = msghdr::zero_msghdr();
944973
h.msg_control = buf.as_mut_ptr().cast();

‎tests/net/unix.rs

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use rustix::net::{
1414
accept, bind_unix, connect_unix, listen, socket, AddressFamily, SocketAddrUnix, SocketType,
1515
};
1616
use rustix::path::DecInt;
17+
use std::mem::MaybeUninit;
1718
use std::path::Path;
1819
use std::str::FromStr;
1920
use std::sync::{Arc, Condvar, Mutex};
@@ -447,13 +448,13 @@ fn test_unix_msg_with_scm_rights() {
447448
let mut pipe_end = None;
448449

449450
let mut buffer = [0; BUFFER_SIZE];
450-
let mut cmsg_space = [0; rustix::cmsg_space!(ScmRights(1))];
451+
let mut cmsg_space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(1))];
451452

452453
'exit: loop {
453454
let data_socket = accept(&connection_socket).unwrap();
454455
let mut sum = 0;
455456
loop {
456-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
457+
let mut cmsg_buffer = RecvAncillaryBuffer::from(cmsg_space.as_mut_slice());
457458
let nread = recvmsg(
458459
&data_socket,
459460
&mut [IoSliceMut::new(&mut buffer)],
@@ -555,8 +556,8 @@ fn test_unix_msg_with_scm_rights() {
555556
// Format the CMSG.
556557
let we = [write_end.as_fd()];
557558
let msg = SendAncillaryMessage::ScmRights(&we);
558-
let mut space = [0; rustix::cmsg_space!(ScmRights(1))];
559-
let mut cmsg_buffer = SendAncillaryBuffer::new(&mut space);
559+
let mut space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(1))];
560+
let mut cmsg_buffer = SendAncillaryBuffer::from(space.as_mut_slice());
560561
assert!(cmsg_buffer.push(msg));
561562

562563
connect_unix(&data_socket, &addr).unwrap();
@@ -615,8 +616,8 @@ fn test_unix_peercred_explicit() {
615616

616617
let ucred = sockopt::get_socket_peercred(&send_sock).unwrap();
617618
let msg = SendAncillaryMessage::ScmCredentials(ucred);
618-
let mut space = [0; rustix::cmsg_space!(ScmCredentials(1))];
619-
let mut cmsg_buffer = SendAncillaryBuffer::new(&mut space);
619+
let mut space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmCredentials(1))];
620+
let mut cmsg_buffer = SendAncillaryBuffer::from(space.as_mut_slice());
620621
assert!(cmsg_buffer.push(msg));
621622

622623
sendmsg(
@@ -627,8 +628,8 @@ fn test_unix_peercred_explicit() {
627628
)
628629
.unwrap();
629630

630-
let mut cmsg_space = [0; rustix::cmsg_space!(ScmCredentials(1))];
631-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
631+
let mut cmsg_space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmCredentials(1))];
632+
let mut cmsg_buffer = RecvAncillaryBuffer::from(cmsg_space.as_mut_slice());
632633

633634
let mut buffer = [0; BUFFER_SIZE];
634635
recvmsg(
@@ -685,8 +686,8 @@ fn test_unix_peercred_implicit() {
685686
)
686687
.unwrap();
687688

688-
let mut cmsg_space = [0; rustix::cmsg_space!(ScmCredentials(1))];
689-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
689+
let mut cmsg_space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmCredentials(1))];
690+
let mut cmsg_buffer = RecvAncillaryBuffer::from(cmsg_space.as_mut_slice());
690691

691692
let mut buffer = [0; BUFFER_SIZE];
692693
recvmsg(
@@ -737,13 +738,14 @@ fn test_unix_msg_with_combo() {
737738
let mut yet_another_pipe_end = None;
738739

739740
let mut buffer = [0; BUFFER_SIZE];
740-
let mut cmsg_space = [0; rustix::cmsg_space!(ScmRights(2), ScmRights(1))];
741+
let mut cmsg_space =
742+
[MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(2), ScmRights(1))];
741743

742744
'exit: loop {
743745
let data_socket = accept(&connection_socket).unwrap();
744746
let mut sum = 0;
745747
loop {
746-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
748+
let mut cmsg_buffer = RecvAncillaryBuffer::from(cmsg_space.as_mut_slice());
747749
let nread = recvmsg(
748750
&data_socket,
749751
&mut [IoSliceMut::new(&mut buffer)],
@@ -859,8 +861,8 @@ fn test_unix_msg_with_combo() {
859861

860862
let data_socket = socket(AddressFamily::UNIX, SocketType::SEQPACKET, None).unwrap();
861863

862-
let mut space = [0; rustix::cmsg_space!(ScmRights(2), ScmRights(1))];
863-
let mut cmsg_buffer = SendAncillaryBuffer::new(&mut space);
864+
let mut space = [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(2), ScmRights(1))];
865+
let mut cmsg_buffer = SendAncillaryBuffer::from(space.as_mut_slice());
864866

865867
// Format a CMSG.
866868
let we = [write_end.as_fd(), another_write_end.as_fd()];

‎tests/net/unix_alloc.rs

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -445,13 +445,14 @@ fn test_unix_msg_with_scm_rights() {
445445
let mut pipe_end = None;
446446

447447
let mut buffer = vec![0; BUFFER_SIZE];
448-
let mut cmsg_space = vec![0; rustix::cmsg_space!(ScmRights(1))];
448+
let mut cmsg_space = Vec::with_capacity(rustix::cmsg_space!(ScmRights(1)));
449449

450450
'exit: loop {
451451
let data_socket = accept(&connection_socket).unwrap();
452452
let mut sum = 0;
453453
loop {
454-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
454+
let mut cmsg_buffer =
455+
RecvAncillaryBuffer::from(cmsg_space.spare_capacity_mut());
455456
let nread = recvmsg(
456457
&data_socket,
457458
&mut [IoSliceMut::new(&mut buffer)],
@@ -553,8 +554,8 @@ fn test_unix_msg_with_scm_rights() {
553554
// Format the CMSG.
554555
let we = [write_end.as_fd()];
555556
let msg = SendAncillaryMessage::ScmRights(&we);
556-
let mut space = vec![0; msg.size()];
557-
let mut cmsg_buffer = SendAncillaryBuffer::new(&mut space);
557+
let mut space = Vec::with_capacity(msg.size());
558+
let mut cmsg_buffer = SendAncillaryBuffer::from(space.spare_capacity_mut());
558559
assert!(cmsg_buffer.push(msg));
559560

560561
connect_unix(&data_socket, &addr).unwrap();
@@ -618,8 +619,8 @@ fn test_unix_peercred() {
618619
assert_eq!(ucred.gid, getgid());
619620

620621
let msg = SendAncillaryMessage::ScmCredentials(ucred);
621-
let mut space = vec![0; msg.size()];
622-
let mut cmsg_buffer = SendAncillaryBuffer::new(&mut space);
622+
let mut space = Vec::with_capacity(msg.size());
623+
let mut cmsg_buffer = SendAncillaryBuffer::from(space.spare_capacity_mut());
623624
assert!(cmsg_buffer.push(msg));
624625

625626
sendmsg(
@@ -630,8 +631,8 @@ fn test_unix_peercred() {
630631
)
631632
.unwrap();
632633

633-
let mut cmsg_space = vec![0; rustix::cmsg_space!(ScmCredentials(1))];
634-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
634+
let mut cmsg_space = Vec::with_capacity(rustix::cmsg_space!(ScmCredentials(1)));
635+
let mut cmsg_buffer = RecvAncillaryBuffer::from(cmsg_space.spare_capacity_mut());
635636

636637
let mut buffer = vec![0; BUFFER_SIZE];
637638
recvmsg(
@@ -682,13 +683,15 @@ fn test_unix_msg_with_combo() {
682683
let mut yet_another_pipe_end = None;
683684

684685
let mut buffer = vec![0; BUFFER_SIZE];
685-
let mut cmsg_space = vec![0; rustix::cmsg_space!(ScmRights(1), ScmRights(2))];
686+
let mut cmsg_space =
687+
Vec::with_capacity(rustix::cmsg_space!(ScmRights(1), ScmRights(2)));
686688

687689
'exit: loop {
688690
let data_socket = accept(&connection_socket).unwrap();
689691
let mut sum = 0;
690692
loop {
691-
let mut cmsg_buffer = RecvAncillaryBuffer::new(&mut cmsg_space);
693+
let mut cmsg_buffer =
694+
RecvAncillaryBuffer::from(cmsg_space.spare_capacity_mut());
692695
let nread = recvmsg(
693696
&data_socket,
694697
&mut [IoSliceMut::new(&mut buffer)],
@@ -804,8 +807,8 @@ fn test_unix_msg_with_combo() {
804807

805808
let data_socket = socket(AddressFamily::UNIX, SocketType::SEQPACKET, None).unwrap();
806809

807-
let mut space = vec![0; rustix::cmsg_space!(ScmRights(1), ScmRights(2))];
808-
let mut cmsg_buffer = SendAncillaryBuffer::new(&mut space);
810+
let mut space = Vec::with_capacity(rustix::cmsg_space!(ScmRights(1), ScmRights(2)));
811+
let mut cmsg_buffer = SendAncillaryBuffer::from(space.spare_capacity_mut());
809812

810813
// Format a CMSG.
811814
let we = [write_end.as_fd(), another_write_end.as_fd()];

0 commit comments

Comments
 (0)