diff --git a/tokio/src/net/udp/socket.rs b/tokio/src/net/udp/socket.rs index 8c1bcf42d23..b2d29ba8041 100644 --- a/tokio/src/net/udp/socket.rs +++ b/tokio/src/net/udp/socket.rs @@ -1,4 +1,4 @@ -use crate::io::PollEvented; +use crate::io::{PollEvented, ReadBuf}; use crate::net::{to_socket_addrs, ToSocketAddrs}; use std::convert::TryFrom; @@ -355,15 +355,23 @@ impl UdpSocket { /// This function may encounter any standard I/O error except `WouldBlock`. /// /// [`connect`]: method@Self::connect - pub fn poll_recv(&self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll> { + pub fn poll_recv( + &self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { loop { let ev = ready!(self.io.poll_read_ready(cx))?; - match self.io.get_ref().recv(buf) { + match self.io.get_ref().recv(buf.initialize_unfilled()) { Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.io.clear_readiness(ev); } - x => return Poll::Ready(x), + Err(e) => return Poll::Ready(Err(e)), + Ok(n) => { + buf.advance(n); + return Poll::Ready(Ok(n)); + } } } } @@ -514,16 +522,20 @@ impl UdpSocket { pub fn poll_recv_from( &self, cx: &mut Context<'_>, - buf: &mut [u8], + buf: &mut ReadBuf<'_>, ) -> Poll> { loop { let ev = ready!(self.io.poll_read_ready(cx))?; - match self.io.get_ref().recv_from(buf) { + match self.io.get_ref().recv_from(buf.initialize_unfilled()) { Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { self.io.clear_readiness(ev); } - x => return Poll::Ready(x), + Err(e) => return Poll::Ready(Err(e)), + Ok(n) => { + buf.advance(n); + return Poll::Ready(Ok(n)); + } } } }