From 592a99bca4e760d00057fc3927c3c6f164e353e2 Mon Sep 17 00:00:00 2001 From: Rick Richardson Date: Sat, 19 Nov 2016 09:05:00 -0800 Subject: [PATCH] completed basic implementation of FramedUdp for streams and sink --- src/io/udp_frame.rs | 258 ++++++++++++++++++++++++++++++++++++++++++++ src/net/udp.rs | 14 +++ 2 files changed, 272 insertions(+) create mode 100644 src/io/udp_frame.rs diff --git a/src/io/udp_frame.rs b/src/io/udp_frame.rs new file mode 100644 index 000000000..9c1eeb958 --- /dev/null +++ b/src/io/udp_frame.rs @@ -0,0 +1,258 @@ +use std::io; +use std::ops::{Deref, DerefMut}; +use std::sync::Arc; +use net::udp::UdpSocket +use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink}; +use futures::sync::BiLock; + +use io::Io; + +/// 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 EncodeUdp { + + /// The type of frames to be encoded. + type Out; + + + /// 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 codec 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) -> SocketAddr; +} + +/// Decoding of frames via buffers. +/// +/// This trait is used when constructing an instance of `FramedUdp`. It provides +/// one type: `In` for decoding incoming frames from a Datagram +/// +/// Because UDP is a connectionless protocol, the decode method will also be +/// supplied with a SocketAddr of the remote host which sent the datagram +/// +/// 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 DecodeUdp { + /// The type of decoded frames. + type In; + + /// 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 Decoder 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. + /// + /// If the bytes look valid, but a frame isn't fully available yet, then + /// `Ok(None)` is returned. This indicates to the `Framed` instance that + /// it needs to read some more bytes before calling this method again. + /// In such a case, it is the decoder's responsibility to copy the data + /// into their own internal buffer for future use. + /// + /// 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. + /// + /// When dealing with connectionless streams, there will likely be some sort + /// of state machine. + fn decode(&mut self, src: &SocketAddr, buf: &mut Vec) -> Result, io::Error>; +} + +/// A unified `Stream` and `Sink` interface to an underlying `Io` object, using +/// the `Encode` and `Decode` traits to encode and decode frames. +/// +/// You can acquire a `Framed` instance by using the `Io::framed` adapter. +pub struct FramedUdp { + socket: UdpSocket, + encoder: E, + decoder: D, + out_addr : Option, + rd: Vec, + wr: Vec, +} + +impl Stream for Framed { + type Item = D::In; + type Error = io::Error; + + fn poll(&mut self) -> Poll, io::Error> { + loop { + + let before = self.rd.len(); + let ret = self.socket.recv_from(self.rd.mut_bytes(), &mut inaddr); + match ret { + Ok((n, addr)) => { + trace!("read {} bytes", n); + trace!("attempting to decode a frame"); + if let Some(frame) = try!(self.decoder.decode(&addr, &mut self.rd)) { + trace!("frame decoded from buffer"); + self.rd.clear(); + return Ok(Async::Ready(Some(frame))); + } + } + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + if self.rd.len() == before { + return Ok(Async::NotReady) + } + } + Err(e) => return Err(e), + } + } + } +} + +impl Sink for Framed { + type SinkItem = E::Out; + type SinkError = io::Error; + + fn start_send(&mut self, item: C::Out) -> StartSend { + 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"); + + while !self.wr.is_empty() { + if let Some(outaddr) = self.out_addr.ref() { + trace!("writing; remaining={}", self.wr.len()); + let n = try_nb!(self.socket.send_to(&self.wr, outaddr)); + self.wr.clear(); + self.out_addr = None; + if n != self.wr.len() { + return Err(io::Error::new(io::ErrorKind::WriteZero, + "failed to write frame datagram to socket")); + } + } + else { + return Err(io::Error::new(io::ErrorKind::Other, + "outbound stream in invalid state: out_addr is not known")); + } + } + + return Ok(Async::Ready(())); + } +} + +pub fn framed_udp(socket : UdpSocket, decoder : D, encoder : E) -> Framed { + Framed { + socket: socket, + encoder: encoder, + decoder: decoder, + rd: Vec::with_capacity(64 * 1024), + wr: Vec::with_capacity(64 * 1024) + } +} + +impl FramedUdp { + /// 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) -> (FramedRead, FramedWrite) { + 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 `Decode` trait +/// to decode frames. +pub struct FramedRead { + framed: BiLock>, +} + +impl Stream for FramedRead { + type Item = D::In; + type Error = io::Error; + + fn poll(&mut self) -> Poll, 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 `Encode` trait +/// to encode frames. +pub struct FramedWrite { + framed: BiLock>, +} + +impl Sink for FramedWrite { + type SinkItem = E::Out; + type SinkError = io::Error; + + fn start_send(&mut self, item: E::Out) -> StartSend { + 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) + } + } +} + diff --git a/src/net/udp.rs b/src/net/udp.rs index a88a8ca16..2a2e0a03a 100644 --- a/src/net/udp.rs +++ b/src/net/udp.rs @@ -42,6 +42,20 @@ impl UdpSocket { UdpSocket::new(udp, handle) } + /// Creates a FramedUdp object, which leverages a supplied `EncodeUdp` + /// and `DecodeUdp` to implement `Stream` and `Sink` + /// This moves the socket into the newly created FramedUdp object + pub fn framed(self, decoder : D, encoder : E) -> Framed { + FramedUdp { + socket: self, + encoder: encoder, + decoder: decoder, + is_readable: false, + rd: Vec::with_capacity(64 * 1024); + wr: Vec::with_capacity(64 * 1024), + } + } + /// Returns the local address that this stream is bound to. pub fn local_addr(&self) -> io::Result { self.io.get_ref().local_addr()