diff --git a/tokio-udp/src/frame.rs b/tokio-udp/src/frame.rs index b7ae49010..ca52717fe 100644 --- a/tokio-udp/src/frame.rs +++ b/tokio-udp/src/frame.rs @@ -33,6 +33,8 @@ pub struct UdpFramed { wr: BytesMut, out_addr: SocketAddr, flushed: bool, + is_readable: bool, + current_addr: Option, } impl Stream for UdpFramed { @@ -42,19 +44,37 @@ impl Stream for UdpFramed { fn poll(&mut self) -> Poll, Self::Error> { self.rd.reserve(INITIAL_RD_CAPACITY); - 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) - }; - 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)) + 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"); + + let current_addr = self + .current_addr + .expect("will always be set before this line is called"); + + 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(); + } + + // 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) + }; + + self.current_addr = Some(addr); + self.is_readable = true; + + trace!("received {} bytes, decoding", n); + } } } @@ -126,6 +146,8 @@ impl UdpFramed { rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY), wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY), flushed: true, + is_readable: false, + current_addr: None, } } diff --git a/tokio-udp/tests/udp.rs b/tokio-udp/tests/udp.rs index aadd15ca8..a7b018bff 100644 --- a/tokio-udp/tests/udp.rs +++ b/tokio-udp/tests/udp.rs @@ -12,7 +12,7 @@ use std::net::SocketAddr; use futures::{Future, Poll, Sink, Stream}; use bytes::{BufMut, BytesMut}; -use tokio_codec::{Decoder, Encoder}; +use tokio_codec::{Decoder, Encoder, LinesCodec}; use tokio_udp::{UdpFramed, UdpSocket}; macro_rules! t { @@ -247,7 +247,7 @@ impl Encoder for ByteCodec { } #[test] -fn send_framed() { +fn send_framed_byte_codec() { drop(env_logger::try_init()); let mut a_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()))); @@ -288,3 +288,27 @@ fn send_framed() { assert_eq!(a_addr, addr); } } + +#[test] +fn send_framed_lines_codec() { + 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::new(b_soc, LinesCodec::new()); + + let msg = b"1\r\n2\r\n3\r\n".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))); +}