udp: Fix UdpFramed decode (#1517)

tokio-udp/src/frame.rs

@@ -34,6 +34,7 @@ pub struct UdpFramed<C> {
     out_addr: SocketAddr,
     flushed: bool,
     is_readable: bool,
+    repeat_decode: bool,
     current_addr: Option<SocketAddr>,
 }
 
@@ -44,36 +45,58 @@ impl<C: Decoder> Stream for UdpFramed<C> {
     fn poll(&mut self) -> Poll<Option<(Self::Item)>, 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<C> UdpFramed<C> {
             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<C> {
+        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,
         }
     }

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