use std::io; use std::ops::{Deref, DerefMut}; use std::sync::Arc; use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink}; use futures::sync::BiLock; use io::Io; /// A reference counted buffer of bytes. /// /// An `EasyBuf` is a representation of a byte buffer where sub-slices of it can /// be handed out efficiently, each with a `'static` lifetime which keeps the /// data alive. The buffer also supports mutation but may require bytes to be /// copied to complete the operation. pub struct EasyBuf { buf: Arc>, start: usize, end: usize, } /// An RAII object returned from `get_mut` which provides mutable access to the /// underlying `Vec`. pub struct EasyBufMut<'a> { buf: &'a mut Vec, end: &'a mut usize, } impl EasyBuf { /// Creates a new EasyBuf with no data and the default capacity. pub fn new() -> EasyBuf { EasyBuf::with_capacity(8 * 1024) } /// Creates a new EasyBuf with `cap` capacity. pub fn with_capacity(cap: usize) -> EasyBuf { EasyBuf { buf: Arc::new(Vec::with_capacity(cap)), start: 0, end: 0, } } /// Changes the starting index of this window to the index specified. /// /// Returns the windows back to chain multiple calls to this method. /// /// # Panics /// /// This method will panic if `start` is out of bounds for the underlying /// slice or if it comes after the `end` configured in this window. fn set_start(&mut self, start: usize) -> &mut EasyBuf { assert!(start <= self.buf.as_ref().len()); assert!(start <= self.end); self.start = start; self } /// Changes the end index of this window to the index specified. /// /// Returns the windows back to chain multiple calls to this method. /// /// # Panics /// /// This method will panic if `end` is out of bounds for the underlying /// slice or if it comes after the `end` configured in this window. fn set_end(&mut self, end: usize) -> &mut EasyBuf { assert!(end <= self.buf.len()); assert!(self.start <= end); self.end = end; self } /// Returns the number of bytes contained in this `EasyBuf`. pub fn len(&self) -> usize { self.end - self.start } /// Returns the inner contents of this `EasyBuf` as a slice. pub fn as_slice(&self) -> &[u8] { self.as_ref() } /// Splits the buffer into two at the given index. /// /// Afterwards `self` contains elements `[0, at)`, and the returned `EasyBuf` /// contains elements `[at, len)`. /// /// This is an O(1) operation that just increases the reference count and /// sets a few indexes. /// /// # Panics /// /// Panics if `at > len` pub fn split_off(&mut self, at: usize) -> EasyBuf { let mut other = EasyBuf { buf: self.buf.clone(), ..*self }; let idx = self.start + at; other.set_start(idx); self.set_end(idx); return other } /// Splits the buffer into two at the given index. /// /// Afterwards `self` contains elements `[at, len)`, and the returned `EasyBuf` /// contains elements `[0, at)`. /// /// This is an O(1) operation that just increases the reference count and /// sets a few indexes. /// /// # Panics /// /// Panics if `at > len` pub fn drain_to(&mut self, at: usize) -> EasyBuf { let mut other = EasyBuf { buf: self.buf.clone(), ..*self }; let idx = self.start + at; other.set_end(idx); self.set_start(idx); return other } /// Returns a mutable reference to the underlying growable buffer of bytes. /// /// If this `EasyBuf` is the only instance pointing at the underlying buffer /// of bytes, a direct mutable reference will be returned. Otherwise the /// contents of this `EasyBuf` will be reallocated in a fresh `Vec` /// allocation with the same capacity as this allocation, and that /// allocation will be returned. /// /// This operation **is not O(1)** as it may clone the entire contents of /// this buffer. /// /// The returned `EasyBufMut` type implement `Deref` and `DerefMut` to /// `Vec` can the byte buffer can be manipulated using the standard /// `Vec` methods. pub fn get_mut(&mut self) -> EasyBufMut { // Fast path if we can get mutable access to our own current // buffer. // // TODO: this should be a match or an if-let if Arc::get_mut(&mut self.buf).is_some() { let buf = Arc::get_mut(&mut self.buf).unwrap(); buf.drain(..self.start); self.start = 0; return EasyBufMut { buf: buf, end: &mut self.end } } // If we couldn't get access above then we give ourself a new buffer // here. let mut v = Vec::with_capacity(self.buf.capacity()); v.extend_from_slice(self.as_ref()); self.start = 0; self.buf = Arc::new(v); EasyBufMut { buf: Arc::get_mut(&mut self.buf).unwrap(), end: &mut self.end, } } } impl AsRef<[u8]> for EasyBuf { fn as_ref(&self) -> &[u8] { &self.buf[self.start..self.end] } } impl<'a> Deref for EasyBufMut<'a> { type Target = Vec; fn deref(&self) -> &Vec { self.buf } } impl<'a> DerefMut for EasyBufMut<'a> { fn deref_mut(&mut self) -> &mut Vec { self.buf } } impl<'a> Drop for EasyBufMut<'a> { fn drop(&mut self) { *self.end = self.buf.len(); } } /// Encoding and decoding of frames via buffers. /// /// This trait is used when constructing an instance of `Framed`. It provides /// two types: `In`, for decoded input frames, and `Out`, for outgoing frames /// that need to be encoded. It also provides methods to actually perform the /// encoding and decoding, which work with corresponding buffer types. /// /// 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 Codec { /// 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 `Framed` whenever bytes are ready to be parsed. /// The provided buffer of bytes is what's been read so far, and this /// instance of `Decode` can determine whether an entire frame is in the /// buffer and is ready to be returned. /// /// If an entire frame is available, then this instance will remove those /// bytes from the buffer provided and return them as a decoded /// frame. Note that removing bytes from the provided buffer doesn't always /// necessarily copy the bytes, so this should be an efficient operation in /// most circumstances. /// /// 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. /// /// 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, buf: &mut EasyBuf) -> Result, io::Error>; /// A default method available to be called when there are no more bytes /// available to be read from the underlying I/O. /// /// This method defaults to calling `decode` and returns an error if /// `Ok(None)` is returned. Typically this doesn't need to be implemented /// unless the framing protocol differs near the end of the stream. fn decode_eof(&mut self, buf: &mut EasyBuf) -> io::Result { match try!(self.decode(buf)) { Some(frame) => Ok(frame), None => Err(io::Error::new(io::ErrorKind::Other, "bytes remaining on stream")), } } /// 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. fn encode(&mut self, msg: Self::Out, buf: &mut Vec); } /// 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 = C::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 = C::Out; type SinkError = io::Error; fn start_send(&mut self, item: C::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) } } } /// 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 Framed { upstream: T, codec: C, eof: bool, is_readable: bool, rd: EasyBuf, wr: Vec, } impl Stream for Framed { type Item = C::In; type Error = io::Error; fn poll(&mut self) -> Poll, io::Error> { loop { // If the read buffer has any pending data, then it could be // possible that `decode` will return a new frame. We leave it to // the decoder to optimize detecting that more data is required. if self.is_readable { if self.eof { if self.rd.len() == 0 { return Ok(None.into()) } else { let frame = try!(self.codec.decode_eof(&mut self.rd)); return Ok(Async::Ready(Some(frame))) } } trace!("attempting to decode a frame"); if let Some(frame) = try!(self.codec.decode(&mut self.rd)) { trace!("frame decoded from buffer"); return Ok(Async::Ready(Some(frame))); } self.is_readable = false; } assert!(!self.eof); // Otherwise, try to read more data and try again // // TODO: shouldn't read_to_end, that may read a lot let before = self.rd.len(); let ret = self.upstream.read_to_end(&mut self.rd.get_mut()); match ret { Ok(_n) => self.eof = true, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { if self.rd.len() == before { return Ok(Async::NotReady) } } Err(e) => return Err(e), } self.is_readable = true; } } } impl Sink for Framed { type SinkItem = C::Out; type SinkError = io::Error; fn start_send(&mut self, item: C::Out) -> StartSend { self.codec.encode(item, &mut self.wr); Ok(AsyncSink::Ready) } fn poll_complete(&mut self) -> Poll<(), io::Error> { // Try flushing the underlying IO try_nb!(self.upstream.flush()); trace!("flushing framed transport"); loop { if self.wr.len() == 0 { trace!("framed transport flushed"); return Ok(Async::Ready(())); } trace!("writing; remaining={:?}", self.wr.len()); let n = try_nb!(self.upstream.write(&self.wr)); self.wr.drain(..n); } } } pub fn framed(io: T, codec: C) -> Framed { Framed { upstream: io, codec: codec, eof: false, is_readable: false, rd: EasyBuf::new(), wr: Vec::with_capacity(8 * 1024), } } impl Framed { /// 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 = FramedRead { framed: a }; let write = FramedWrite { 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) -> &T { &self.upstream } /// 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 T { &mut self.upstream } /// 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) -> T { self.upstream } }