blob: 151648ac23611f434aecba3e9306f8b979da9b4d [file] [log] [blame]
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);
}
}