From 3658e10045664e816a65851d1e25c68965ec4359 Mon Sep 17 00:00:00 2001 From: Patrick Barrett Date: Tue, 20 Nov 2018 11:19:34 -0600 Subject: [PATCH] uds: implement UnixDatagramFramed (#453) Implement `Stream + Sink` layer on top of unix domain sockets using codecs. --- tokio-uds/Cargo.toml | 1 + tokio-uds/src/frame.rs | 160 ++++++++++++++++++++++++++++++++++++ tokio-uds/src/lib.rs | 4 + tokio-uds/tests/datagram.rs | 83 +++++++++++++++++++ 4 files changed, 248 insertions(+) create mode 100644 tokio-uds/src/frame.rs create mode 100644 tokio-uds/tests/datagram.rs diff --git a/tokio-uds/Cargo.toml b/tokio-uds/Cargo.toml index 72e222f8c..2935e7e45 100644 --- a/tokio-uds/Cargo.toml +++ b/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" } diff --git a/tokio-uds/src/frame.rs b/tokio-uds/src/frame.rs new file mode 100644 index 000000000..49c8948c6 --- /dev/null +++ b/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 { + socket: UnixDatagram, + codec: C, + rd: BytesMut, + wr: BytesMut, + out_addr: Option, + flushed: bool, +} + +impl Stream for UnixDatagramFramed { + type Item = (C::Item, SocketAddr); + type Error = C::Error; + + fn poll(&mut self) -> Poll, 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, C: Encoder> Sink for UnixDatagramFramed { + type SinkItem = (C::Item, A); + type SinkError = C::Error; + + fn start_send(&mut self, item: Self::SinkItem) -> StartSend { + 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 UnixDatagramFramed { + /// 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 { + 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 + } +} diff --git a/tokio-uds/src/lib.rs b/tokio-uds/src/lib.rs index c4afed2c5..c751e3966 100644 --- a/tokio-uds/src/lib.rs +++ b/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; diff --git a/tokio-uds/tests/datagram.rs b/tokio-uds/tests/datagram.rs new file mode 100644 index 000000000..6f58f7c07 --- /dev/null +++ b/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, 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"); + } +}