Touch up codes for UDP

* Move to `std::net` as it's all purely UDP related
* Rename to `UdpCodec` and `UdpFramed` to give a consistent `Udp` prefix
* Add `RecvDgram`, rename `SendDGram` to `SendDgram`
* Touch up some style here and there

examples/udp-codec.rs

@@ -8,8 +8,7 @@ extern crate log;
 use std::io;
 use std::net::{SocketAddr};
 use futures::{future, Future, Stream, Sink};
-use tokio_core::io::{CodecUdp};
+use tokio_core::net::{UdpSocket, UdpCodec};
-use tokio_core::net::{UdpSocket};
 use tokio_core::reactor::{Core, Timeout};
 use std::time::Duration;
 use std::str;
@@ -32,7 +31,7 @@ pub struct LineCodec {
     addr : Option<SocketAddr>
 }
 
-impl CodecUdp for LineCodec {
+impl UdpCodec for LineCodec {
     type In = Vec<Vec<u8>>;
     type Out = Vec<u8>;
 
@@ -49,7 +48,7 @@ impl CodecUdp for LineCodec {
         }
     }
 
-    fn encode(&mut self, item: &Vec<u8>, into: &mut Vec<u8>) -> SocketAddr {
+    fn encode(&mut self, item: Vec<u8>, into: &mut Vec<u8>) -> SocketAddr {
         trace!("encoding {}", str::from_utf8(item.as_slice()).unwrap());
         into.extend_from_slice(item.as_slice());
         into.push('\n' as u8);
@@ -84,7 +83,7 @@ fn main() {
     //We create a FramedUdp instance, which associates a socket
     //with a codec. We then immediate split that into the
     //receiving side `Stream` and the writing side `Sink`
-    let (srvstream, srvsink) = server.framed(srvcodec).split();
+    let (srvsink, srvstream) = server.framed(srvcodec).split();
 
     //`Stream::fold` runs once per every received datagram.
     //Note that we pass srvsink into fold, so that it can be
@@ -96,7 +95,7 @@ fn main() {
     }).map(|_| ());
 
     //We create another FramedUdp instance, this time for the client socket
-    let (clistream, clisink) = client.framed(clicodec).split();
+    let (clisink, clistream) = client.framed(clicodec).split();
 
     //And another infinite iteration
     let cliloop = clistream.fold(clisink, move |sink, lines| {

src/io/mod.rs

@@ -33,7 +33,6 @@ macro_rules! try_nb {
 
 mod copy;
 mod frame;
-mod udp_frame;
 mod flush;
 mod read_exact;
 mod read_to_end;
@@ -44,7 +43,6 @@ mod window;
 mod write_all;
 pub use self::copy::{copy, Copy};
 pub use self::frame::{EasyBuf, EasyBufMut, Framed, Codec};
-pub use self::udp_frame::{FramedUdp, framed_udp, FramedUdpRead, FramedUdpWrite, CodecUdp, VecDGramCodec};
 pub use self::flush::{flush, Flush};
 pub use self::read_exact::{read_exact, ReadExact};
 pub use self::read_to_end::{read_to_end, ReadToEnd};

src/io/udp_frame.rs

@@ -1,262 +0,0 @@
-use std::io;
-use std::net::SocketAddr;
-use net::UdpSocket;
-use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};
-use futures::sync::BiLock;
-
-/// Encoding of frames via buffers.
-///
-/// This trait is used when constructing an instance of `FramedUdp`. It provides
-/// one type: `Out` for encoding outgoing frames according to a protocol.
-///
-/// Because UDP is a connectionless protocol, the encode method will also be
-/// responsible for determining the remote host to which the datagram should be
-/// sent
-///
-/// The trait itself is implemented on a type that can track state for decoding
-/// or encoding, which is particularly useful for streaming parsers. In many
-/// cases, though, this type will simply be a unit struct (e.g. `struct
-/// HttpCodec`).
-pub trait CodecUdp {
-
-    /// The type of frames to be encoded.
-    type Out;
-    
-    /// The type of decoded frames.
-    type In;
-
-
-    /// Encodes a frame into the buffer provided.
-    ///
-    /// This method will encode `msg` into the byte buffer provided by `buf`.
-    /// The `buf` provided is an internal buffer of the `Framed` instance and
-    /// will be written out when possible. 
-    ///
-    /// The encode method also determines the destination to which the buffer should
-    /// be directed, which will be returned as a SocketAddr;
-    fn encode(&mut self, msg: &Self::Out, buf: &mut Vec<u8>) -> SocketAddr;
-    
-    /// Attempts to decode a frame from the provided buffer of bytes.
-    ///
-    /// This method is called by `FramedUdp` on a single datagram which has been
-    /// read from a socket. 
-    ///
-    /// It is required that the decode method empty the read buffer in every call to
-    /// decode, as the next poll_read that occurs will write the next datagram
-    /// into the buffer, without regard for what is already there. 
-    ///
-    /// Finally, if the bytes in the buffer are malformed then an error is
-    /// returned indicating why. This informs `Framed` that the stream is now
-    /// corrupt and should be terminated.
-    ///
-    fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> Result<Self::In, io::Error>;
-}
-
-/// A unified `Stream` and `Sink` interface to an underlying `Io` object, using
-/// the `CodecUdp` trait to encode and decode frames.
-///
-/// You can acquire a `Framed` instance by using the `Io::framed` adapter.
-pub struct FramedUdp<C> {
-    socket: UdpSocket,
-    codec: C,
-    out_addr : Option<SocketAddr>,
-    rd: Vec<u8>,
-    wr: Vec<u8>,
-}
-
-impl<C : CodecUdp> Stream for FramedUdp<C> {
-    type Item = C::In;
-    type Error = io::Error;
-
-    fn poll(&mut self) -> Poll<Option<C::In>, io::Error> {
-        loop {
-
-            let ret = self.socket.recv_from(self.rd.as_mut_slice());
-            match ret {
-                Ok((n, addr)) => { 
-                    trace!("read {} bytes", n);
-                    trace!("attempting to decode a frame");
-                    let frame = try!(self.codec.decode(&addr, & self.rd[.. n]));
-                    trace!("frame decoded from buffer");
-                    return Ok(Async::Ready(Some(frame)));
-                }
-                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                    return Ok(Async::NotReady)
-                }
-                Err(e) => return Err(e),
-            }
-        }
-    }
-}
-
-impl<C : CodecUdp> Sink for FramedUdp<C> {
-    type SinkItem = C::Out;
-    type SinkError = io::Error;
-
-    fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
-        if self.wr.len() > 0 {
-            try!(self.poll_complete());
-            if self.wr.len() > 0 {
-                return Ok(AsyncSink::NotReady(item));
-            }
-        }
-
-        self.out_addr = Some(self.codec.encode(&item, &mut self.wr));
-        Ok(AsyncSink::Ready)
-    }
-
-    fn poll_complete(&mut self) -> Poll<(), io::Error> {
-        trace!("flushing framed transport");
-
-        if !self.wr.is_empty() {
-            if let Some(outaddr) = self.out_addr {
-                let remaining = self.wr.len();
-                trace!("writing; remaining={}", remaining);
-                let n = try_nb!(self.socket.send_to(self.wr.as_slice(), &outaddr));
-                trace!("written {}", n);
-                self.wr.clear();
-                self.out_addr = None;
-                if n != remaining {
-                    return Err(io::Error::new(io::ErrorKind::Other,
-                                              "failed to write entire datagram to socket"));
-                }
-            }
-            else {
-                panic!("outbound stream in invalid state: out_addr is not known");
-            }
-        }
-
-        return Ok(Async::Ready(()));
-    }
-}
-
-/// Helper function that Creates a new FramedUdp object.  It moves the supplied socket, codec
-/// into the resulting FramedUdp
-pub fn framed_udp<C>(socket : UdpSocket, codec : C) -> FramedUdp<C> {
-    FramedUdp::new(
-        socket,
-        codec,
-        vec![0; 64 * 1024],
-        Vec::with_capacity(64 * 1024)
-    )
-}
-
-impl<C> FramedUdp<C> {
-
-    /// Creates a new FramedUdp object.  It moves the supplied socket, codec
-    /// supplied vecs. 
-    fn new(sock : UdpSocket, codec : C, rd_buf : Vec<u8>, wr_buf : Vec<u8>) -> FramedUdp<C> {
-        FramedUdp {
-            socket: sock,
-            codec : codec,
-            out_addr: None,
-            rd: rd_buf,
-            wr: wr_buf
-        }
-    }
-
-
-    /// Splits this `Stream + Sink` object into separate `Stream` and `Sink`
-    /// objects, which can be useful when you want to split ownership between
-    /// tasks, or allow direct interaction between the two objects (e.g. via
-    /// `Sink::send_all`).
-    pub fn split(self) -> (FramedUdpRead<C>, FramedUdpWrite<C>) {
-        let (a, b) = BiLock::new(self);
-        let read = FramedUdpRead { framed: a };
-        let write = FramedUdpWrite { framed: b };
-        (read, write)
-    }
-
-    /// Returns a reference to the underlying I/O stream wrapped by `Framed`.
-    ///
-    /// Note that care should be taken to not tamper with the underlying stream
-    /// of data coming in as it may corrupt the stream of frames otherwise being
-    /// worked with.
-    pub fn get_ref(&self) -> &UdpSocket {
-        &self.socket
-    }
-
-    /// Returns a mutable reference to the underlying I/O stream wrapped by
-    /// `Framed`.
-    ///
-    /// Note that care should be taken to not tamper with the underlying stream
-    /// of data coming in as it may corrupt the stream of frames otherwise being
-    /// worked with.
-    pub fn get_mut(&mut self) -> &mut UdpSocket {
-        &mut self.socket
-    }
-
-    /// Consumes the `Framed`, returning its underlying I/O stream.
-    ///
-    /// Note that care should be taken to not tamper with the underlying stream
-    /// of data coming in as it may corrupt the stream of frames otherwise being
-    /// worked with.
-    pub fn into_inner(self) -> UdpSocket {
-        self.socket
-    }
-}
-/// A `Stream` interface to an underlying `Io` object, using the `CodecUdp` trait
-/// to decode frames.
-pub struct FramedUdpRead<C> {
-    framed: BiLock<FramedUdp<C>>,
-}
-
-impl<C : CodecUdp> Stream for FramedUdpRead<C> {
-    type Item = C::In;
-    type Error = io::Error;
-
-    fn poll(&mut self) -> Poll<Option<C::In>, io::Error> {
-        if let Async::Ready(mut guard) = self.framed.poll_lock() {
-            guard.poll()
-        } else {
-            Ok(Async::NotReady)
-        }
-    }
-}
-
-/// A `Sink` interface to an underlying `Io` object, using the `CodecUdp` trait
-/// to encode frames.
-pub struct FramedUdpWrite<C> {
-    framed: BiLock<FramedUdp<C>>,
-}
-
-impl<C : CodecUdp> Sink for FramedUdpWrite<C> {
-    type SinkItem = C::Out;
-    type SinkError = io::Error;
-
-    fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
-        if let Async::Ready(mut guard) = self.framed.poll_lock() {
-            guard.start_send(item)
-        } else {
-            Ok(AsyncSink::NotReady(item))
-        }
-    }
-
-    fn poll_complete(&mut self) -> Poll<(), io::Error> {
-        if let Async::Ready(mut guard) = self.framed.poll_lock() {
-            guard.poll_complete()
-        } else {
-            Ok(Async::NotReady)
-        }
-    }
-}
-
-/// Default implementation of a DGram "parser"
-/// This receives and produces a tuple of ('SocketAddr', `Vec`<u8>)
-pub struct VecDGramCodec;
-
-impl CodecUdp for VecDGramCodec {
-    type In = (SocketAddr, Vec<u8>);
-    type Out = (SocketAddr, Vec<u8>);
-
-    fn decode(&mut self, addr : &SocketAddr, buf: &[u8]) -> Result<Self::In, io::Error> {
-        Ok((*addr, buf.into()))
-    }
-    
-    fn encode(&mut self, item: &Self::Out, into: &mut Vec<u8>) -> SocketAddr {
-        into.extend_from_slice(item.1.as_slice());
-        into.push('\n' as u8);
-        item.0
-    }
-}
-

src/net/mod.rs

@@ -8,4 +8,4 @@ mod udp;
 
 pub use self::tcp::{TcpStream, TcpStreamNew};
 pub use self::tcp::{TcpListener, Incoming};
-pub use self::udp::UdpSocket;
+pub use self::udp::{UdpSocket, UdpCodec, UdpFramed};

src/net/udp/frame.rs

@@ -0,0 +1,154 @@
+use std::io;
+use std::net::{SocketAddr, Ipv4Addr, SocketAddrV4};
+
+use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};
+
+use net::UdpSocket;
+
+/// Encoding of frames via buffers.
+///
+/// This trait is used when constructing an instance of `UdpFramed` and provides
+/// the `In` and `Out` types which are decoded and encoded from the socket,
+/// respectively.
+///
+/// Because UDP is a connectionless protocol, the `decode` method receives the
+/// address where data came from and the `encode` method is also responsible for
+/// determining the remote host to which the datagram should be sent
+///
+/// The trait itself is implemented on a type that can track state for decoding
+/// or encoding, which is particularly useful for streaming parsers. In many
+/// cases, though, this type will simply be a unit struct (e.g. `struct
+/// HttpCodec`).
+pub trait UdpCodec {
+    /// The type of decoded frames.
+    type In;
+
+    /// The type of frames to be encoded.
+    type Out;
+
+    /// Attempts to decode a frame from the provided buffer of bytes.
+    ///
+    /// This method is called by `UdpFramed` on a single datagram which has been
+    /// read from a socket. The `buf` argument contains the data that was
+    /// received from the remote address, and `src` is the address the data came
+    /// from. Note that typically this method should require the entire contents
+    /// of `buf` to be valid or otherwise return an error with trailing data.
+    ///
+    /// Finally, if the bytes in the buffer are malformed then an error is
+    /// returned indicating why. This informs `Framed` that the stream is now
+    /// corrupt and should be terminated.
+    fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> io::Result<Self::In>;
+
+    /// Encodes a frame into the buffer provided.
+    ///
+    /// This method will encode `msg` into the byte buffer provided by `buf`.
+    /// The `buf` provided is an internal buffer of the `Framed` instance and
+    /// will be written out when possible.
+    ///
+    /// The encode method also determines the destination to which the buffer
+    /// should be directed, which will be returned as a `SocketAddr`.
+    fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr;
+}
+
+/// A unified `Stream` and `Sink` interface to an underlying `UdpSocket`, using
+/// the `UdpCodec` trait to encode and decode frames.
+///
+/// You can acquire a `UdpFramed` instance by using the `UdpSocket::framed`
+/// adapter.
+pub struct UdpFramed<C> {
+    socket: UdpSocket,
+    codec: C,
+    rd: Vec<u8>,
+    wr: Vec<u8>,
+    out_addr: SocketAddr,
+}
+
+impl<C: UdpCodec> Stream for UdpFramed<C> {
+    type Item = C::In;
+    type Error = io::Error;
+
+    fn poll(&mut self) -> Poll<Option<C::In>, io::Error> {
+        let (n, addr) = try_nb!(self.socket.recv_from(&mut self.rd));
+        trace!("received {} bytes, decoding", n);
+        let frame = try!(self.codec.decode(&addr, &self.rd[..n]));
+        trace!("frame decoded from buffer");
+        Ok(Async::Ready(Some(frame)))
+    }
+}
+
+impl<C: UdpCodec> Sink for UdpFramed<C> {
+    type SinkItem = C::Out;
+    type SinkError = io::Error;
+
+    fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
+        if self.wr.len() > 0 {
+            try!(self.poll_complete());
+            if self.wr.len() > 0 {
+                return Ok(AsyncSink::NotReady(item));
+            }
+        }
+
+        self.out_addr = self.codec.encode(item, &mut self.wr);
+        Ok(AsyncSink::Ready)
+    }
+
+    fn poll_complete(&mut self) -> Poll<(), io::Error> {
+        trace!("flushing framed transport");
+
+        if self.wr.is_empty() {
+            return Ok(Async::Ready(()))
+        }
+
+        trace!("writing; remaining={}", self.wr.len());
+        let n = try_nb!(self.socket.send_to(&self.wr, &self.out_addr));
+        trace!("written {}", n);
+        let wrote_all = n == self.wr.len();
+        self.wr.clear();
+        if wrote_all {
+            Ok(Async::Ready(()))
+        } else {
+            Err(io::Error::new(io::ErrorKind::Other,
+                               "failed to write entire datagram to socket"))
+        }
+    }
+}
+
+pub fn new<C: UdpCodec>(socket: UdpSocket, codec: C) -> UdpFramed<C> {
+    UdpFramed {
+        socket: socket,
+        codec: codec,
+        out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
+        rd: vec![0; 64 * 1024],
+        wr: Vec::with_capacity(8 * 1024),
+    }
+}
+
+impl<C> UdpFramed<C> {
+    /// Returns a reference to the underlying I/O stream wrapped by `Framed`.
+    ///
+    /// Note that care should be taken to not tamper with the underlying stream
+    /// of data coming in as it may corrupt the stream of frames otherwise being
+    /// worked with.
+    pub fn get_ref(&self) -> &UdpSocket {
+        &self.socket
+    }
+
+    /// Returns a mutable reference to the underlying I/O stream wrapped by
+    /// `Framed`.
+    ///
+    /// Note that care should be taken to not tamper with the underlying stream
+    /// of data coming in as it may corrupt the stream of frames otherwise being
+    /// worked with.
+    pub fn get_mut(&mut self) -> &mut UdpSocket {
+        &mut self.socket
+    }
+
+    /// Consumes the `Framed`, returning its underlying I/O stream.
+    ///
+    /// Note that care should be taken to not tamper with the underlying stream
+    /// of data coming in as it may corrupt the stream of frames otherwise being
+    /// worked with.
+    pub fn into_inner(self) -> UdpSocket {
+        self.socket
+    }
+}

src/net/udp/mod.rs

@@ -1,10 +1,10 @@
 use std::io;
+use std::mem;
 use std::net::{self, SocketAddr, Ipv4Addr, Ipv6Addr};
 use std::fmt;
-use io::{FramedUdp, framed_udp};
+
 use futures::{Async, Future, Poll};
 use mio;
-use std::mem;
 
 use reactor::{Handle, PollEvented};
 
@@ -13,6 +13,9 @@ pub struct UdpSocket {
     io: PollEvented<mio::udp::UdpSocket>,
 }
 
+mod frame;
+pub use self::frame::{UdpFramed, UdpCodec};
+
 impl UdpSocket {
     /// Create a new UDP socket bound to the specified address.
     ///
@@ -43,11 +46,27 @@ impl UdpSocket {
         UdpSocket::new(udp, handle)
     }
 
-    /// Creates a FramedUdp object, which leverages a supplied `EncodeUdp`
+    /// Provides a `Stream` and `Sink` interface for reading and writing to this
-    /// and `DecodeUdp` to implement `Stream` and `Sink` 
+    /// `UdpSocket` object, using the provided `UdpCodec` to read and write the
-    /// This moves the socket into the newly created FramedUdp object
+    /// raw data.
-    pub fn framed<C>(self, codec : C) -> FramedUdp<C> {
+    ///
-        framed_udp(self, codec)
+    /// Raw UDP sockets work with datagrams, but higher-level code usually
+    /// wants to batch these into meaningful chunks, called "frames". This
+    /// method layers framing on top of this socket by using the `UdpCodec`
+    /// trait to handle encoding and decoding of messages frames. Note that
+    /// the incoming and outgoing frame types may be distinct.
+    ///
+    /// This function returns a *single* object that is both `Stream` and
+    /// `Sink`; grouping this into a single object is often useful for layering
+    /// things which require both read and write access to the underlying
+    /// object.
+    ///
+    /// If you want to work more directly with the streams and sink, consider
+    /// calling `split` on the `UdpFramed` returned by this method, which will
+    /// break them into separate objects, allowing them to interact more
+    /// easily.
+    pub fn framed<C: UdpCodec>(self, codec: C) -> UdpFramed<C> {
+        frame::new(self, codec)
     }
 
     /// Returns the local address that this stream is bound to.
@@ -94,26 +113,26 @@ impl UdpSocket {
         }
     }
 
-    /// Creates a future that will write the entire contents of the buffer `buf` to
+    /// Creates a future that will write the entire contents of the buffer
-    /// the stream `a` provided.
+    /// `buf` provided as a datagram to this socket.
     ///
-    /// The returned future will return after data has been written to the outbound
+    /// The returned future will return after data has been written to the
-    /// socket. 
+    /// outbound socket.  The future will resolve to the stream as well as the
-    /// The future will resolve to the stream as well as the buffer (for reuse if
+    /// buffer (for reuse if needed).
-    /// needed).
     ///
-    /// Any error which happens during writing will cause both the stream and the
+    /// Any error which happens during writing will cause both the stream and
-    /// buffer to get destroyed.
+    /// the buffer to get destroyed. Note that failure to write the entire
+    /// buffer is considered an error for the purposes of sending a datagram.
     ///
-    /// The `buf` parameter here only requires the `AsRef<[u8]>` trait, which should
+    /// The `buf` parameter here only requires the `AsRef<[u8]>` trait, which
-    /// be broadly applicable to accepting data which can be converted to a slice.
+    /// should be broadly applicable to accepting data which can be converted
-    /// The `Window` struct is also available in this crate to provide a different
+    /// to a slice.  The `Window` struct is also available in this crate to
-    /// window into a slice if necessary.
+    /// provide a different window into a slice if necessary.
-    pub fn send_dgram<T>(self, buf: T, addr : SocketAddr) -> SendDGram<T>
+    pub fn send_dgram<T>(self, buf: T, addr: SocketAddr) -> SendDgram<T>
         where T: AsRef<[u8]>,
     {
-        SendDGram {
+        SendDgram {
-            state: UdpState::Writing {
+            state: SendState::Writing {
                 sock: self,
                 addr: addr,
                 buf: buf,
@@ -137,6 +156,31 @@ impl UdpSocket {
         }
     }
 
+    /// Creates a future that receive a datagram to be written to the buffer
+    /// provided.
+    ///
+    /// The returned future will return after a datagram has been received on
+    /// this socket. The future will resolve to the socket, the buffer, the
+    /// amount of data read, and the address the data was received from.
+    ///
+    /// An error during reading will cause the socket and buffer to get
+    /// destroyed and the socket will be returned.
+    ///
+    /// The `buf` parameter here only requires the `AsMut<[u8]>` trait, which
+    /// should be broadly applicable to accepting data which can be converted
+    /// to a slice.  The `Window` struct is also available in this crate to
+    /// provide a different window into a slice if necessary.
+    pub fn recv_dgram<T>(self, buf: T) -> RecvDgram<T>
+        where T: AsMut<[u8]>,
+    {
+        RecvDgram {
+            state: RecvState::Reading {
+                sock: self,
+                buf: buf,
+            },
+        }
+    }
+
     /// Gets the value of the `SO_BROADCAST` option for this socket.
     ///
     /// For more information about this option, see
@@ -284,16 +328,14 @@ impl fmt::Debug for UdpSocket {
     }
 }
 
-/// A future used to write the entire contents of some data to a stream.
+/// A future used to write the entire contents of some data to a UDP socket.
 ///
-/// This is created by the [`write_all`] top-level method.
+/// This is created by the `UdpSocket::send_dgram` method.
-///
+pub struct SendDgram<T> {
-/// [`write_all`]: fn.write_all.html
+    state: SendState<T>,
-pub struct SendDGram<T> {
-    state: UdpState<T>,
 }
 
-enum UdpState<T> {
+enum SendState<T> {
     Writing {
         sock: UdpSocket,
         buf: T,
@@ -306,7 +348,7 @@ fn incomplete_write(reason : &str) -> io::Error {
     io::Error::new(io::ErrorKind::Other, reason)
 }
 
-impl<T> Future for SendDGram<T>
+impl<T> Future for SendDgram<T>
     where T: AsRef<[u8]>,
 {
     type Item = (UdpSocket, T);
@@ -314,19 +356,59 @@ impl<T> Future for SendDGram<T>
 
     fn poll(&mut self) -> Poll<(UdpSocket, T), io::Error> {
         match self.state {
-            UdpState::Writing { ref sock, ref buf, ref addr} => {
+            SendState::Writing { ref sock, ref buf, ref addr } => {
-                let buf = buf.as_ref();
+                let n = try_nb!(sock.send_to(buf.as_ref(), addr));
-                let n = try_nb!(sock.send_to(&buf, addr));
+                if n != buf.as_ref().len() {
-                if n != buf.len() {
+                    return Err(incomplete_write("failed to send entire message \
-                    return Err(incomplete_write("Failed to send entire message in datagram"))
+                                                 in datagram"))
                 }
             }
-            UdpState::Empty => panic!("poll a SendDGram after it's done"),
+            SendState::Empty => panic!("poll a SendDgram after it's done"),
         }
 
-        match mem::replace(&mut self.state, UdpState::Empty) {
+        match mem::replace(&mut self.state, SendState::Empty) {
-            UdpState::Writing { sock, buf, .. } => Ok(Async::Ready((sock, (buf).into()))),
+            SendState::Writing { sock, buf, addr: _ } => {
-            UdpState::Empty => panic!(),
+                Ok(Async::Ready((sock, buf)))
+            }
+            SendState::Empty => panic!(),
+        }
+    }
+}
+
+/// A future used to receive a datagram from a UDP socket.
+///
+/// This is created by the `UdpSocket::recv_dgram` method.
+pub struct RecvDgram<T> {
+    state: RecvState<T>,
+}
+
+enum RecvState<T> {
+    Reading {
+        sock: UdpSocket,
+        buf: T,
+    },
+    Empty,
+}
+
+impl<T> Future for RecvDgram<T>
+    where T: AsMut<[u8]>,
+{
+    type Item = (UdpSocket, T, usize, SocketAddr);
+    type Error = io::Error;
+
+    fn poll(&mut self) -> Poll<Self::Item, io::Error> {
+        let (n, addr) = match self.state {
+            RecvState::Reading { ref sock, ref mut buf } => {
+                try_nb!(sock.recv_from(buf.as_mut()))
+            }
+            RecvState::Empty => panic!("poll a RecvDgram after it's done"),
+        };
+
+        match mem::replace(&mut self.state, RecvState::Empty) {
+            RecvState::Reading { sock, buf } => {
+                Ok(Async::Ready((sock, buf, n, addr)))
+            }
+            RecvState::Empty => panic!(),
         }
     }
 }