UDP 0-length datagrams

Resolves #248, preventing UdpFramed from sending 0-length datagrams.
Also, adds 0-length tests for UdpSocket.

1 files changed, 145 insertions, 23 deletions

diff --git a/tests/udp.rs b/tests/udp.rs
index c5b0ef57..f22266d3 100644
--- a/tests/udp.rs
+++ b/tests/udp.rs
@@ -5,8 +5,8 @@ extern crate tokio_core;
 use std::io;
 use std::net::SocketAddr;
 
-use futures::{Future, Poll};
-use tokio_core::net::UdpSocket;
+use futures::{Future, Poll, Stream, Sink};
+use tokio_core::net::{UdpSocket, UdpCodec};
 use tokio_core::reactor::Core;
 
 macro_rules! t {
@@ -19,47 +19,169 @@ macro_rules! t {
 #[test]
 fn send_messages() {
     let mut l = t!(Core::new());
-    let a = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
-    let b = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
+    let mut a = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
+    let mut b = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
     let a_addr = t!(a.local_addr());
     let b_addr = t!(b.local_addr());
 
-    let send = SendMessage { socket: a, addr: b_addr };
-    let recv = RecvMessage { socket: b, expected_addr: a_addr };
-    t!(l.run(send.join(recv)));
+    {
+        let send = SendMessage::new(a, b_addr, b"1234");
+        let recv = RecvMessage::new(b, a_addr, b"1234");
+        let (sendt, received) = t!(l.run(send.join(recv)));
+        a = sendt;
+        b = received;
+    }
+
+    {
+        let send = SendMessage::new(a, b_addr, b"");
+        let recv = RecvMessage::new(b, a_addr, b"");
+        t!(l.run(send.join(recv)));
+    }
 }
 
 struct SendMessage {
-    socket: UdpSocket,
+    socket: Option<UdpSocket>,
     addr: SocketAddr,
+    data: &'static [u8],
+}
+
+impl SendMessage {
+    fn new(socket: UdpSocket, addr: SocketAddr, data: &'static [u8]) -> SendMessage {
+        SendMessage {
+            socket: Some(socket),
+            addr: addr,
+            data: data,
+        }
+    }
 }
 
 impl Future for SendMessage {
-    type Item = ();
+    type Item = UdpSocket;
     type Error = io::Error;
 
-    fn poll(&mut self) -> Poll<(), io::Error> {
-        let n = try_nb!(self.socket.send_to(b"1234", &self.addr));
-        assert_eq!(n, 4);
-        Ok(().into())
+    fn poll(&mut self) -> Poll<UdpSocket, io::Error> {
+        let n = try_nb!(self.socket.as_ref().unwrap()
+            .send_to(&self.data[..], &self.addr));
+
+        assert_eq!(n, self.data.len());
+
+        Ok(self.socket.take().unwrap().into())
     }
 }
 
 struct RecvMessage {
-    socket: UdpSocket,
-    expected_addr: SocketAddr,
+    socket: Option<UdpSocket>,
+    addr: SocketAddr,
+    data: &'static [u8],
+}
+
+impl RecvMessage {
+    fn new(socket: UdpSocket, expected_addr: SocketAddr,
+           expected_data: &'static [u8]) -> RecvMessage
+    {
+        RecvMessage {
+            socket: Some(socket),
+            addr: expected_addr,
+            data: expected_data,
+        }
+    }
 }
 
 impl Future for RecvMessage {
-    type Item = ();
+    type Item = UdpSocket;
     type Error = io::Error;
 
-    fn poll(&mut self) -> Poll<(), io::Error> {
-        let mut buf = [0; 32];
-        let (n, addr) = try_nb!(self.socket.recv_from(&mut buf));
-        assert_eq!(n, 4);
-        assert_eq!(&buf[..4], b"1234");
-        assert_eq!(addr, self.expected_addr);
-        Ok(().into())
+    fn poll(&mut self) -> Poll<UdpSocket, io::Error> {
+        let mut buf = vec![0u8; 10 + self.data.len() * 10];
+        let (n, addr) = try_nb!(self.socket.as_ref().unwrap()
+            .recv_from(&mut buf[..]));
+
+        assert_eq!(n, self.data.len());
+        assert_eq!(&buf[..self.data.len()], &self.data[..]);
+        assert_eq!(addr, self.addr);
+
+        Ok(self.socket.take().unwrap().into())
+    }
+}
+
+#[test]
+fn send_dgrams() {
+    let mut l = t!(Core::new());
+    let mut a = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
+    let mut b = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
+    let mut buf = [0u8; 50];
+    let b_addr = t!(b.local_addr());
+
+    {
+        let send = a.send_dgram(&b"4321"[..], b_addr);
+        let recv = b.recv_dgram(&mut buf[..]);
+        let (sendt, received) = t!(l.run(send.join(recv)));
+        assert_eq!(received.2, 4);
+        assert_eq!(&received.1[..4], b"4321");
+        a = sendt.0;
+        b = received.0;
+    }
+
+    {
+        let send = a.send_dgram(&b""[..], b_addr);
+        let recv = b.recv_dgram(&mut buf[..]);
+        let received = t!(l.run(send.join(recv))).1;
+        assert_eq!(received.2, 0);
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Codec {
+    data: &'static [u8],
+    from: SocketAddr,
+    to: SocketAddr,
+}
+
+impl UdpCodec for Codec {
+    type In = ();
+    type Out = &'static [u8];
+
+    fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> io::Result<Self::In> {
+        assert_eq!(src, &self.from);
+        assert_eq!(buf, self.data);
+        Ok(())
+    }
+
+    fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
+        assert_eq!(msg, self.data);
+        buf.extend_from_slice(msg);
+        self.to
+    }
+}
+
+#[test]
+fn send_framed() {
+    let mut l = t!(Core::new());
+    let mut a_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
+    let mut b_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
+    let a_addr = t!(a_soc.local_addr());
+    let b_addr = t!(b_soc.local_addr());
+
+    {
+        let a = a_soc.framed(Codec { data: &b"4567"[..], from: a_addr, to: b_addr});
+        let b = b_soc.framed(Codec { data: &b"4567"[..], from: a_addr, to: b_addr});
+
+        let send = a.send(&b"4567"[..]);
+        let recv = b.into_future().map_err(|e| e.0);
+        let (sendt, received) = t!(l.run(send.join(recv)));
+        assert_eq!(received.0, Some(()));
+
+        a_soc = sendt.into_inner();
+        b_soc = received.1.into_inner();
+    }
+
+    {
+        let a = a_soc.framed(Codec { data: &b""[..], from: a_addr, to: b_addr});
+        let b = b_soc.framed(Codec { data: &b""[..], from: a_addr, to: b_addr});
+
+        let send = a.send(&b""[..]);
+        let recv = b.into_future().map_err(|e| e.0);
+        let received = t!(l.run(send.join(recv))).1;
+        assert_eq!(received.0, Some(()));
     }
 }