v0.1.x: Fix UdpFramed with regards to Decode (#1444)

* add test for using LinesCodec with UdpFramed

* fix UdpFramed decode

* rustfmt

Author: JohnDoneth

tokio-udp/src/frame.rs

@@ -33,6 +33,8 @@ pub struct UdpFramed<C> {
     wr: BytesMut,
     out_addr: SocketAddr,
     flushed: bool,
+    is_readable: bool,
+    current_addr: Option<SocketAddr>,
 }
 
 impl<C: Decoder> Stream for UdpFramed<C> {
@@ -42,19 +44,37 @@ impl<C: Decoder> Stream for UdpFramed<C> {
     fn poll(&mut self) -> Poll<Option<(Self::Item)>, 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<C> UdpFramed<C> {
             rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY),
             wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY),
             flushed: true,
+            is_readable: false,
+            current_addr: None,
         }
     }
 

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)));
+}