uds: implement UnixDatagramFramed (#453)

Implement `Stream + Sink` layer on top of unix domain sockets
using codecs.

tokio-uds/Cargo.toml

@@ -24,6 +24,7 @@ libc = "0.2.42"
 log = "0.4.2"
 mio = "0.6.14"
 mio-uds = "0.6.5"
+tokio-codec = { version = "0.1.0", path = "../tokio-codec" }
 tokio-reactor = { version = "0.1.1", path = "../tokio-reactor" }
 tokio-io = { version = "0.1.6", path = "../tokio-io" }
 

tokio-uds/src/frame.rs

@@ -0,0 +1,160 @@
+use std::io;
+use std::os::unix::net::SocketAddr;
+use std::path::Path;
+
+use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};
+
+use super::UnixDatagram;
+
+use tokio_codec::{Decoder, Encoder};
+use bytes::{BytesMut, BufMut};
+
+/// A unified `Stream` and `Sink` interface to an underlying `UnixDatagram`, using
+/// the `Encoder` and `Decoder` traits to encode and decode frames.
+///
+/// Unix datagram sockets work with datagrams, but higher-level code may wants to
+/// batch these into meaningful chunks, called "frames". This method layers
+/// framing on top of this socket by using the `Encoder` and `Decoder` traits 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 `UnixDatagramFramed` returned by this method, which will break
+/// them into separate objects, allowing them to interact more easily.
+#[must_use = "sinks do nothing unless polled"]
+#[derive(Debug)]
+pub struct UnixDatagramFramed<A, C> {
+    socket: UnixDatagram,
+    codec: C,
+    rd: BytesMut,
+    wr: BytesMut,
+    out_addr: Option<A>,
+    flushed: bool,
+}
+
+impl<A, C: Decoder> Stream for UnixDatagramFramed<A, C> {
+    type Item = (C::Item, SocketAddr);
+    type Error = C::Error;
+
+    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        self.rd.reserve(INITIAL_RD_CAPACITY);
+
+        let (n, addr) = unsafe {
+            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));
+        trace!("frame decoded from buffer");
+        Ok(Async::Ready(result))
+    }
+}
+
+impl<A: AsRef<Path>, C: Encoder> Sink for UnixDatagramFramed<A, C> {
+    type SinkItem = (C::Item, A);
+    type SinkError = C::Error;
+
+    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
+        trace!("sending frame");
+
+        if !self.flushed {
+            match try!(self.poll_complete()) {
+                Async::Ready(()) => {},
+                Async::NotReady => return Ok(AsyncSink::NotReady(item)),
+            }
+        }
+
+        let (frame, out_addr) = item;
+        self.codec.encode(frame, &mut self.wr)?;
+        self.out_addr = Some(out_addr);
+        self.flushed = false;
+        trace!("frame encoded; length={}", self.wr.len());
+
+        Ok(AsyncSink::Ready)
+    }
+
+    fn poll_complete(&mut self) -> Poll<(), C::Error> {
+        if self.flushed {
+            return Ok(Async::Ready(()))
+        }
+
+        let n = {
+            let out_path = match self.out_addr {
+                Some(ref out_path) => out_path.as_ref(),
+                None => return Err(io::Error::new(io::ErrorKind::Other,
+                                                      "internal error: addr not available while data not flushed").into()),
+            };
+
+            trace!("flushing frame; length={}", self.wr.len());
+            try_ready!(self.socket.poll_send_to(&self.wr, out_path))
+        };
+
+        trace!("written {}", n);
+
+        let wrote_all = n == self.wr.len();
+        self.wr.clear();
+        self.flushed = true;
+
+        if wrote_all {
+            self.out_addr = None;
+            Ok(Async::Ready(()))
+        } else {
+            Err(io::Error::new(io::ErrorKind::Other,
+                               "failed to write entire datagram to socket").into())
+        }
+    }
+
+    fn close(&mut self) -> Poll<(), C::Error> {
+        self.poll_complete()
+    }
+}
+
+const INITIAL_RD_CAPACITY: usize = 64 * 1024;
+const INITIAL_WR_CAPACITY: usize = 8 * 1024;
+
+impl<A, C> UnixDatagramFramed<A, C> {
+    /// Create a new `UnixDatagramFramed` backed by the given socket and codec.
+    ///
+    /// See struct level documentation for more details.
+    pub fn new(socket: UnixDatagram, codec: C) -> UnixDatagramFramed<A, C> {
+        UnixDatagramFramed {
+            socket: socket,
+            codec: codec,
+            out_addr: None,
+            rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY),
+            wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY),
+            flushed: true,
+        }
+    }
+
+    /// Returns a reference to the underlying I/O stream wrapped by `Framed`.
+    ///
+    /// # Note
+    ///
+    /// 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) -> &UnixDatagram {
+        &self.socket
+    }
+
+    /// Returns a mutable reference to the underlying I/O stream wrapped by
+    /// `Framed`.
+    ///
+    /// # Note
+    ///
+    /// 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 UnixDatagram {
+        &mut self.socket
+    }
+}

tokio-uds/src/lib.rs

@@ -11,13 +11,16 @@ extern crate bytes;
 extern crate futures;
 extern crate iovec;
 extern crate libc;
+#[macro_use]
 extern crate log;
 extern crate mio;
 extern crate mio_uds;
+extern crate tokio_codec;
 extern crate tokio_io;
 extern crate tokio_reactor;
 
 mod datagram;
+mod frame;
 mod incoming;
 mod listener;
 mod recv_dgram;
@@ -26,6 +29,7 @@ mod stream;
 mod ucred;
 
 pub use datagram::UnixDatagram;
+pub use frame::UnixDatagramFramed;
 pub use incoming::Incoming;
 pub use listener::UnixListener;
 pub use recv_dgram::RecvDgram;

tokio-uds/tests/datagram.rs

@@ -0,0 +1,83 @@
+#![cfg(unix)]
+
+extern crate bytes;
+extern crate futures;
+extern crate tempfile;
+extern crate tokio;
+extern crate tokio_codec;
+extern crate tokio_uds;
+
+use tokio_uds::*;
+
+use std::str;
+
+use bytes::BytesMut;
+
+use tokio::io;
+use tokio::runtime::current_thread::Runtime;
+
+use tokio_codec::{Decoder, Encoder};
+
+use futures::{Future, Sink, Stream};
+
+struct StringDatagramCodec;
+
+/// A codec to decode datagrams from a unix domain socket as utf-8 text messages.
+impl Encoder for StringDatagramCodec {
+    type Item = String;
+    type Error = io::Error;
+
+    fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
+        dst.extend_from_slice(&item.into_bytes());
+        Ok(())
+    }
+}
+
+/// A codec to decode datagrams from a unix domain socket as utf-8 text messages.
+impl Decoder for StringDatagramCodec {
+    type Item = String;
+    type Error = io::Error;
+
+    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        let decoded = str::from_utf8(buf)
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
+            .to_string();
+
+        Ok(Some(decoded))
+    }
+}
+
+#[test]
+fn framed_echo() {
+    let dir = tempfile::tempdir().unwrap();
+    let server_path = dir.path().join("server.sock");
+    let client_path = dir.path().join("client.sock");
+
+    let mut rt = Runtime::new().unwrap();
+
+    {
+        let socket = UnixDatagram::bind(&server_path).unwrap();
+        let server = UnixDatagramFramed::new(socket, StringDatagramCodec);
+
+        let (sink, stream) = server.split();
+
+        let echo_stream = stream.map(|(msg, addr)| (msg, addr.as_pathname().unwrap().to_path_buf()));
+
+        // spawn echo server
+        rt.spawn(echo_stream.forward(sink)
+                 .map_err(|e| panic!("err={:?}", e))
+                 .map(|_| ()));
+    }
+
+    {
+        let socket = UnixDatagram::bind(&client_path).unwrap();
+        let client = UnixDatagramFramed::new(socket, StringDatagramCodec);
+
+        let (sink, stream) = client.split();
+
+        rt.block_on(sink.send(("ECHO".to_string(), server_path))).unwrap();
+
+        let response = rt.block_on(stream.take(1).collect()).unwrap();
+        assert_eq!(response[0].0, "ECHO");
+    }
+}