tokio_core/io/

frame.rs

use std::fmt;
use std::io;
use std::hash;
use std::mem;
use std::cmp;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;

use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};

use io::Io;

const INITIAL_CAPACITY: usize = 8 * 1024;

/// 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.
#[derive(Clone, Eq)]
pub struct EasyBuf {
    buf: Arc<Vec<u8>>,
    start: usize,
    end: usize,
}

/// An RAII object returned from `get_mut` which provides mutable access to the
/// underlying `Vec<u8>`.
pub struct EasyBufMut<'a> {
    buf: &'a mut Vec<u8>,
    end: &'a mut usize,
}

impl EasyBuf {
    /// Creates a new EasyBuf with no data and the default capacity.
    pub fn new() -> EasyBuf {
        EasyBuf::with_capacity(INITIAL_CAPACITY)
    }

    /// 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<u8>`
    /// allocation with the same capacity as an `EasyBuf` created with `EasyBuf::new()`,
    /// 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<u8>` can the byte buffer can be manipulated using the standard
    /// `Vec<u8>` 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.end..);
            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(cmp::max(INITIAL_CAPACITY, self.as_ref().len()));
        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<u8>;

    fn deref(&self) -> &Vec<u8> {
        self.buf
    }
}

impl<'a> DerefMut for EasyBufMut<'a> {
    fn deref_mut(&mut self) -> &mut Vec<u8> {
        self.buf
    }
}

impl From<Vec<u8>> for EasyBuf {
    fn from(vec: Vec<u8>) -> EasyBuf {
        let end = vec.len();
        EasyBuf {
            buf: Arc::new(vec),
            start: 0,
            end: end,
        }
    }
}

impl<T: AsRef<[u8]>> PartialEq<T> for EasyBuf {
    fn eq(&self, other: &T) -> bool {
        self.as_slice().eq(other.as_ref())
    }
}

impl Ord for EasyBuf {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.as_slice().cmp(other.as_slice())
    }
}

impl<T: AsRef<[u8]>> PartialOrd<T> for EasyBuf {
    fn partial_cmp(&self, other: &T) -> Option<cmp::Ordering> {
        self.as_slice().partial_cmp(other.as_ref())
    }
}

impl hash::Hash for EasyBuf {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.as_slice().hash(state)
    }
}

impl<'a> Drop for EasyBufMut<'a> {
    fn drop(&mut self) {
        *self.end = self.buf.len();
    }
}

impl fmt::Debug for EasyBuf {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        let bytes = self.as_ref();
        let len = self.len();
        if len < 10 {
            write!(formatter, "EasyBuf{{len={}/{} {:?}}}", self.len(), self.buf.len(), bytes)
        } else { // choose a more compact representation
            write!(formatter, "EasyBuf{{len={}/{} [{}, {}, {}, {}, ..., {}, {}, {}, {}]}}", self.len(), self.buf.len(), bytes[0], bytes[1], bytes[2], bytes[3], bytes[len-4], bytes[len-3], bytes[len-2], bytes[len-1])
        }
    }
}

impl Into<Vec<u8>> for EasyBuf {
    fn into(mut self) -> Vec<u8> {
        mem::replace(self.get_mut().buf, vec![])
    }
}

/// 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) -> io::Result<Option<Self::In>>;

    /// 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<Self::In> {
        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<u8>) -> io::Result<()>;
}

/// A unified `Stream` and `Sink` interface to an underlying `Io` object, using
/// the `Codec` trait to encode and decode frames.
///
/// You can acquire a `Framed` instance by using the `Io::framed` adapter.
#[must_use = "streams do nothing unless polled"]
pub struct Framed<T, C> {
    upstream: T,
    codec: C,
    eof: bool,
    is_readable: bool,
    rd: EasyBuf,
    wr: Vec<u8>,
}

impl<T: Io, C: Codec> Stream for Framed<T, C> {
    type Item = C::In;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<C::In>, 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<T: Io, C: Codec> Sink for Framed<T, C> {
    type SinkItem = C::Out;
    type SinkError = io::Error;

    fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
        // If the buffer is already over 8KiB, then attempt to flush it. If after flushing it's
        // *still* over 8KiB, then apply backpressure (reject the send).
        const BACKPRESSURE_BOUNDARY: usize = INITIAL_CAPACITY;
        if self.wr.len() > BACKPRESSURE_BOUNDARY {
            try!(self.poll_complete());
            if self.wr.len() > BACKPRESSURE_BOUNDARY {
                return Ok(AsyncSink::NotReady(item));
            }
        }

        try!(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() {
            trace!("writing; remaining={}", self.wr.len());
            let n = try_nb!(self.upstream.write(&self.wr));
            if n == 0 {
                return Err(io::Error::new(io::ErrorKind::WriteZero,
                                          "failed to write frame to transport"));
            }
            self.wr.drain(..n);
        }

        // Try flushing the underlying IO
        try_nb!(self.upstream.flush());

        trace!("framed transport flushed");
        return Ok(Async::Ready(()));
    }

    fn close(&mut self) -> Poll<(), io::Error> {
        try_ready!(self.poll_complete());
        Ok(().into())
    }
}

pub fn framed<T, C>(io: T, codec: C) -> Framed<T, C> {
    Framed {
        upstream: io,
        codec: codec,
        eof: false,
        is_readable: false,
        rd: EasyBuf::new(),
        wr: Vec::with_capacity(INITIAL_CAPACITY),
    }
}

impl<T, C> Framed<T, 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) -> &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
    }
}

#[cfg(test)]
mod tests {
    use super::{INITIAL_CAPACITY, EasyBuf};
    use std::mem;

    #[test]
    fn debug_empty_easybuf() {
        let buf: EasyBuf = vec![].into();
        assert_eq!("EasyBuf{len=0/0 []}", format!("{:?}", buf));
    }

    #[test]
    fn debug_small_easybuf() {
        let buf: EasyBuf = vec![1, 2, 3, 4, 5, 6].into();
        assert_eq!("EasyBuf{len=6/6 [1, 2, 3, 4, 5, 6]}", format!("{:?}", buf));
    }

    #[test]
    fn debug_small_easybuf_split() {
        let mut buf: EasyBuf = vec![1, 2, 3, 4, 5, 6].into();
        let split = buf.split_off(4);
        assert_eq!("EasyBuf{len=4/6 [1, 2, 3, 4]}", format!("{:?}", buf));
        assert_eq!("EasyBuf{len=2/6 [5, 6]}", format!("{:?}", split));
    }

    #[test]
    fn debug_large_easybuf() {
        let vec: Vec<u8> = (0u8..255u8).collect();
        let buf: EasyBuf = vec.into();
        assert_eq!("EasyBuf{len=255/255 [0, 1, 2, 3, ..., 251, 252, 253, 254]}", format!("{:?}", buf));
    }

    #[test]
    fn easybuf_get_mut_sliced() {
        let vec: Vec<u8> = (0u8..10u8).collect();
        let mut buf: EasyBuf = vec.into();
        buf.split_off(9);
        buf.drain_to(3);
        assert_eq!(*buf.get_mut(), [3, 4, 5, 6, 7, 8]);
    }

    #[test]
    fn easybuf_get_mut_sliced_allocating_at_least_initial_capacity() {
        let vec: Vec<u8> = (0u8..10u8).collect();
        let mut buf: EasyBuf = vec.into();
        buf.split_off(9);
        buf.drain_to(3);
        // Clone to make shared
        let clone = buf.clone();
        assert_eq!(*buf.get_mut(), [3, 4, 5, 6, 7, 8]);
        assert_eq!(buf.get_mut().buf.capacity(), INITIAL_CAPACITY);
        mem::drop(clone); // prevent unused warning
    }

    #[test]
    fn easybuf_get_mut_sliced_allocating_required_capacity() {
        let vec: Vec<u8> = (0..INITIAL_CAPACITY * 2).map(|_|0u8).collect();
        let mut buf: EasyBuf = vec.into();
        buf.drain_to(INITIAL_CAPACITY / 2);
        let clone = buf.clone();
        assert_eq!(buf.get_mut().buf.capacity(), INITIAL_CAPACITY + INITIAL_CAPACITY / 2);
        mem::drop(clone)
    }

    #[test]
    fn easybuf_into_vec_simple() {
        let vec: Vec<u8> = (0u8..10u8).collect();
        let reference = vec.clone();
        let buf: EasyBuf = vec.into();
        let original_pointer = buf.buf.as_ref().as_ptr();
        let result: Vec<u8> = buf.into();
        assert_eq!(result, reference);
        let new_pointer = result.as_ptr();
        assert_eq!(original_pointer, new_pointer, "Into<Vec<u8>> should reuse the exclusive Vec");
    }

    #[test]
    fn easybuf_into_vec_sliced() {
        let vec: Vec<u8> = (0u8..10u8).collect();
        let mut buf: EasyBuf = vec.into();
        let original_pointer = buf.buf.as_ref().as_ptr();
        buf.split_off(9);
        buf.drain_to(3);
        let result: Vec<u8> = buf.into();
        let reference: Vec<u8> = (3u8..9u8).collect();
        assert_eq!(result, reference);
        let new_pointer = result.as_ptr();
        assert_eq!(original_pointer, new_pointer, "Into<Vec<u8>> should reuse the exclusive Vec");
    }

    #[test]
    fn easybuf_into_vec_sliced_allocating() {
        let vec: Vec<u8> = (0u8..10u8).collect();
        let mut buf: EasyBuf = vec.into();
        let original_pointer = buf.buf.as_ref().as_ptr();
        // Create a clone to create second reference to this EasyBuf and force allocation
        let original = buf.clone();
        buf.split_off(9);
        buf.drain_to(3);
        let result: Vec<u8> = buf.into();
        let reference: Vec<u8> = (3u8..9u8).collect();
        assert_eq!(result, reference);
        let original_reference: EasyBuf =(0u8..10u8).collect::<Vec<u8>>().into();
        assert_eq!(original.as_ref(), original_reference.as_ref());
        let new_pointer = result.as_ptr();
        assert_ne!(original_pointer, new_pointer, "A new vec should be allocated");
    }

    #[test]
    fn easybuf_equality_same_underlying_vec() {
        let mut buf: EasyBuf = (0u8..10).collect::<Vec<_>>().into();
        assert_eq!(buf, buf);
        let other = buf.drain_to(5);
        assert_ne!(buf, other);

        let buf: EasyBuf = (0u8..5).collect::<Vec<_>>().into();
        assert_eq!(buf, other);
    }

    #[test]
    fn easybuf_equality_different_underlying_vec() {
        let mut buf: EasyBuf = (0u8..10).collect::<Vec<_>>().into();
        let mut other: EasyBuf = (0u8..10).collect::<Vec<_>>().into();
        assert_eq!(buf, other);

        buf = buf.drain_to(5);
        assert_ne!(buf, other);

        other = other.drain_to(5);
        assert_eq!(buf, other);
    }
}