Skip to content

Commit d24ce51

Browse files
George-MiaoGeorge-Miao
authored
Merge pull request #283 from AsakuraMizu/master
fix(driver/net): ancillary data
2 parents 84ed77b + 970f328 commit d24ce51

File tree

6 files changed

+69
-33
lines changed

6 files changed

+69
-33
lines changed

‎compio-driver/src/iocp/op.rs

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -786,6 +786,7 @@ pub struct RecvMsg<T: IoVectoredBufMut, C: IoBufMut, S> {
786786
fd: SharedFd<S>,
787787
buffer: T,
788788
control: C,
789+
control_len: u32,
789790
_p: PhantomPinned,
790791
}
791792

@@ -806,16 +807,22 @@ impl<T: IoVectoredBufMut, C: IoBufMut, S> RecvMsg<T, C, S> {
806807
fd,
807808
buffer,
808809
control,
810+
control_len: 0,
809811
_p: PhantomPinned,
810812
}
811813
}
812814
}
813815

814816
impl<T: IoVectoredBufMut, C: IoBufMut, S> IntoInner for RecvMsg<T, C, S> {
815-
type Inner = ((T, C), SOCKADDR_STORAGE, socklen_t);
817+
type Inner = ((T, C), SOCKADDR_STORAGE, socklen_t, usize);
816818

817819
fn into_inner(self) -> Self::Inner {
818-
((self.buffer, self.control), self.addr, self.addr_len)
820+
(
821+
(self.buffer, self.control),
822+
self.addr,
823+
self.addr_len,
824+
self.control_len as _,
825+
)
819826
}
820827
}
821828

@@ -828,7 +835,7 @@ impl<T: IoVectoredBufMut, C: IoBufMut, S: AsRawFd> OpCode for RecvMsg<T, C, S> {
828835
})?;
829836

830837
let this = self.get_unchecked_mut();
831-
let mut slices = this.buffer.as_io_slices_mut();
838+
let mut slices = this.buffer.io_slices_mut();
832839
let mut msg = WSAMSG {
833840
name: &mut this.addr as *mut _ as _,
834841
namelen: this.addr_len,
@@ -837,6 +844,7 @@ impl<T: IoVectoredBufMut, C: IoBufMut, S: AsRawFd> OpCode for RecvMsg<T, C, S> {
837844
Control: std::mem::transmute::<IoSliceMut, WSABUF>(this.control.as_io_slice_mut()),
838845
dwFlags: 0,
839846
};
847+
this.control_len = 0;
840848

841849
let mut received = 0;
842850
let res = recvmsg_fn(
@@ -846,6 +854,7 @@ impl<T: IoVectoredBufMut, C: IoBufMut, S: AsRawFd> OpCode for RecvMsg<T, C, S> {
846854
optr,
847855
None,
848856
);
857+
this.control_len = msg.Control.len;
849858
winsock_result(res, received)
850859
}
851860

@@ -897,7 +906,7 @@ impl<T: IoVectoredBuf, C: IoBuf, S: AsRawFd> OpCode for SendMsg<T, C, S> {
897906
unsafe fn operate(self: Pin<&mut Self>, optr: *mut OVERLAPPED) -> Poll<io::Result<usize>> {
898907
let this = self.get_unchecked_mut();
899908

900-
let slices = this.buffer.as_io_slices();
909+
let slices = this.buffer.io_slices();
901910
let msg = WSAMSG {
902911
name: this.addr.as_ptr() as _,
903912
namelen: this.addr.len(),

‎compio-driver/src/op.rs

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,23 +49,47 @@ impl<T: SetBufInit, O> BufResultExt for BufResult<(usize, O), T> {
4949
}
5050
}
5151

52-
/// Helper trait for [`RecvFrom`] and [`RecvFromVectored`].
52+
impl<T: SetBufInit, C: SetBufInit, O> BufResultExt for BufResult<(usize, usize, O), (T, C)> {
53+
fn map_advanced(self) -> Self {
54+
self.map(
55+
|(init_buffer, init_control, obj), (mut buffer, mut control)| {
56+
unsafe {
57+
buffer.set_buf_init(init_buffer);
58+
control.set_buf_init(init_control);
59+
}
60+
((init_buffer, init_control, obj), (buffer, control))
61+
},
62+
)
63+
}
64+
}
65+
66+
/// Helper trait for [`RecvFrom`], [`RecvFromVectored`] and [`RecvMsg`].
5367
pub trait RecvResultExt {
5468
/// The mapped result.
55-
type RecvFromResult;
69+
type RecvResult;
5670

5771
/// Create [`SockAddr`] if the result is [`Ok`].
58-
fn map_addr(self) -> Self::RecvFromResult;
72+
fn map_addr(self) -> Self::RecvResult;
5973
}
6074

6175
impl<T> RecvResultExt for BufResult<usize, (T, sockaddr_storage, socklen_t)> {
62-
type RecvFromResult = BufResult<(usize, SockAddr), T>;
76+
type RecvResult = BufResult<(usize, SockAddr), T>;
77+
78+
fn map_addr(self) -> Self::RecvResult {
79+
self.map_buffer(|(buffer, addr_buffer, addr_size)| (buffer, addr_buffer, addr_size, 0))
80+
.map_addr()
81+
.map_res(|(res, _, addr)| (res, addr))
82+
}
83+
}
84+
85+
impl<T> RecvResultExt for BufResult<usize, (T, sockaddr_storage, socklen_t, usize)> {
86+
type RecvResult = BufResult<(usize, usize, SockAddr), T>;
6387

64-
fn map_addr(self) -> Self::RecvFromResult {
88+
fn map_addr(self) -> Self::RecvResult {
6589
self.map2(
66-
|res, (buffer, addr_buffer, addr_size)| {
90+
|res, (buffer, addr_buffer, addr_size, len)| {
6791
let addr = unsafe { SockAddr::new(addr_buffer, addr_size) };
68-
((res, addr), buffer)
92+
((res, len, addr), buffer)
6993
},
7094
|(buffer, ..)| buffer,
7195
)

‎compio-driver/src/unix/op.rs

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -404,22 +404,27 @@ impl<T: IoVectoredBufMut, C: IoBufMut, S> RecvMsg<T, C, S> {
404404
}
405405

406406
pub(crate) unsafe fn set_msg(&mut self) {
407-
self.slices = self.buffer.as_io_slices_mut();
407+
self.slices = self.buffer.io_slices_mut();
408408

409409
self.msg.msg_name = std::ptr::addr_of_mut!(self.addr) as _;
410410
self.msg.msg_namelen = std::mem::size_of_val(&self.addr) as _;
411411
self.msg.msg_iov = self.slices.as_mut_ptr() as _;
412412
self.msg.msg_iovlen = self.slices.len() as _;
413413
self.msg.msg_control = self.control.as_buf_mut_ptr() as _;
414-
self.msg.msg_controllen = self.control.buf_len() as _;
414+
self.msg.msg_controllen = self.control.buf_capacity() as _;
415415
}
416416
}
417417

418418
impl<T: IoVectoredBufMut, C: IoBufMut, S> IntoInner for RecvMsg<T, C, S> {
419-
type Inner = ((T, C), sockaddr_storage, socklen_t);
419+
type Inner = ((T, C), sockaddr_storage, socklen_t, usize);
420420

421421
fn into_inner(self) -> Self::Inner {
422-
((self.buffer, self.control), self.addr, self.msg.msg_namelen)
422+
(
423+
(self.buffer, self.control),
424+
self.addr,
425+
self.msg.msg_namelen,
426+
self.msg.msg_controllen as _,
427+
)
423428
}
424429
}
425430

@@ -458,7 +463,7 @@ impl<T: IoVectoredBuf, C: IoBuf, S> SendMsg<T, C, S> {
458463
}
459464

460465
pub(crate) unsafe fn set_msg(&mut self) {
461-
self.slices = self.buffer.as_io_slices();
466+
self.slices = self.buffer.io_slices();
462467

463468
self.msg.msg_name = self.addr.as_ptr() as _;
464469
self.msg.msg_namelen = self.addr.len();

‎compio-net/src/socket.rs

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,7 @@ impl Socket {
261261
&self,
262262
buffer: T,
263263
control: C,
264-
) -> BufResult<(usize, SockAddr), (T, C)> {
264+
) -> BufResult<(usize, usize, SockAddr), (T, C)> {
265265
self.recv_msg_vectored([buffer], control)
266266
.await
267267
.map_buffer(|([buffer], control)| (buffer, control))
@@ -271,20 +271,14 @@ impl Socket {
271271
&self,
272272
buffer: T,
273273
control: C,
274-
) -> BufResult<(usize, SockAddr), (T, C)> {
274+
) -> BufResult<(usize, usize, SockAddr), (T, C)> {
275275
let fd = self.to_shared_fd();
276276
let op = RecvMsg::new(fd, buffer, control);
277277
compio_runtime::submit(op)
278278
.await
279279
.into_inner()
280280
.map_addr()
281-
.map(|(init, obj), (mut buffer, control)| {
282-
// SAFETY: The number of bytes received would not bypass the buffer capacity.
283-
unsafe {
284-
buffer.set_buf_init(init);
285-
}
286-
((init, obj), (buffer, control))
287-
})
281+
.map_advanced()
288282
}
289283

290284
pub async fn send_to<T: IoBuf>(&self, buffer: T, addr: &SockAddr) -> BufResult<usize, T> {

‎compio-net/src/udp.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -228,11 +228,11 @@ impl UdpSocket {
228228
&self,
229229
buffer: T,
230230
control: C,
231-
) -> BufResult<(usize, SocketAddr), (T, C)> {
231+
) -> BufResult<(usize, usize, SocketAddr), (T, C)> {
232232
self.inner
233233
.recv_msg(buffer, control)
234234
.await
235-
.map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
235+
.map_res(|(n, m, addr)| (n, m, addr.as_socket().expect("should be SocketAddr")))
236236
}
237237

238238
/// Receives a single datagram message and ancillary data on the socket. On
@@ -241,11 +241,11 @@ impl UdpSocket {
241241
&self,
242242
buffer: T,
243243
control: C,
244-
) -> BufResult<(usize, SocketAddr), (T, C)> {
244+
) -> BufResult<(usize, usize, SocketAddr), (T, C)> {
245245
self.inner
246246
.recv_msg_vectored(buffer, control)
247247
.await
248-
.map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
248+
.map_res(|(n, m, addr)| (n, m, addr.as_socket().expect("should be SocketAddr")))
249249
}
250250

251251
/// Sends data on the socket to the given address. On success, returns the

‎compio-net/tests/udp.rs

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -103,14 +103,18 @@ async fn send_msg_with_ipv6_ecn() {
103103

104104
active.send_msg(MSG, control, passive_addr).await.unwrap();
105105

106-
let res = passive.recv_msg(Vec::with_capacity(20), [0u8; 32]).await;
107-
assert_eq!(res.0.unwrap().1, active_addr);
108-
assert_eq!(res.1.0, MSG.as_bytes());
106+
let ((_, _, addr), (buffer, control)) = passive
107+
.recv_msg(Vec::with_capacity(20), Vec::with_capacity(32))
108+
.await
109+
.unwrap();
110+
assert_eq!(addr, active_addr);
111+
assert_eq!(buffer, MSG.as_bytes());
109112
unsafe {
110-
let mut iter = CMsgIter::new(&res.1.1);
113+
let mut iter = CMsgIter::new(&control);
111114
let cmsg = iter.next().unwrap();
112115
assert_eq!(cmsg.level(), IPPROTO_IPV6);
113116
assert_eq!(cmsg.ty(), IPV6_TCLASS);
114117
assert_eq!(cmsg.data::<i32>(), &ECN_BITS);
118+
assert!(iter.next().is_none());
115119
}
116120
}

0 commit comments

Comments
 (0)