From c3c3481d74ed9c83dd86fc62a418e28b17e55e84 Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Fri, 30 Aug 2019 14:13:54 -0400 Subject: [PATCH] udp: Fix UdpFramed decode (#1517) --- tokio-udp/src/frame.rs | 78 ++++++++++++++++++++++++++++++++---------- tokio-udp/tests/udp.rs | 28 ++++++++++++++- 2 files changed, 87 insertions(+), 19 deletions(-) diff --git a/tokio-udp/src/frame.rs b/tokio-udp/src/frame.rs index ca52717fe..d9b4c46de 100644 --- a/tokio-udp/src/frame.rs +++ b/tokio-udp/src/frame.rs @@ -34,6 +34,7 @@ pub struct UdpFramed { out_addr: SocketAddr, flushed: bool, is_readable: bool, + repeat_decode: bool, current_addr: Option, } @@ -44,36 +45,58 @@ impl Stream for UdpFramed { fn poll(&mut self) -> Poll, Self::Error> { self.rd.reserve(INITIAL_RD_CAPACITY); - loop { - // Are there are still bytes left in the read buffer to decode? - if self.is_readable { - if let Some(frame) = self.codec.decode(&mut self.rd)? { - trace!("frame decoded from buffer"); + if self.repeat_decode { + loop { + // Are there are still bytes left in the read buffer to decode? + if self.is_readable { + // Use deocde_eof since every datagram contains its own + // eof which is just the end of the datagram. This supports + // the lines use case where there may not be a terminating + // delimiter and thus you may never get the end of the frame. + // This is generally fine for most implementations of codec + // since by default this will defer to calling decode. + if let Some(frame) = self.codec.decode_eof(&mut self.rd)? { + trace!("frame decoded from buffer"); - let current_addr = self - .current_addr - .expect("will always be set before this line is called"); + let current_addr = self + .current_addr + .expect("will always be set before this line is called"); - return Ok(Async::Ready(Some((frame, current_addr)))); + return Ok(Async::Ready(Some((frame, current_addr)))); + } + + // if this line has been reached then decode has returned `None`. + self.is_readable = false; + self.rd.clear(); } - // if this line has been reached then decode has returned `None`. - self.is_readable = false; - self.rd.clear(); - } + // We're out of data. Try and fetch more data to decode + let (n, addr) = unsafe { + // Read into the buffer without having to initialize the memory. + let (n, addr) = try_ready!(self.socket.poll_recv_from(self.rd.bytes_mut())); + self.rd.advance_mut(n); + (n, addr) + }; - // We're out of data. Try and fetch more data to decode + self.current_addr = Some(addr); + self.is_readable = true; + + trace!("received {} bytes, decoding", n); + } + } else { let (n, addr) = unsafe { // Read into the buffer without having to initialize the memory. let (n, addr) = try_ready!(self.socket.poll_recv_from(self.rd.bytes_mut())); self.rd.advance_mut(n); (n, addr) }; - - self.current_addr = Some(addr); - self.is_readable = true; - trace!("received {} bytes, decoding", n); + let frame_res = self.codec.decode(&mut self.rd); + self.rd.clear(); + let frame = frame_res?; + let result = frame.map(|frame| (frame, addr)); // frame -> (frame, addr) + trace!("frame decoded from buffer"); + Ok(Async::Ready(result)) } } } @@ -147,6 +170,25 @@ impl UdpFramed { wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY), flushed: true, is_readable: false, + repeat_decode: false, + current_addr: None, + } + } + + /// Create a new `UdpFramed` backed by the given socket and codec. That will + /// continue to call `decode_eof` until the decoder has cleared the entire buffer. + /// + /// See struct level documentation for more details. + pub fn with_decode(socket: UdpSocket, codec: C, repeat_decode: bool) -> UdpFramed { + UdpFramed { + socket: socket, + codec: codec, + out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)), + rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY), + wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY), + flushed: true, + is_readable: false, + repeat_decode, current_addr: None, } } diff --git a/tokio-udp/tests/udp.rs b/tokio-udp/tests/udp.rs index 7215977ac..f63fc9ad3 100644 --- a/tokio-udp/tests/udp.rs +++ b/tokio-udp/tests/udp.rs @@ -301,7 +301,7 @@ fn send_framed_lines_codec() { let b_addr = t!(b_soc.local_addr()); let a = UdpFramed::new(a_soc, ByteCodec); - let b = UdpFramed::new(b_soc, LinesCodec::new()); + let b = UdpFramed::with_decode(b_soc, LinesCodec::new(), true); let msg = b"1\r\n2\r\n3\r\n".to_vec(); @@ -314,3 +314,29 @@ fn send_framed_lines_codec() { assert_eq!(recv.next(), Some(("2".to_string(), a_addr))); assert_eq!(recv.next(), Some(("3".to_string(), a_addr))); } + +#[test] +fn send_framed_lines_codec_with_non_terminating_frame() { + drop(env_logger::try_init()); + + let a_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()))); + let b_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()))); + let a_addr = t!(a_soc.local_addr()); + let b_addr = t!(b_soc.local_addr()); + + let a = UdpFramed::new(a_soc, ByteCodec); + let b = UdpFramed::with_decode(b_soc, LinesCodec::new(), true); + + // This has no terminating delimiter thus we want to return the rest of the + // frame and this tests that if decode fails, we try to decode_eof. + let msg = b"1\r\n2\r\n3".to_vec(); + + let send = a.send((msg.clone(), b_addr)); + t!(send.wait()); + + let mut recv = Stream::wait(b).map(|e| e.unwrap()); + + assert_eq!(recv.next(), Some(("1".to_string(), a_addr))); + assert_eq!(recv.next(), Some(("2".to_string(), a_addr))); + assert_eq!(recv.next(), Some(("3".to_string(), a_addr))); +}