mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-25 00:00:18 +02:00
codec: move into tokio-util (#1675)
Related to #1318, Tokio APIs that are "less stable" are moved into a new `tokio-util` crate. This crate will mirror `tokio` and provide additional APIs that may require a greater rate of breaking changes. As examples require `tokio-util`, they are moved into a separate crate (`examples`). This has the added advantage of being able to avoid example only dependencies in the `tokio` crate.
This commit is contained in:
Carl Lerche
@@ -0,0 +1,41 @@
|
||||
use crate::codec::decoder::Decoder;
|
||||
use crate::codec::encoder::Encoder;
|
||||
|
||||
use bytes::{BufMut, Bytes, BytesMut};
|
||||
use std::io;
|
||||
|
||||
/// A simple `Codec` implementation that just ships bytes around.
|
||||
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Default)]
|
||||
pub struct BytesCodec(());
|
||||
|
||||
impl BytesCodec {
|
||||
/// Creates a new `BytesCodec` for shipping around raw bytes.
|
||||
pub fn new() -> BytesCodec {
|
||||
BytesCodec(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for BytesCodec {
|
||||
type Item = BytesMut;
|
||||
type Error = io::Error;
|
||||
|
||||
fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<BytesMut>, io::Error> {
|
||||
if !buf.is_empty() {
|
||||
let len = buf.len();
|
||||
Ok(Some(buf.split_to(len)))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for BytesCodec {
|
||||
type Item = Bytes;
|
||||
type Error = io::Error;
|
||||
|
||||
fn encode(&mut self, data: Bytes, buf: &mut BytesMut) -> Result<(), io::Error> {
|
||||
buf.reserve(data.len());
|
||||
buf.put(data);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
use crate::codec::encoder::Encoder;
|
||||
use crate::codec::Framed;
|
||||
|
||||
use tokio_io::{AsyncRead, AsyncWrite};
|
||||
|
||||
use bytes::BytesMut;
|
||||
use std::io;
|
||||
|
||||
/// Decoding of frames via buffers.
|
||||
///
|
||||
/// This trait is used when constructing an instance of `Framed` or
|
||||
/// `FramedRead`. An implementation of `Decoder` takes a byte stream that has
|
||||
/// already been buffered in `src` and decodes the data into a stream of
|
||||
/// `Self::Item` frames.
|
||||
///
|
||||
/// Implementations are able to track state on `self`, which enables
|
||||
/// implementing stateful streaming parsers. In many cases, though, this type
|
||||
/// will simply be a unit struct (e.g. `struct HttpDecoder`).
|
||||
pub trait Decoder {
|
||||
/// The type of decoded frames.
|
||||
type Item;
|
||||
|
||||
/// The type of unrecoverable frame decoding errors.
|
||||
///
|
||||
/// If an individual message is ill-formed but can be ignored without
|
||||
/// interfering with the processing of future messages, it may be more
|
||||
/// useful to report the failure as an `Item`.
|
||||
///
|
||||
/// `From<io::Error>` is required in the interest of making `Error` suitable
|
||||
/// for returning directly from a `FramedRead`, and to enable the default
|
||||
/// implementation of `decode_eof` to yield an `io::Error` when the decoder
|
||||
/// fails to consume all available data.
|
||||
///
|
||||
/// Note that implementors of this trait can simply indicate `type Error =
|
||||
/// io::Error` to use I/O errors as this type.
|
||||
type Error: From<io::Error>;
|
||||
|
||||
/// Attempts to decode a frame from the provided buffer of bytes.
|
||||
///
|
||||
/// This method is called by `FramedRead` 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.
|
||||
///
|
||||
/// Note that the bytes provided may be empty. If a previous call to
|
||||
/// `decode` consumed all the bytes in the buffer then `decode` will be
|
||||
/// called again until it returns `Ok(None)`, indicating that more bytes need to
|
||||
/// be read.
|
||||
///
|
||||
/// 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.
|
||||
///
|
||||
/// # Buffer management
|
||||
///
|
||||
/// Before returning from the function, implementations should ensure that
|
||||
/// the buffer has appropriate capacity in anticipation of future calls to
|
||||
/// `decode`. Failing to do so leads to inefficiency.
|
||||
///
|
||||
/// For example, if frames have a fixed length, or if the length of the
|
||||
/// current frame is known from a header, a possible buffer management
|
||||
/// strategy is:
|
||||
///
|
||||
/// ```no_run
|
||||
/// # use std::io;
|
||||
/// #
|
||||
/// # use bytes::BytesMut;
|
||||
/// # use tokio_util::codec::Decoder;
|
||||
/// #
|
||||
/// # struct MyCodec;
|
||||
/// #
|
||||
/// impl Decoder for MyCodec {
|
||||
/// // ...
|
||||
/// # type Item = BytesMut;
|
||||
/// # type Error = io::Error;
|
||||
///
|
||||
/// fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
|
||||
/// // ...
|
||||
///
|
||||
/// // Reserve enough to complete decoding of the current frame.
|
||||
/// let current_frame_len: usize = 1000; // Example.
|
||||
/// // And to start decoding the next frame.
|
||||
/// let next_frame_header_len: usize = 10; // Example.
|
||||
/// src.reserve(current_frame_len + next_frame_header_len);
|
||||
///
|
||||
/// return Ok(None);
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// An optimal buffer management strategy minimizes reallocations and
|
||||
/// over-allocations.
|
||||
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::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 while there is unconsumed data in `buf`.
|
||||
/// Typically this doesn't need to be implemented unless the framing
|
||||
/// protocol differs near the end of the stream.
|
||||
///
|
||||
/// Note that the `buf` argument may be empty. If a previous call to
|
||||
/// `decode_eof` consumed all the bytes in the buffer, `decode_eof` will be
|
||||
/// called again until it returns `None`, indicating that there are no more
|
||||
/// frames to yield. This behavior enables returning finalization frames
|
||||
/// that may not be based on inbound data.
|
||||
fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
|
||||
match self.decode(buf)? {
|
||||
Some(frame) => Ok(Some(frame)),
|
||||
None => {
|
||||
if buf.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Err(io::Error::new(io::ErrorKind::Other, "bytes remaining on stream").into())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provides a `Stream` and `Sink` interface for reading and writing to this
|
||||
/// `Io` object, using `Decode` and `Encode` to read and write the raw data.
|
||||
///
|
||||
/// Raw I/O objects work with byte sequences, but higher-level code usually
|
||||
/// wants to batch these into meaningful chunks, called "frames". This
|
||||
/// method layers framing on top of an I/O object, by using the `Codec`
|
||||
/// 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 like gzip or TLS, 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 `Framed` returned by this method, which will
|
||||
/// break them into separate objects, allowing them to interact more easily.
|
||||
fn framed<T: AsyncRead + AsyncWrite + Sized>(self, io: T) -> Framed<T, Self>
|
||||
where
|
||||
Self: Encoder + Sized,
|
||||
{
|
||||
Framed::new(io, self)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
use bytes::BytesMut;
|
||||
use std::io;
|
||||
|
||||
/// Trait of helper objects to write out messages as bytes, for use with
|
||||
/// `FramedWrite`.
|
||||
pub trait Encoder {
|
||||
/// The type of items consumed by the `Encoder`
|
||||
type Item;
|
||||
|
||||
/// The type of encoding errors.
|
||||
///
|
||||
/// `FramedWrite` requires `Encoder`s errors to implement `From<io::Error>`
|
||||
/// in the interest letting it return `Error`s directly.
|
||||
type Error: From<io::Error>;
|
||||
|
||||
/// Encodes a frame into the buffer provided.
|
||||
///
|
||||
/// This method will encode `item` into the byte buffer provided by `dst`.
|
||||
/// The `dst` provided is an internal buffer of the `Framed` instance and
|
||||
/// will be written out when possible.
|
||||
fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error>;
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
use crate::codec::decoder::Decoder;
|
||||
use crate::codec::encoder::Encoder;
|
||||
use crate::codec::framed_read::{framed_read2, framed_read2_with_buffer, FramedRead2};
|
||||
use crate::codec::framed_write::{framed_write2, framed_write2_with_buffer, FramedWrite2};
|
||||
|
||||
use tokio_io::{AsyncBufRead, AsyncRead, AsyncWrite};
|
||||
|
||||
use bytes::BytesMut;
|
||||
use futures_core::Stream;
|
||||
use futures_sink::Sink;
|
||||
use std::fmt;
|
||||
use std::io::{self, BufRead, Read, Write};
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// A unified `Stream` and `Sink` interface to an underlying I/O object, using
|
||||
/// the `Encoder` and `Decoder` traits to encode and decode frames.
|
||||
///
|
||||
/// You can create a `Framed` instance by using the `AsyncRead::framed` adapter.
|
||||
pub struct Framed<T, U> {
|
||||
inner: FramedRead2<FramedWrite2<Fuse<T, U>>>,
|
||||
}
|
||||
|
||||
pub(crate) struct Fuse<T, U>(pub(crate) T, pub(crate) U);
|
||||
|
||||
impl<T, U> Framed<T, U>
|
||||
where
|
||||
T: AsyncRead + AsyncWrite,
|
||||
U: Decoder + Encoder,
|
||||
{
|
||||
/// Provides a `Stream` and `Sink` interface for reading and writing to this
|
||||
/// `Io` object, using `Decode` and `Encode` to read and write the raw data.
|
||||
///
|
||||
/// Raw I/O objects work with byte sequences, but higher-level code usually
|
||||
/// wants to batch these into meaningful chunks, called "frames". This
|
||||
/// method layers framing on top of an I/O object, by using the `Codec`
|
||||
/// 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 like gzip or TLS, 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 `Framed` returned by this method, which will
|
||||
/// break them into separate objects, allowing them to interact more easily.
|
||||
pub fn new(inner: T, codec: U) -> Framed<T, U> {
|
||||
Framed {
|
||||
inner: framed_read2(framed_write2(Fuse(inner, codec))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> Framed<T, U> {
|
||||
/// Provides a `Stream` and `Sink` interface for reading and writing to this
|
||||
/// `Io` object, using `Decode` and `Encode` to read and write the raw data.
|
||||
///
|
||||
/// Raw I/O objects work with byte sequences, but higher-level code usually
|
||||
/// wants to batch these into meaningful chunks, called "frames". This
|
||||
/// method layers framing on top of an I/O object, by using the `Codec`
|
||||
/// 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 like gzip or TLS, which require both read and write access to the
|
||||
/// underlying object.
|
||||
///
|
||||
/// This objects takes a stream and a readbuffer and a writebuffer. These field
|
||||
/// can be obtained from an existing `Framed` with the `into_parts` method.
|
||||
///
|
||||
/// If you want to work more directly with the streams and sink, consider
|
||||
/// calling `split` on the `Framed` returned by this method, which will
|
||||
/// break them into separate objects, allowing them to interact more easily.
|
||||
pub fn from_parts(parts: FramedParts<T, U>) -> Framed<T, U> {
|
||||
Framed {
|
||||
inner: framed_read2_with_buffer(
|
||||
framed_write2_with_buffer(Fuse(parts.io, parts.codec), parts.write_buf),
|
||||
parts.read_buf,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a reference to the underlying I/O stream wrapped by
|
||||
/// `Frame`.
|
||||
///
|
||||
/// 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.inner.get_ref().get_ref().0
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying I/O stream wrapped by
|
||||
/// `Frame`.
|
||||
///
|
||||
/// 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.inner.get_mut().get_mut().0
|
||||
}
|
||||
|
||||
/// Returns a reference to the underlying codec wrapped by
|
||||
/// `Frame`.
|
||||
///
|
||||
/// Note that care should be taken to not tamper with the underlying codec
|
||||
/// as it may corrupt the stream of frames otherwise being worked with.
|
||||
pub fn codec(&self) -> &U {
|
||||
&self.inner.get_ref().get_ref().1
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying codec wrapped by
|
||||
/// `Frame`.
|
||||
///
|
||||
/// Note that care should be taken to not tamper with the underlying codec
|
||||
/// as it may corrupt the stream of frames otherwise being worked with.
|
||||
pub fn codec_mut(&mut self) -> &mut U {
|
||||
&mut self.inner.get_mut().get_mut().1
|
||||
}
|
||||
|
||||
/// Consumes the `Frame`, 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.inner.into_inner().into_inner().0
|
||||
}
|
||||
|
||||
/// Consumes the `Frame`, returning its underlying I/O stream, the buffer
|
||||
/// with unprocessed data, and the codec.
|
||||
///
|
||||
/// 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_parts(self) -> FramedParts<T, U> {
|
||||
let (inner, read_buf) = self.inner.into_parts();
|
||||
let (inner, write_buf) = inner.into_parts();
|
||||
|
||||
FramedParts {
|
||||
io: inner.0,
|
||||
codec: inner.1,
|
||||
read_buf,
|
||||
write_buf,
|
||||
_priv: (),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> Stream for Framed<T, U>
|
||||
where
|
||||
T: AsyncRead + Unpin,
|
||||
U: Decoder + Unpin,
|
||||
{
|
||||
type Item = Result<U::Item, U::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
pin!(self.get_mut().inner).poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, I, U> Sink<I> for Framed<T, U>
|
||||
where
|
||||
T: AsyncWrite + Unpin,
|
||||
U: Encoder<Item = I> + Unpin,
|
||||
U::Error: From<io::Error>,
|
||||
{
|
||||
type Error = U::Error;
|
||||
|
||||
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Pin::new(Pin::get_mut(self).inner.get_mut()).poll_ready(cx)
|
||||
}
|
||||
|
||||
fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
|
||||
Pin::new(Pin::get_mut(self).inner.get_mut()).start_send(item)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Pin::new(Pin::get_mut(self).inner.get_mut()).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Pin::new(Pin::get_mut(self).inner.get_mut()).poll_close(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> fmt::Debug for Framed<T, U>
|
||||
where
|
||||
T: fmt::Debug,
|
||||
U: fmt::Debug,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("Framed")
|
||||
.field("io", &self.inner.get_ref().get_ref().0)
|
||||
.field("codec", &self.inner.get_ref().get_ref().1)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl Fuse =====
|
||||
|
||||
impl<T: Read, U> Read for Fuse<T, U> {
|
||||
fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
|
||||
self.0.read(dst)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: BufRead, U> BufRead for Fuse<T, U> {
|
||||
fn fill_buf(&mut self) -> io::Result<&[u8]> {
|
||||
self.0.fill_buf()
|
||||
}
|
||||
|
||||
fn consume(&mut self, amt: usize) {
|
||||
self.0.consume(amt)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsyncRead + Unpin, U: Unpin> AsyncRead for Fuse<T, U> {
|
||||
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
|
||||
self.0.prepare_uninitialized_buffer(buf)
|
||||
}
|
||||
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &mut [u8],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
pin!(self.get_mut().0).poll_read(cx, buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsyncBufRead + Unpin, U: Unpin> AsyncBufRead for Fuse<T, U> {
|
||||
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
|
||||
pin!(self.get_mut().0).poll_fill_buf(cx)
|
||||
}
|
||||
|
||||
fn consume(self: Pin<&mut Self>, amt: usize) {
|
||||
pin!(self.get_mut().0).consume(amt)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Write, U> Write for Fuse<T, U> {
|
||||
fn write(&mut self, src: &[u8]) -> io::Result<usize> {
|
||||
self.0.write(src)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.0.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsyncWrite + Unpin, U: Unpin> AsyncWrite for Fuse<T, U> {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
pin!(self.get_mut().0).poll_write(cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
|
||||
pin!(self.get_mut().0).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
|
||||
pin!(self.get_mut().0).poll_shutdown(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U: Decoder> Decoder for Fuse<T, U> {
|
||||
type Item = U::Item;
|
||||
type Error = U::Error;
|
||||
|
||||
fn decode(&mut self, buffer: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
|
||||
self.1.decode(buffer)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, buffer: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
|
||||
self.1.decode_eof(buffer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U: Encoder> Encoder for Fuse<T, U> {
|
||||
type Item = U::Item;
|
||||
type Error = U::Error;
|
||||
|
||||
fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
|
||||
self.1.encode(item, dst)
|
||||
}
|
||||
}
|
||||
|
||||
/// `FramedParts` contains an export of the data of a Framed transport.
|
||||
/// It can be used to construct a new `Framed` with a different codec.
|
||||
/// It contains all current buffers and the inner transport.
|
||||
#[derive(Debug)]
|
||||
pub struct FramedParts<T, U> {
|
||||
/// The inner transport used to read bytes to and write bytes to
|
||||
pub io: T,
|
||||
|
||||
/// The codec
|
||||
pub codec: U,
|
||||
|
||||
/// The buffer with read but unprocessed data.
|
||||
pub read_buf: BytesMut,
|
||||
|
||||
/// A buffer with unprocessed data which are not written yet.
|
||||
pub write_buf: BytesMut,
|
||||
|
||||
/// This private field allows us to add additional fields in the future in a
|
||||
/// backwards compatible way.
|
||||
_priv: (),
|
||||
}
|
||||
|
||||
impl<T, U> FramedParts<T, U> {
|
||||
/// Create a new, default, `FramedParts`
|
||||
pub fn new(io: T, codec: U) -> FramedParts<T, U> {
|
||||
FramedParts {
|
||||
io,
|
||||
codec,
|
||||
read_buf: BytesMut::new(),
|
||||
write_buf: BytesMut::new(),
|
||||
_priv: (),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
use crate::codec::framed::Fuse;
|
||||
use crate::codec::Decoder;
|
||||
|
||||
use tokio_io::AsyncRead;
|
||||
|
||||
use bytes::BytesMut;
|
||||
use futures_core::Stream;
|
||||
use futures_sink::Sink;
|
||||
use log::trace;
|
||||
use std::fmt;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// A `Stream` of messages decoded from an `AsyncRead`.
|
||||
pub struct FramedRead<T, D> {
|
||||
inner: FramedRead2<Fuse<T, D>>,
|
||||
}
|
||||
|
||||
pub(crate) struct FramedRead2<T> {
|
||||
inner: T,
|
||||
eof: bool,
|
||||
is_readable: bool,
|
||||
buffer: BytesMut,
|
||||
}
|
||||
|
||||
const INITIAL_CAPACITY: usize = 8 * 1024;
|
||||
|
||||
// ===== impl FramedRead =====
|
||||
|
||||
impl<T, D> FramedRead<T, D>
|
||||
where
|
||||
T: AsyncRead,
|
||||
D: Decoder,
|
||||
{
|
||||
/// Creates a new `FramedRead` with the given `decoder`.
|
||||
pub fn new(inner: T, decoder: D) -> FramedRead<T, D> {
|
||||
FramedRead {
|
||||
inner: framed_read2(Fuse(inner, decoder)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, D> FramedRead<T, D> {
|
||||
/// Returns a reference to the underlying I/O stream wrapped by
|
||||
/// `FramedRead`.
|
||||
///
|
||||
/// 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.inner.inner.0
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying I/O stream wrapped by
|
||||
/// `FramedRead`.
|
||||
///
|
||||
/// 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.inner.inner.0
|
||||
}
|
||||
|
||||
/// Consumes the `FramedRead`, 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.inner.inner.0
|
||||
}
|
||||
|
||||
/// Returns a reference to the underlying decoder.
|
||||
pub fn decoder(&self) -> &D {
|
||||
&self.inner.inner.1
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying decoder.
|
||||
pub fn decoder_mut(&mut self) -> &mut D {
|
||||
&mut self.inner.inner.1
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, D> Stream for FramedRead<T, D>
|
||||
where
|
||||
T: AsyncRead + Unpin,
|
||||
D: Decoder + Unpin,
|
||||
{
|
||||
type Item = Result<D::Item, D::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
pin!(self.get_mut().inner).poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
// This impl just defers to the underlying T: Sink
|
||||
impl<T, I, D> Sink<I> for FramedRead<T, D>
|
||||
where
|
||||
T: Sink<I> + Unpin,
|
||||
D: Unpin,
|
||||
{
|
||||
type Error = T::Error;
|
||||
|
||||
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
pin!(Pin::get_mut(self).inner.inner.0).poll_ready(cx)
|
||||
}
|
||||
|
||||
fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
|
||||
pin!(Pin::get_mut(self).inner.inner.0).start_send(item)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
pin!(Pin::get_mut(self).inner.inner.0).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
pin!(Pin::get_mut(self).inner.inner.0).poll_close(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, D> fmt::Debug for FramedRead<T, D>
|
||||
where
|
||||
T: fmt::Debug,
|
||||
D: fmt::Debug,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("FramedRead")
|
||||
.field("inner", &self.inner.inner.0)
|
||||
.field("decoder", &self.inner.inner.1)
|
||||
.field("eof", &self.inner.eof)
|
||||
.field("is_readable", &self.inner.is_readable)
|
||||
.field("buffer", &self.inner.buffer)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl FramedRead2 =====
|
||||
|
||||
pub(crate) fn framed_read2<T>(inner: T) -> FramedRead2<T> {
|
||||
FramedRead2 {
|
||||
inner,
|
||||
eof: false,
|
||||
is_readable: false,
|
||||
buffer: BytesMut::with_capacity(INITIAL_CAPACITY),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn framed_read2_with_buffer<T>(inner: T, mut buf: BytesMut) -> FramedRead2<T> {
|
||||
if buf.capacity() < INITIAL_CAPACITY {
|
||||
let bytes_to_reserve = INITIAL_CAPACITY - buf.capacity();
|
||||
buf.reserve(bytes_to_reserve);
|
||||
}
|
||||
FramedRead2 {
|
||||
inner,
|
||||
eof: false,
|
||||
is_readable: !buf.is_empty(),
|
||||
buffer: buf,
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> FramedRead2<T> {
|
||||
pub(crate) fn get_ref(&self) -> &T {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
pub(crate) fn into_inner(self) -> T {
|
||||
self.inner
|
||||
}
|
||||
|
||||
pub(crate) fn into_parts(self) -> (T, BytesMut) {
|
||||
(self.inner, self.buffer)
|
||||
}
|
||||
|
||||
pub(crate) fn get_mut(&mut self) -> &mut T {
|
||||
&mut self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Stream for FramedRead2<T>
|
||||
where
|
||||
T: AsyncRead + Decoder + Unpin,
|
||||
{
|
||||
type Item = Result<T::Item, T::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let pinned = Pin::get_mut(self);
|
||||
loop {
|
||||
// Repeatedly call `decode` or `decode_eof` as long as it is
|
||||
// "readable". Readable is defined as not having returned `None`. If
|
||||
// the upstream has returned EOF, and the decoder is no longer
|
||||
// readable, it can be assumed that the decoder will never become
|
||||
// readable again, at which point the stream is terminated.
|
||||
if pinned.is_readable {
|
||||
if pinned.eof {
|
||||
let frame = pinned.inner.decode_eof(&mut pinned.buffer)?;
|
||||
return Poll::Ready(frame.map(Ok));
|
||||
}
|
||||
|
||||
trace!("attempting to decode a frame");
|
||||
|
||||
if let Some(frame) = pinned.inner.decode(&mut pinned.buffer)? {
|
||||
trace!("frame decoded from buffer");
|
||||
return Poll::Ready(Some(Ok(frame)));
|
||||
}
|
||||
|
||||
pinned.is_readable = false;
|
||||
}
|
||||
|
||||
assert!(!pinned.eof);
|
||||
|
||||
// Otherwise, try to read more data and try again. Make sure we've
|
||||
// got room for at least one byte to read to ensure that we don't
|
||||
// get a spurious 0 that looks like EOF
|
||||
pinned.buffer.reserve(1);
|
||||
let bytect = match pin!(pinned.inner).poll_read_buf(cx, &mut pinned.buffer)? {
|
||||
Poll::Ready(ct) => ct,
|
||||
Poll::Pending => return Poll::Pending,
|
||||
};
|
||||
if bytect == 0 {
|
||||
pinned.eof = true;
|
||||
}
|
||||
|
||||
pinned.is_readable = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,288 @@
|
||||
use crate::codec::decoder::Decoder;
|
||||
use crate::codec::encoder::Encoder;
|
||||
use crate::codec::framed::Fuse;
|
||||
|
||||
use tokio_io::{AsyncBufRead, AsyncRead, AsyncWrite};
|
||||
|
||||
use bytes::BytesMut;
|
||||
use futures_core::{ready, Stream};
|
||||
use futures_sink::Sink;
|
||||
use log::trace;
|
||||
use std::fmt;
|
||||
use std::io::{self, BufRead, Read};
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// A `Sink` of frames encoded to an `AsyncWrite`.
|
||||
pub struct FramedWrite<T, E> {
|
||||
inner: FramedWrite2<Fuse<T, E>>,
|
||||
}
|
||||
|
||||
pub(crate) struct FramedWrite2<T> {
|
||||
inner: T,
|
||||
buffer: BytesMut,
|
||||
}
|
||||
|
||||
const INITIAL_CAPACITY: usize = 8 * 1024;
|
||||
const BACKPRESSURE_BOUNDARY: usize = INITIAL_CAPACITY;
|
||||
|
||||
impl<T, E> FramedWrite<T, E>
|
||||
where
|
||||
T: AsyncWrite,
|
||||
E: Encoder,
|
||||
{
|
||||
/// Creates a new `FramedWrite` with the given `encoder`.
|
||||
pub fn new(inner: T, encoder: E) -> FramedWrite<T, E> {
|
||||
FramedWrite {
|
||||
inner: framed_write2(Fuse(inner, encoder)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, E> FramedWrite<T, E> {
|
||||
/// Returns a reference to the underlying I/O stream wrapped by
|
||||
/// `FramedWrite`.
|
||||
///
|
||||
/// 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.inner.inner.0
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying I/O stream wrapped by
|
||||
/// `FramedWrite`.
|
||||
///
|
||||
/// 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.inner.inner.0
|
||||
}
|
||||
|
||||
/// Consumes the `FramedWrite`, 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.inner.inner.0
|
||||
}
|
||||
|
||||
/// Returns a reference to the underlying decoder.
|
||||
pub fn encoder(&self) -> &E {
|
||||
&self.inner.inner.1
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying decoder.
|
||||
pub fn encoder_mut(&mut self) -> &mut E {
|
||||
&mut self.inner.inner.1
|
||||
}
|
||||
}
|
||||
|
||||
// This impl just defers to the underlying FramedWrite2
|
||||
impl<T, I, E> Sink<I> for FramedWrite<T, E>
|
||||
where
|
||||
T: AsyncWrite + Unpin,
|
||||
E: Encoder<Item = I> + Unpin,
|
||||
E::Error: From<io::Error>,
|
||||
{
|
||||
type Error = E::Error;
|
||||
|
||||
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
pin!(Pin::get_mut(self).inner).poll_ready(cx)
|
||||
}
|
||||
|
||||
fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
|
||||
pin!(Pin::get_mut(self).inner).start_send(item)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
pin!(Pin::get_mut(self).inner).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
pin!(Pin::get_mut(self).inner).poll_close(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, D> Stream for FramedWrite<T, D>
|
||||
where
|
||||
T: Stream + Unpin,
|
||||
D: Unpin,
|
||||
{
|
||||
type Item = T::Item;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
Pin::new(Pin::get_mut(self).get_mut()).poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> fmt::Debug for FramedWrite<T, U>
|
||||
where
|
||||
T: fmt::Debug,
|
||||
U: fmt::Debug,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("FramedWrite")
|
||||
.field("inner", &self.inner.get_ref().0)
|
||||
.field("encoder", &self.inner.get_ref().1)
|
||||
.field("buffer", &self.inner.buffer)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl FramedWrite2 =====
|
||||
|
||||
pub(crate) fn framed_write2<T>(inner: T) -> FramedWrite2<T> {
|
||||
FramedWrite2 {
|
||||
inner,
|
||||
buffer: BytesMut::with_capacity(INITIAL_CAPACITY),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn framed_write2_with_buffer<T>(inner: T, mut buf: BytesMut) -> FramedWrite2<T> {
|
||||
if buf.capacity() < INITIAL_CAPACITY {
|
||||
let bytes_to_reserve = INITIAL_CAPACITY - buf.capacity();
|
||||
buf.reserve(bytes_to_reserve);
|
||||
}
|
||||
FramedWrite2 { inner, buffer: buf }
|
||||
}
|
||||
|
||||
impl<T> FramedWrite2<T> {
|
||||
pub(crate) fn get_ref(&self) -> &T {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
pub(crate) fn into_inner(self) -> T {
|
||||
self.inner
|
||||
}
|
||||
|
||||
pub(crate) fn into_parts(self) -> (T, BytesMut) {
|
||||
(self.inner, self.buffer)
|
||||
}
|
||||
|
||||
pub(crate) fn get_mut(&mut self) -> &mut T {
|
||||
&mut self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl<I, T> Sink<I> for FramedWrite2<T>
|
||||
where
|
||||
T: AsyncWrite + Encoder<Item = I> + Unpin,
|
||||
{
|
||||
type Error = T::Error;
|
||||
|
||||
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::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).
|
||||
if self.buffer.len() >= BACKPRESSURE_BOUNDARY {
|
||||
match self.as_mut().poll_flush(cx) {
|
||||
Poll::Pending => return Poll::Pending,
|
||||
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
|
||||
Poll::Ready(Ok(())) => (),
|
||||
};
|
||||
|
||||
if self.buffer.len() >= BACKPRESSURE_BOUNDARY {
|
||||
return Poll::Pending;
|
||||
}
|
||||
}
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
|
||||
let pinned = Pin::get_mut(self);
|
||||
pinned.inner.encode(item, &mut pinned.buffer)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
trace!("flushing framed transport");
|
||||
let pinned = Pin::get_mut(self);
|
||||
|
||||
while !pinned.buffer.is_empty() {
|
||||
trace!("writing; remaining={}", pinned.buffer.len());
|
||||
|
||||
let buf = &pinned.buffer;
|
||||
let n = ready!(pin!(pinned.inner).poll_write(cx, &buf))?;
|
||||
|
||||
if n == 0 {
|
||||
return Poll::Ready(Err(io::Error::new(
|
||||
io::ErrorKind::WriteZero,
|
||||
"failed to \
|
||||
write frame to transport",
|
||||
)
|
||||
.into()));
|
||||
}
|
||||
|
||||
// TODO: Add a way to `bytes` to do this w/o returning the drained data.
|
||||
let _ = pinned.buffer.split_to(n);
|
||||
}
|
||||
|
||||
// Try flushing the underlying IO
|
||||
ready!(pin!(pinned.inner).poll_flush(cx))?;
|
||||
|
||||
trace!("framed transport flushed");
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
ready!(pin!(self).poll_flush(cx))?;
|
||||
ready!(pin!(self.inner).poll_shutdown(cx))?;
|
||||
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Decoder> Decoder for FramedWrite2<T> {
|
||||
type Item = T::Item;
|
||||
type Error = T::Error;
|
||||
|
||||
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<T::Item>, T::Error> {
|
||||
self.inner.decode(src)
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, src: &mut BytesMut) -> Result<Option<T::Item>, T::Error> {
|
||||
self.inner.decode_eof(src)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Read> Read for FramedWrite2<T> {
|
||||
fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
|
||||
self.inner.read(dst)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: BufRead> BufRead for FramedWrite2<T> {
|
||||
fn fill_buf(&mut self) -> io::Result<&[u8]> {
|
||||
self.inner.fill_buf()
|
||||
}
|
||||
|
||||
fn consume(&mut self, amt: usize) {
|
||||
self.inner.consume(amt)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsyncRead + Unpin> AsyncRead for FramedWrite2<T> {
|
||||
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
|
||||
self.inner.prepare_uninitialized_buffer(buf)
|
||||
}
|
||||
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &mut [u8],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
pin!(self.get_mut().inner).poll_read(cx, buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsyncBufRead + Unpin> AsyncBufRead for FramedWrite2<T> {
|
||||
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
|
||||
pin!(self.get_mut().inner).poll_fill_buf(cx)
|
||||
}
|
||||
|
||||
fn consume(self: Pin<&mut Self>, amt: usize) {
|
||||
pin!(self.get_mut().inner).consume(amt)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,962 @@
|
||||
//! Frame a stream of bytes based on a length prefix
|
||||
//!
|
||||
//! Many protocols delimit their frames by prefacing frame data with a
|
||||
//! frame head that specifies the length of the frame. The
|
||||
//! `length_delimited` module provides utilities for handling the length
|
||||
//! based framing. This allows the consumer to work with entire frames
|
||||
//! without having to worry about buffering or other framing logic.
|
||||
//!
|
||||
//! # Getting started
|
||||
//!
|
||||
//! If implementing a protocol from scratch, using length delimited framing
|
||||
//! is an easy way to get started. [`LengthDelimitedCodec::new()`] will
|
||||
//! return a length delimited codec using default configuration values.
|
||||
//! This can then be used to construct a framer to adapt a full-duplex
|
||||
//! byte stream into a stream of frames.
|
||||
//!
|
||||
//! ```
|
||||
//! use tokio_io::{AsyncRead, AsyncWrite};
|
||||
//! use tokio_util::codec::{Framed, LengthDelimitedCodec};
|
||||
//!
|
||||
//! fn bind_transport<T: AsyncRead + AsyncWrite>(io: T)
|
||||
//! -> Framed<T, LengthDelimitedCodec>
|
||||
//! {
|
||||
//! Framed::new(io, LengthDelimitedCodec::new())
|
||||
//! }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! The returned transport implements `Sink + Stream` for `BytesMut`. It
|
||||
//! encodes the frame with a big-endian `u32` header denoting the frame
|
||||
//! payload length:
|
||||
//!
|
||||
//! ```text
|
||||
//! +----------+--------------------------------+
|
||||
//! | len: u32 | frame payload |
|
||||
//! +----------+--------------------------------+
|
||||
//! ```
|
||||
//!
|
||||
//! Specifically, given the following:
|
||||
//!
|
||||
//! ```
|
||||
//! use tokio::prelude::*;
|
||||
//! use tokio_util::codec::{Framed, LengthDelimitedCodec};
|
||||
//!
|
||||
//! use bytes::Bytes;
|
||||
//!
|
||||
//! async fn write_frame<T>(io: T) -> Result<(), Box<dyn std::error::Error>>
|
||||
//! where
|
||||
//! T: AsyncRead + AsyncWrite + Unpin,
|
||||
//! {
|
||||
//! let mut transport = Framed::new(io, LengthDelimitedCodec::new());
|
||||
//! let frame = Bytes::from("hello world");
|
||||
//!
|
||||
//! transport.send(frame).await?;
|
||||
//! Ok(())
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! The encoded frame will look like this:
|
||||
//!
|
||||
//! ```text
|
||||
//! +---- len: u32 ----+---- data ----+
|
||||
//! | \x00\x00\x00\x0b | hello world |
|
||||
//! +------------------+--------------+
|
||||
//! ```
|
||||
//!
|
||||
//! # Decoding
|
||||
//!
|
||||
//! [`FramedRead`] adapts an [`AsyncRead`] into a `Stream` of [`BytesMut`],
|
||||
//! such that each yielded [`BytesMut`] value contains the contents of an
|
||||
//! entire frame. There are many configuration parameters enabling
|
||||
//! [`FramedRead`] to handle a wide range of protocols. Here are some
|
||||
//! examples that will cover the various options at a high level.
|
||||
//!
|
||||
//! ## Example 1
|
||||
//!
|
||||
//! The following will parse a `u16` length field at offset 0, including the
|
||||
//! frame head in the yielded `BytesMut`.
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncRead;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn bind_read<T: AsyncRead>(io: T) {
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_offset(0) // default value
|
||||
//! .length_field_length(2)
|
||||
//! .length_adjustment(0) // default value
|
||||
//! .num_skip(0) // Do not strip frame header
|
||||
//! .new_read(io);
|
||||
//! # }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! The following frame will be decoded as such:
|
||||
//!
|
||||
//! ```text
|
||||
//! INPUT DECODED
|
||||
//! +-- len ---+--- Payload ---+ +-- len ---+--- Payload ---+
|
||||
//! | \x00\x0B | Hello world | --> | \x00\x0B | Hello world |
|
||||
//! +----------+---------------+ +----------+---------------+
|
||||
//! ```
|
||||
//!
|
||||
//! The value of the length field is 11 (`\x0B`) which represents the length
|
||||
//! of the payload, `hello world`. By default, [`FramedRead`] assumes that
|
||||
//! the length field represents the number of bytes that **follows** the
|
||||
//! length field. Thus, the entire frame has a length of 13: 2 bytes for the
|
||||
//! frame head + 11 bytes for the payload.
|
||||
//!
|
||||
//! ## Example 2
|
||||
//!
|
||||
//! The following will parse a `u16` length field at offset 0, omitting the
|
||||
//! frame head in the yielded `BytesMut`.
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncRead;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn bind_read<T: AsyncRead>(io: T) {
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_offset(0) // default value
|
||||
//! .length_field_length(2)
|
||||
//! .length_adjustment(0) // default value
|
||||
//! // `num_skip` is not needed, the default is to skip
|
||||
//! .new_read(io);
|
||||
//! # }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! The following frame will be decoded as such:
|
||||
//!
|
||||
//! ```text
|
||||
//! INPUT DECODED
|
||||
//! +-- len ---+--- Payload ---+ +--- Payload ---+
|
||||
//! | \x00\x0B | Hello world | --> | Hello world |
|
||||
//! +----------+---------------+ +---------------+
|
||||
//! ```
|
||||
//!
|
||||
//! This is similar to the first example, the only difference is that the
|
||||
//! frame head is **not** included in the yielded `BytesMut` value.
|
||||
//!
|
||||
//! ## Example 3
|
||||
//!
|
||||
//! The following will parse a `u16` length field at offset 0, including the
|
||||
//! frame head in the yielded `BytesMut`. In this case, the length field
|
||||
//! **includes** the frame head length.
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncRead;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn bind_read<T: AsyncRead>(io: T) {
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_offset(0) // default value
|
||||
//! .length_field_length(2)
|
||||
//! .length_adjustment(-2) // size of head
|
||||
//! .num_skip(0)
|
||||
//! .new_read(io);
|
||||
//! # }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! The following frame will be decoded as such:
|
||||
//!
|
||||
//! ```text
|
||||
//! INPUT DECODED
|
||||
//! +-- len ---+--- Payload ---+ +-- len ---+--- Payload ---+
|
||||
//! | \x00\x0D | Hello world | --> | \x00\x0D | Hello world |
|
||||
//! +----------+---------------+ +----------+---------------+
|
||||
//! ```
|
||||
//!
|
||||
//! In most cases, the length field represents the length of the payload
|
||||
//! only, as shown in the previous examples. However, in some protocols the
|
||||
//! length field represents the length of the whole frame, including the
|
||||
//! head. In such cases, we specify a negative `length_adjustment` to adjust
|
||||
//! the value provided in the frame head to represent the payload length.
|
||||
//!
|
||||
//! ## Example 4
|
||||
//!
|
||||
//! The following will parse a 3 byte length field at offset 0 in a 5 byte
|
||||
//! frame head, including the frame head in the yielded `BytesMut`.
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncRead;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn bind_read<T: AsyncRead>(io: T) {
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_offset(0) // default value
|
||||
//! .length_field_length(3)
|
||||
//! .length_adjustment(2) // remaining head
|
||||
//! .num_skip(0)
|
||||
//! .new_read(io);
|
||||
//! # }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! The following frame will be decoded as such:
|
||||
//!
|
||||
//! ```text
|
||||
//! INPUT
|
||||
//! +---- len -----+- head -+--- Payload ---+
|
||||
//! | \x00\x00\x0B | \xCAFE | Hello world |
|
||||
//! +--------------+--------+---------------+
|
||||
//!
|
||||
//! DECODED
|
||||
//! +---- len -----+- head -+--- Payload ---+
|
||||
//! | \x00\x00\x0B | \xCAFE | Hello world |
|
||||
//! +--------------+--------+---------------+
|
||||
//! ```
|
||||
//!
|
||||
//! A more advanced example that shows a case where there is extra frame
|
||||
//! head data between the length field and the payload. In such cases, it is
|
||||
//! usually desirable to include the frame head as part of the yielded
|
||||
//! `BytesMut`. This lets consumers of the length delimited framer to
|
||||
//! process the frame head as needed.
|
||||
//!
|
||||
//! The positive `length_adjustment` value lets `FramedRead` factor in the
|
||||
//! additional head into the frame length calculation.
|
||||
//!
|
||||
//! ## Example 5
|
||||
//!
|
||||
//! The following will parse a `u16` length field at offset 1 of a 4 byte
|
||||
//! frame head. The first byte and the length field will be omitted from the
|
||||
//! yielded `BytesMut`, but the trailing 2 bytes of the frame head will be
|
||||
//! included.
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncRead;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn bind_read<T: AsyncRead>(io: T) {
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_offset(1) // length of hdr1
|
||||
//! .length_field_length(2)
|
||||
//! .length_adjustment(1) // length of hdr2
|
||||
//! .num_skip(3) // length of hdr1 + LEN
|
||||
//! .new_read(io);
|
||||
//! # }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! The following frame will be decoded as such:
|
||||
//!
|
||||
//! ```text
|
||||
//! INPUT
|
||||
//! +- hdr1 -+-- len ---+- hdr2 -+--- Payload ---+
|
||||
//! | \xCA | \x00\x0B | \xFE | Hello world |
|
||||
//! +--------+----------+--------+---------------+
|
||||
//!
|
||||
//! DECODED
|
||||
//! +- hdr2 -+--- Payload ---+
|
||||
//! | \xFE | Hello world |
|
||||
//! +--------+---------------+
|
||||
//! ```
|
||||
//!
|
||||
//! The length field is situated in the middle of the frame head. In this
|
||||
//! case, the first byte in the frame head could be a version or some other
|
||||
//! identifier that is not needed for processing. On the other hand, the
|
||||
//! second half of the head is needed.
|
||||
//!
|
||||
//! `length_field_offset` indicates how many bytes to skip before starting
|
||||
//! to read the length field. `length_adjustment` is the number of bytes to
|
||||
//! skip starting at the end of the length field. In this case, it is the
|
||||
//! second half of the head.
|
||||
//!
|
||||
//! ## Example 6
|
||||
//!
|
||||
//! The following will parse a `u16` length field at offset 1 of a 4 byte
|
||||
//! frame head. The first byte and the length field will be omitted from the
|
||||
//! yielded `BytesMut`, but the trailing 2 bytes of the frame head will be
|
||||
//! included. In this case, the length field **includes** the frame head
|
||||
//! length.
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncRead;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn bind_read<T: AsyncRead>(io: T) {
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_offset(1) // length of hdr1
|
||||
//! .length_field_length(2)
|
||||
//! .length_adjustment(-3) // length of hdr1 + LEN, negative
|
||||
//! .num_skip(3)
|
||||
//! .new_read(io);
|
||||
//! # }
|
||||
//! ```
|
||||
//!
|
||||
//! The following frame will be decoded as such:
|
||||
//!
|
||||
//! ```text
|
||||
//! INPUT
|
||||
//! +- hdr1 -+-- len ---+- hdr2 -+--- Payload ---+
|
||||
//! | \xCA | \x00\x0F | \xFE | Hello world |
|
||||
//! +--------+----------+--------+---------------+
|
||||
//!
|
||||
//! DECODED
|
||||
//! +- hdr2 -+--- Payload ---+
|
||||
//! | \xFE | Hello world |
|
||||
//! +--------+---------------+
|
||||
//! ```
|
||||
//!
|
||||
//! Similar to the example above, the difference is that the length field
|
||||
//! represents the length of the entire frame instead of just the payload.
|
||||
//! The length of `hdr1` and `len` must be counted in `length_adjustment`.
|
||||
//! Note that the length of `hdr2` does **not** need to be explicitly set
|
||||
//! anywhere because it already is factored into the total frame length that
|
||||
//! is read from the byte stream.
|
||||
//!
|
||||
//! # Encoding
|
||||
//!
|
||||
//! [`FramedWrite`] adapts an [`AsyncWrite`] into a `Sink` of [`BytesMut`],
|
||||
//! such that each submitted [`BytesMut`] is prefaced by a length field.
|
||||
//! There are fewer configuration options than [`FramedRead`]. Given
|
||||
//! protocols that have more complex frame heads, an encoder should probably
|
||||
//! be written by hand using [`Encoder`].
|
||||
//!
|
||||
//! Here is a simple example, given a `FramedWrite` with the following
|
||||
//! configuration:
|
||||
//!
|
||||
//! ```
|
||||
//! # use tokio_io::AsyncWrite;
|
||||
//! # use tokio_util::codec::LengthDelimitedCodec;
|
||||
//! # fn write_frame<T: AsyncWrite>(io: T) {
|
||||
//! # let _ =
|
||||
//! LengthDelimitedCodec::builder()
|
||||
//! .length_field_length(2)
|
||||
//! .new_write(io);
|
||||
//! # }
|
||||
//! # pub fn main() {}
|
||||
//! ```
|
||||
//!
|
||||
//! A payload of `hello world` will be encoded as:
|
||||
//!
|
||||
//! ```text
|
||||
//! +- len: u16 -+---- data ----+
|
||||
//! | \x00\x0b | hello world |
|
||||
//! +------------+--------------+
|
||||
//! ```
|
||||
//!
|
||||
//! [`LengthDelimitedCodec::new()`]: struct.LengthDelimitedCodec.html#method.new
|
||||
//! [`FramedRead`]: struct.FramedRead.html
|
||||
//! [`FramedWrite`]: struct.FramedWrite.html
|
||||
//! [`AsyncRead`]: ../../trait.AsyncRead.html
|
||||
//! [`AsyncWrite`]: ../../trait.AsyncWrite.html
|
||||
//! [`Encoder`]: ../trait.Encoder.html
|
||||
//! [`BytesMut`]: https://docs.rs/bytes/0.4/bytes/struct.BytesMut.html
|
||||
|
||||
use crate::codec::{Decoder, Encoder, Framed, FramedRead, FramedWrite};
|
||||
|
||||
use tokio_io::{AsyncRead, AsyncWrite};
|
||||
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut, IntoBuf};
|
||||
use std::error::Error as StdError;
|
||||
use std::io::{self, Cursor};
|
||||
use std::{cmp, fmt};
|
||||
|
||||
/// Configure length delimited `LengthDelimitedCodec`s.
|
||||
///
|
||||
/// `Builder` enables constructing configured length delimited codecs. Note
|
||||
/// that not all configuration settings apply to both encoding and decoding. See
|
||||
/// the documentation for specific methods for more detail.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct Builder {
|
||||
// Maximum frame length
|
||||
max_frame_len: usize,
|
||||
|
||||
// Number of bytes representing the field length
|
||||
length_field_len: usize,
|
||||
|
||||
// Number of bytes in the header before the length field
|
||||
length_field_offset: usize,
|
||||
|
||||
// Adjust the length specified in the header field by this amount
|
||||
length_adjustment: isize,
|
||||
|
||||
// Total number of bytes to skip before reading the payload, if not set,
|
||||
// `length_field_len + length_field_offset`
|
||||
num_skip: Option<usize>,
|
||||
|
||||
// Length field byte order (little or big endian)
|
||||
length_field_is_big_endian: bool,
|
||||
}
|
||||
|
||||
/// An error when the number of bytes read is more than max frame length.
|
||||
pub struct LengthDelimitedCodecError {
|
||||
_priv: (),
|
||||
}
|
||||
|
||||
/// A codec for frames delimited by a frame head specifying their lengths.
|
||||
///
|
||||
/// This allows the consumer to work with entire frames without having to worry
|
||||
/// about buffering or other framing logic.
|
||||
///
|
||||
/// See [module level] documentation for more detail.
|
||||
///
|
||||
/// [module level]: index.html
|
||||
#[derive(Debug)]
|
||||
pub struct LengthDelimitedCodec {
|
||||
// Configuration values
|
||||
builder: Builder,
|
||||
|
||||
// Read state
|
||||
state: DecodeState,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum DecodeState {
|
||||
Head,
|
||||
Data(usize),
|
||||
}
|
||||
|
||||
// ===== impl LengthDelimitedCodec ======
|
||||
|
||||
impl LengthDelimitedCodec {
|
||||
/// Creates a new `LengthDelimitedCodec` with the default configuration values.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
builder: Builder::new(),
|
||||
state: DecodeState::Head,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a new length delimited codec builder with default configuration
|
||||
/// values.
|
||||
pub fn builder() -> Builder {
|
||||
Builder::new()
|
||||
}
|
||||
|
||||
/// Returns the current max frame setting
|
||||
///
|
||||
/// This is the largest size this codec will accept from the wire. Larger
|
||||
/// frames will be rejected.
|
||||
pub fn max_frame_length(&self) -> usize {
|
||||
self.builder.max_frame_len
|
||||
}
|
||||
|
||||
/// Updates the max frame setting.
|
||||
///
|
||||
/// The change takes effect the next time a frame is decoded. In other
|
||||
/// words, if a frame is currently in process of being decoded with a frame
|
||||
/// size greater than `val` but less than the max frame length in effect
|
||||
/// before calling this function, then the frame will be allowed.
|
||||
pub fn set_max_frame_length(&mut self, val: usize) {
|
||||
self.builder.max_frame_length(val);
|
||||
}
|
||||
|
||||
fn decode_head(&mut self, src: &mut BytesMut) -> io::Result<Option<usize>> {
|
||||
let head_len = self.builder.num_head_bytes();
|
||||
let field_len = self.builder.length_field_len;
|
||||
|
||||
if src.len() < head_len {
|
||||
// Not enough data
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let n = {
|
||||
let mut src = Cursor::new(&mut *src);
|
||||
|
||||
// Skip the required bytes
|
||||
src.advance(self.builder.length_field_offset);
|
||||
|
||||
// match endianess
|
||||
let n = if self.builder.length_field_is_big_endian {
|
||||
src.get_uint_be(field_len)
|
||||
} else {
|
||||
src.get_uint_le(field_len)
|
||||
};
|
||||
|
||||
if n > self.builder.max_frame_len as u64 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
LengthDelimitedCodecError { _priv: () },
|
||||
));
|
||||
}
|
||||
|
||||
// The check above ensures there is no overflow
|
||||
let n = n as usize;
|
||||
|
||||
// Adjust `n` with bounds checking
|
||||
let n = if self.builder.length_adjustment < 0 {
|
||||
n.checked_sub(-self.builder.length_adjustment as usize)
|
||||
} else {
|
||||
n.checked_add(self.builder.length_adjustment as usize)
|
||||
};
|
||||
|
||||
// Error handling
|
||||
match n {
|
||||
Some(n) => n,
|
||||
None => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"provided length would overflow after adjustment",
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let num_skip = self.builder.get_num_skip();
|
||||
|
||||
if num_skip > 0 {
|
||||
let _ = src.split_to(num_skip);
|
||||
}
|
||||
|
||||
// Ensure that the buffer has enough space to read the incoming
|
||||
// payload
|
||||
src.reserve(n);
|
||||
|
||||
Ok(Some(n))
|
||||
}
|
||||
|
||||
fn decode_data(&self, n: usize, src: &mut BytesMut) -> io::Result<Option<BytesMut>> {
|
||||
// At this point, the buffer has already had the required capacity
|
||||
// reserved. All there is to do is read.
|
||||
if src.len() < n {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(src.split_to(n)))
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for LengthDelimitedCodec {
|
||||
type Item = BytesMut;
|
||||
type Error = io::Error;
|
||||
|
||||
fn decode(&mut self, src: &mut BytesMut) -> io::Result<Option<BytesMut>> {
|
||||
let n = match self.state {
|
||||
DecodeState::Head => match self.decode_head(src)? {
|
||||
Some(n) => {
|
||||
self.state = DecodeState::Data(n);
|
||||
n
|
||||
}
|
||||
None => return Ok(None),
|
||||
},
|
||||
DecodeState::Data(n) => n,
|
||||
};
|
||||
|
||||
match self.decode_data(n, src)? {
|
||||
Some(data) => {
|
||||
// Update the decode state
|
||||
self.state = DecodeState::Head;
|
||||
|
||||
// Make sure the buffer has enough space to read the next head
|
||||
src.reserve(self.builder.num_head_bytes());
|
||||
|
||||
Ok(Some(data))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for LengthDelimitedCodec {
|
||||
type Item = Bytes;
|
||||
type Error = io::Error;
|
||||
|
||||
fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> Result<(), io::Error> {
|
||||
let n = (&data).into_buf().remaining();
|
||||
|
||||
if n > self.builder.max_frame_len {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
LengthDelimitedCodecError { _priv: () },
|
||||
));
|
||||
}
|
||||
|
||||
// Adjust `n` with bounds checking
|
||||
let n = if self.builder.length_adjustment < 0 {
|
||||
n.checked_add(-self.builder.length_adjustment as usize)
|
||||
} else {
|
||||
n.checked_sub(self.builder.length_adjustment as usize)
|
||||
};
|
||||
|
||||
let n = n.ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"provided length would overflow after adjustment",
|
||||
)
|
||||
})?;
|
||||
|
||||
// Reserve capacity in the destination buffer to fit the frame and
|
||||
// length field (plus adjustment).
|
||||
dst.reserve(self.builder.length_field_len + n);
|
||||
|
||||
if self.builder.length_field_is_big_endian {
|
||||
dst.put_uint_be(n as u64, self.builder.length_field_len);
|
||||
} else {
|
||||
dst.put_uint_le(n as u64, self.builder.length_field_len);
|
||||
}
|
||||
|
||||
// Write the frame to the buffer
|
||||
dst.extend_from_slice(&data[..]);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LengthDelimitedCodec {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl Builder =====
|
||||
|
||||
impl Builder {
|
||||
/// Creates a new length delimited codec builder with default configuration
|
||||
/// values.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_offset(0)
|
||||
/// .length_field_length(2)
|
||||
/// .length_adjustment(0)
|
||||
/// .num_skip(0)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn new() -> Builder {
|
||||
Builder {
|
||||
// Default max frame length of 8MB
|
||||
max_frame_len: 8 * 1_024 * 1_024,
|
||||
|
||||
// Default byte length of 4
|
||||
length_field_len: 4,
|
||||
|
||||
// Default to the header field being at the start of the header.
|
||||
length_field_offset: 0,
|
||||
|
||||
length_adjustment: 0,
|
||||
|
||||
// Total number of bytes to skip before reading the payload, if not set,
|
||||
// `length_field_len + length_field_offset`
|
||||
num_skip: None,
|
||||
|
||||
// Default to reading the length field in network (big) endian.
|
||||
length_field_is_big_endian: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the length field as a big endian integer
|
||||
///
|
||||
/// This is the default setting.
|
||||
///
|
||||
/// This configuration option applies to both encoding and decoding.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .big_endian()
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn big_endian(&mut self) -> &mut Self {
|
||||
self.length_field_is_big_endian = true;
|
||||
self
|
||||
}
|
||||
|
||||
/// Read the length field as a little endian integer
|
||||
///
|
||||
/// The default setting is big endian.
|
||||
///
|
||||
/// This configuration option applies to both encoding and decoding.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .little_endian()
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn little_endian(&mut self) -> &mut Self {
|
||||
self.length_field_is_big_endian = false;
|
||||
self
|
||||
}
|
||||
|
||||
/// Read the length field as a native endian integer
|
||||
///
|
||||
/// The default setting is big endian.
|
||||
///
|
||||
/// This configuration option applies to both encoding and decoding.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .native_endian()
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn native_endian(&mut self) -> &mut Self {
|
||||
if cfg!(target_endian = "big") {
|
||||
self.big_endian()
|
||||
} else {
|
||||
self.little_endian()
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets the max frame length
|
||||
///
|
||||
/// This configuration option applies to both encoding and decoding. The
|
||||
/// default value is 8MB.
|
||||
///
|
||||
/// When decoding, the length field read from the byte stream is checked
|
||||
/// against this setting **before** any adjustments are applied. When
|
||||
/// encoding, the length of the submitted payload is checked against this
|
||||
/// setting.
|
||||
///
|
||||
/// When frames exceed the max length, an `io::Error` with the custom value
|
||||
/// of the `LengthDelimitedCodecError` type will be returned.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .max_frame_length(8 * 1024)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn max_frame_length(&mut self, val: usize) -> &mut Self {
|
||||
self.max_frame_len = val;
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets the number of bytes used to represent the length field
|
||||
///
|
||||
/// The default value is `4`. The max value is `8`.
|
||||
///
|
||||
/// This configuration option applies to both encoding and decoding.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_length(4)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn length_field_length(&mut self, val: usize) -> &mut Self {
|
||||
assert!(val > 0 && val <= 8, "invalid length field length");
|
||||
self.length_field_len = val;
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets the number of bytes in the header before the length field
|
||||
///
|
||||
/// This configuration option only applies to decoding.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_offset(1)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn length_field_offset(&mut self, val: usize) -> &mut Self {
|
||||
self.length_field_offset = val;
|
||||
self
|
||||
}
|
||||
|
||||
/// Delta between the payload length specified in the header and the real
|
||||
/// payload length
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_adjustment(-2)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn length_adjustment(&mut self, val: isize) -> &mut Self {
|
||||
self.length_adjustment = val;
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets the number of bytes to skip before reading the payload
|
||||
///
|
||||
/// Default value is `length_field_len + length_field_offset`
|
||||
///
|
||||
/// This configuration option only applies to decoding
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .num_skip(4)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn num_skip(&mut self, val: usize) -> &mut Self {
|
||||
self.num_skip = Some(val);
|
||||
self
|
||||
}
|
||||
|
||||
/// Create a configured length delimited `LengthDelimitedCodec`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
/// # pub fn main() {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_offset(0)
|
||||
/// .length_field_length(2)
|
||||
/// .length_adjustment(0)
|
||||
/// .num_skip(0)
|
||||
/// .new_codec();
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn new_codec(&self) -> LengthDelimitedCodec {
|
||||
LengthDelimitedCodec {
|
||||
builder: *self,
|
||||
state: DecodeState::Head,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a configured length delimited `FramedRead`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncRead;
|
||||
/// use tokio_util::codec::LengthDelimitedCodec;
|
||||
///
|
||||
/// # fn bind_read<T: AsyncRead>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_offset(0)
|
||||
/// .length_field_length(2)
|
||||
/// .length_adjustment(0)
|
||||
/// .num_skip(0)
|
||||
/// .new_read(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn new_read<T>(&self, upstream: T) -> FramedRead<T, LengthDelimitedCodec>
|
||||
where
|
||||
T: AsyncRead,
|
||||
{
|
||||
FramedRead::new(upstream, self.new_codec())
|
||||
}
|
||||
|
||||
/// Create a configured length delimited `FramedWrite`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::AsyncWrite;
|
||||
/// # use tokio_util::codec::LengthDelimitedCodec;
|
||||
/// # fn write_frame<T: AsyncWrite>(io: T) {
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_length(2)
|
||||
/// .new_write(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn new_write<T>(&self, inner: T) -> FramedWrite<T, LengthDelimitedCodec>
|
||||
where
|
||||
T: AsyncWrite,
|
||||
{
|
||||
FramedWrite::new(inner, self.new_codec())
|
||||
}
|
||||
|
||||
/// Create a configured length delimited `Framed`
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # use tokio_io::{AsyncRead, AsyncWrite};
|
||||
/// # use tokio_util::codec::LengthDelimitedCodec;
|
||||
/// # fn write_frame<T: AsyncRead + AsyncWrite>(io: T) {
|
||||
/// # let _ =
|
||||
/// LengthDelimitedCodec::builder()
|
||||
/// .length_field_length(2)
|
||||
/// .new_framed(io);
|
||||
/// # }
|
||||
/// # pub fn main() {}
|
||||
/// ```
|
||||
pub fn new_framed<T>(&self, inner: T) -> Framed<T, LengthDelimitedCodec>
|
||||
where
|
||||
T: AsyncRead + AsyncWrite,
|
||||
{
|
||||
Framed::new(inner, self.new_codec())
|
||||
}
|
||||
|
||||
fn num_head_bytes(&self) -> usize {
|
||||
let num = self.length_field_offset + self.length_field_len;
|
||||
cmp::max(num, self.num_skip.unwrap_or(0))
|
||||
}
|
||||
|
||||
fn get_num_skip(&self) -> usize {
|
||||
self.num_skip
|
||||
.unwrap_or(self.length_field_offset + self.length_field_len)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Builder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl LengthDelimitedCodecError =====
|
||||
|
||||
impl fmt::Debug for LengthDelimitedCodecError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("LengthDelimitedCodecError").finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for LengthDelimitedCodecError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str("frame size too big")
|
||||
}
|
||||
}
|
||||
|
||||
impl StdError for LengthDelimitedCodecError {}
|
||||
@@ -0,0 +1,224 @@
|
||||
use crate::codec::decoder::Decoder;
|
||||
use crate::codec::encoder::Encoder;
|
||||
|
||||
use bytes::{BufMut, BytesMut};
|
||||
use std::{cmp, fmt, io, str, usize};
|
||||
|
||||
/// A simple `Codec` implementation that splits up data into lines.
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
|
||||
pub struct LinesCodec {
|
||||
// Stored index of the next index to examine for a `\n` character.
|
||||
// This is used to optimize searching.
|
||||
// For example, if `decode` was called with `abc`, it would hold `3`,
|
||||
// because that is the next index to examine.
|
||||
// The next time `decode` is called with `abcde\n`, the method will
|
||||
// only look at `de\n` before returning.
|
||||
next_index: usize,
|
||||
|
||||
/// The maximum length for a given line. If `usize::MAX`, lines will be
|
||||
/// read until a `\n` character is reached.
|
||||
max_length: usize,
|
||||
|
||||
/// Are we currently discarding the remainder of a line which was over
|
||||
/// the length limit?
|
||||
is_discarding: bool,
|
||||
}
|
||||
|
||||
impl LinesCodec {
|
||||
/// Returns a `LinesCodec` for splitting up data into lines.
|
||||
///
|
||||
/// # Note
|
||||
///
|
||||
/// The returned `LinesCodec` will not have an upper bound on the length
|
||||
/// of a buffered line. See the documentation for [`new_with_max_length`]
|
||||
/// for information on why this could be a potential security risk.
|
||||
///
|
||||
/// [`new_with_max_length`]: #method.new_with_max_length
|
||||
pub fn new() -> LinesCodec {
|
||||
LinesCodec {
|
||||
next_index: 0,
|
||||
max_length: usize::MAX,
|
||||
is_discarding: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a `LinesCodec` with a maximum line length limit.
|
||||
///
|
||||
/// If this is set, calls to `LinesCodec::decode` will return a
|
||||
/// [`LengthError`] when a line exceeds the length limit. Subsequent calls
|
||||
/// will discard up to `limit` bytes from that line until a newline
|
||||
/// character is reached, returning `None` until the line over the limit
|
||||
/// has been fully discarded. After that point, calls to `decode` will
|
||||
/// function as normal.
|
||||
///
|
||||
/// # Note
|
||||
///
|
||||
/// Setting a length limit is highly recommended for any `LinesCodec` which
|
||||
/// will be exposed to untrusted input. Otherwise, the size of the buffer
|
||||
/// that holds the line currently being read is unbounded. An attacker could
|
||||
/// exploit this unbounded buffer by sending an unbounded amount of input
|
||||
/// without any `\n` characters, causing unbounded memory consumption.
|
||||
///
|
||||
/// [`LengthError`]: ../struct.LengthError
|
||||
pub fn new_with_max_length(max_length: usize) -> Self {
|
||||
LinesCodec {
|
||||
max_length,
|
||||
..LinesCodec::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the maximum line length when decoding.
|
||||
///
|
||||
/// ```
|
||||
/// use std::usize;
|
||||
/// use tokio_util::codec::LinesCodec;
|
||||
///
|
||||
/// let codec = LinesCodec::new();
|
||||
/// assert_eq!(codec.max_length(), usize::MAX);
|
||||
/// ```
|
||||
/// ```
|
||||
/// use tokio_util::codec::LinesCodec;
|
||||
///
|
||||
/// let codec = LinesCodec::new_with_max_length(256);
|
||||
/// assert_eq!(codec.max_length(), 256);
|
||||
/// ```
|
||||
pub fn max_length(&self) -> usize {
|
||||
self.max_length
|
||||
}
|
||||
}
|
||||
|
||||
fn utf8(buf: &[u8]) -> Result<&str, io::Error> {
|
||||
str::from_utf8(buf)
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Unable to decode input as UTF8"))
|
||||
}
|
||||
|
||||
fn without_carriage_return(s: &[u8]) -> &[u8] {
|
||||
if let Some(&b'\r') = s.last() {
|
||||
&s[..s.len() - 1]
|
||||
} else {
|
||||
s
|
||||
}
|
||||
}
|
||||
|
||||
impl Decoder for LinesCodec {
|
||||
type Item = String;
|
||||
type Error = LinesCodecError;
|
||||
|
||||
fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<String>, LinesCodecError> {
|
||||
loop {
|
||||
// Determine how far into the buffer we'll search for a newline. If
|
||||
// there's no max_length set, we'll read to the end of the buffer.
|
||||
let read_to = cmp::min(self.max_length.saturating_add(1), buf.len());
|
||||
|
||||
let newline_offset = buf[self.next_index..read_to]
|
||||
.iter()
|
||||
.position(|b| *b == b'\n');
|
||||
|
||||
match (self.is_discarding, newline_offset) {
|
||||
(true, Some(offset)) => {
|
||||
// If we found a newline, discard up to that offset and
|
||||
// then stop discarding. On the next iteration, we'll try
|
||||
// to read a line normally.
|
||||
buf.advance(offset + self.next_index + 1);
|
||||
self.is_discarding = false;
|
||||
self.next_index = 0;
|
||||
}
|
||||
(true, None) => {
|
||||
// Otherwise, we didn't find a newline, so we'll discard
|
||||
// everything we read. On the next iteration, we'll continue
|
||||
// discarding up to max_len bytes unless we find a newline.
|
||||
buf.advance(read_to);
|
||||
self.next_index = 0;
|
||||
if buf.is_empty() {
|
||||
return Err(LinesCodecError::MaxLineLengthExceeded);
|
||||
}
|
||||
}
|
||||
(false, Some(offset)) => {
|
||||
// Found a line!
|
||||
let newline_index = offset + self.next_index;
|
||||
self.next_index = 0;
|
||||
let line = buf.split_to(newline_index + 1);
|
||||
let line = &line[..line.len() - 1];
|
||||
let line = without_carriage_return(line);
|
||||
let line = utf8(line)?;
|
||||
return Ok(Some(line.to_string()));
|
||||
}
|
||||
(false, None) if buf.len() > self.max_length => {
|
||||
// Reached the maximum length without finding a
|
||||
// newline, return an error and start discarding on the
|
||||
// next call.
|
||||
self.is_discarding = true;
|
||||
return Err(LinesCodecError::MaxLineLengthExceeded);
|
||||
}
|
||||
(false, None) => {
|
||||
// We didn't find a line or reach the length limit, so the next
|
||||
// call will resume searching at the current offset.
|
||||
self.next_index = read_to;
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<String>, LinesCodecError> {
|
||||
Ok(match self.decode(buf)? {
|
||||
Some(frame) => Some(frame),
|
||||
None => {
|
||||
// No terminating newline - return remaining data, if any
|
||||
if buf.is_empty() || buf == &b"\r"[..] {
|
||||
None
|
||||
} else {
|
||||
let line = buf.take();
|
||||
let line = without_carriage_return(&line);
|
||||
let line = utf8(line)?;
|
||||
self.next_index = 0;
|
||||
Some(line.to_string())
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Encoder for LinesCodec {
|
||||
type Item = String;
|
||||
type Error = LinesCodecError;
|
||||
|
||||
fn encode(&mut self, line: String, buf: &mut BytesMut) -> Result<(), LinesCodecError> {
|
||||
buf.reserve(line.len() + 1);
|
||||
buf.put(line);
|
||||
buf.put_u8(b'\n');
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LinesCodec {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// An error occured while encoding or decoding a line.
|
||||
#[derive(Debug)]
|
||||
pub enum LinesCodecError {
|
||||
/// The maximum line length was exceeded.
|
||||
MaxLineLengthExceeded,
|
||||
/// An IO error occured.
|
||||
Io(io::Error),
|
||||
}
|
||||
|
||||
impl fmt::Display for LinesCodecError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
LinesCodecError::MaxLineLengthExceeded => write!(f, "max line length exceeded"),
|
||||
LinesCodecError::Io(e) => write!(f, "{}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<io::Error> for LinesCodecError {
|
||||
fn from(e: io::Error) -> LinesCodecError {
|
||||
LinesCodecError::Io(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for LinesCodecError {}
|
||||
@@ -0,0 +1,7 @@
|
||||
/// A macro to reduce some of the boilerplate for projecting from
|
||||
/// `Pin<&mut T>` to `Pin<&mut T.field>`
|
||||
macro_rules! pin {
|
||||
($e:expr) => {
|
||||
std::pin::Pin::new(&mut $e)
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
//! Utilities for encoding and decoding frames.
|
||||
//!
|
||||
//! Contains adapters to go from streams of bytes, [`AsyncRead`] and
|
||||
//! [`AsyncWrite`], to framed streams implementing [`Sink`] and [`Stream`].
|
||||
//! Framed streams are also known as transports.
|
||||
//!
|
||||
//! [`AsyncRead`]: https://docs.rs/tokio/*/tokio/io/trait.AsyncRead.html
|
||||
//! [`AsyncWrite`]: https://docs.rs/tokio/*/tokio/io/trait.AsyncWrite.html
|
||||
//! [`Sink`]: https://docs.rs/futures-sink-preview/*/futures_sink/trait.Sink.html
|
||||
//! [`Stream`]: https://docs.rs/futures-core-preview/*/futures_core/stream/trait.Stream.html
|
||||
|
||||
#[macro_use]
|
||||
mod macros;
|
||||
|
||||
mod bytes_codec;
|
||||
pub use self::bytes_codec::BytesCodec;
|
||||
|
||||
mod decoder;
|
||||
pub use self::decoder::Decoder;
|
||||
|
||||
mod encoder;
|
||||
pub use self::encoder::Encoder;
|
||||
|
||||
mod framed;
|
||||
pub use self::framed::{Framed, FramedParts};
|
||||
|
||||
mod framed_read;
|
||||
pub use self::framed_read::FramedRead;
|
||||
|
||||
mod framed_write;
|
||||
pub use self::framed_write::FramedWrite;
|
||||
|
||||
pub mod length_delimited;
|
||||
pub use self::length_delimited::{LengthDelimitedCodec, LengthDelimitedCodecError};
|
||||
|
||||
mod lines_codec;
|
||||
pub use self::lines_codec::{LinesCodec, LinesCodecError};
|
||||
Reference in New Issue
Block a user