| 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}, |
| stream::Stream, |
| }; |
| |
| use bytes::BytesMut; |
| use futures_sink::Sink; |
| use pin_project_lite::pin_project; |
| use std::fmt; |
| use std::io::{self, BufRead, Read, Write}; |
| use std::mem::MaybeUninit; |
| use std::pin::Pin; |
| use std::task::{Context, Poll}; |
| |
| pin_project! { |
| /// 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 [`Decoder::framed`] adapter, or |
| /// by using the `new` function seen below. |
| /// |
| /// [`Stream`]: tokio::stream::Stream |
| /// [`Sink`]: futures_sink::Sink |
| /// [`AsyncRead`]: tokio::io::AsyncRead |
| /// [`Decoder::framed`]: crate::codec::Decoder::framed() |
| pub struct Framed<T, U> { |
| #[pin] |
| inner: FramedRead2<FramedWrite2<Fuse<T, U>>>, |
| } |
| } |
| |
| pin_project! { |
| pub(crate) struct Fuse<T, U> { |
| #[pin] |
| pub(crate) io: T, |
| pub(crate) codec: U, |
| } |
| } |
| |
| /// Abstracts over `FramedRead2` being either `FramedRead2<FramedWrite2<Fuse<T, U>>>` or |
| /// `FramedRead2<Fuse<T, U>>` and lets the io and codec parts be extracted in either case. |
| pub(crate) trait ProjectFuse { |
| type Io; |
| type Codec; |
| |
| fn project(self: Pin<&mut Self>) -> Fuse<Pin<&mut Self::Io>, &mut Self::Codec>; |
| } |
| |
| impl<T, U> ProjectFuse for Fuse<T, U> { |
| type Io = T; |
| type Codec = U; |
| |
| fn project(self: Pin<&mut Self>) -> Fuse<Pin<&mut Self::Io>, &mut Self::Codec> { |
| let self_ = self.project(); |
| Fuse { |
| io: self_.io, |
| codec: self_.codec, |
| } |
| } |
| } |
| |
| impl<T, U> Framed<T, U> |
| where |
| T: AsyncRead + AsyncWrite, |
| { |
| /// Provides a [`Stream`] and [`Sink`] interface for reading and writing to this |
| /// I/O object, using [`Decoder`] and [`Encoder`] 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. |
| /// |
| /// [`Stream`]: tokio::stream::Stream |
| /// [`Sink`]: futures_sink::Sink |
| /// [`Decode`]: crate::codec::Decoder |
| /// [`Encoder`]: crate::codec::Encoder |
| /// [`split`]: https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.split |
| pub fn new(inner: T, codec: U) -> Framed<T, U> { |
| Framed { |
| inner: framed_read2(framed_write2(Fuse { io: inner, codec })), |
| } |
| } |
| |
| /// Provides a [`Stream`] and [`Sink`] interface for reading and writing to this |
| /// I/O object, using [`Decoder`] and [`Encoder`] to read and write the raw data, |
| /// with a specific read buffer initial capacity. |
| /// |
| /// 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. |
| /// |
| /// [`Stream`]: tokio::stream::Stream |
| /// [`Sink`]: futures_sink::Sink |
| /// [`Decode`]: crate::codec::Decoder |
| /// [`Encoder`]: crate::codec::Encoder |
| /// [`split`]: https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.split |
| pub fn with_capacity(inner: T, codec: U, capacity: usize) -> Framed<T, U> { |
| Framed { |
| inner: framed_read2_with_buffer( |
| framed_write2(Fuse { io: inner, codec }), |
| BytesMut::with_capacity(capacity), |
| ), |
| } |
| } |
| } |
| |
| impl<T, U> Framed<T, U> { |
| /// Provides a [`Stream`] and [`Sink`] interface for reading and writing to this |
| /// I/O object, using [`Decoder`] and [`Encoder`] 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. |
| /// |
| /// [`Stream`]: tokio::stream::Stream |
| /// [`Sink`]: futures_sink::Sink |
| /// [`Decoder`]: crate::codec::Decoder |
| /// [`Encoder`]: crate::codec::Encoder |
| /// [`into_parts`]: crate::codec::Framed::into_parts() |
| /// [`split`]: https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.split |
| pub fn from_parts(parts: FramedParts<T, U>) -> Framed<T, U> { |
| Framed { |
| inner: framed_read2_with_buffer( |
| framed_write2_with_buffer( |
| Fuse { |
| io: parts.io, |
| codec: parts.codec, |
| }, |
| parts.write_buf, |
| ), |
| parts.read_buf, |
| ), |
| } |
| } |
| |
| /// 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.inner.get_ref().get_ref().io |
| } |
| |
| /// 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.inner.get_mut().get_mut().io |
| } |
| |
| /// Returns a reference to the underlying codec wrapped by |
| /// `Framed`. |
| /// |
| /// 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().codec |
| } |
| |
| /// Returns a mutable reference to the underlying codec wrapped by |
| /// `Framed`. |
| /// |
| /// 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().codec |
| } |
| |
| /// Returns a reference to the read buffer. |
| pub fn read_buffer(&self) -> &BytesMut { |
| self.inner.buffer() |
| } |
| |
| /// 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.inner.into_inner().into_inner().io |
| } |
| |
| /// Consumes the `Framed`, 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.io, |
| codec: inner.codec, |
| read_buf, |
| write_buf, |
| _priv: (), |
| } |
| } |
| } |
| |
| impl<T, U> Stream for Framed<T, U> |
| where |
| T: AsyncRead, |
| U: Decoder, |
| { |
| type Item = Result<U::Item, U::Error>; |
| |
| fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { |
| self.project().inner.poll_next(cx) |
| } |
| } |
| |
| impl<T, I, U> Sink<I> for Framed<T, U> |
| where |
| T: AsyncWrite, |
| U: Encoder<I>, |
| U::Error: From<io::Error>, |
| { |
| type Error = U::Error; |
| |
| fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
| self.project().inner.get_pin_mut().poll_ready(cx) |
| } |
| |
| fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> { |
| self.project().inner.get_pin_mut().start_send(item) |
| } |
| |
| fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
| self.project().inner.get_pin_mut().poll_flush(cx) |
| } |
| |
| fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { |
| self.project().inner.get_pin_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().io) |
| .field("codec", &self.inner.get_ref().get_ref().codec) |
| .finish() |
| } |
| } |
| |
| // ===== impl Fuse ===== |
| |
| impl<T: Read, U> Read for Fuse<T, U> { |
| fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> { |
| self.io.read(dst) |
| } |
| } |
| |
| impl<T: BufRead, U> BufRead for Fuse<T, U> { |
| fn fill_buf(&mut self) -> io::Result<&[u8]> { |
| self.io.fill_buf() |
| } |
| |
| fn consume(&mut self, amt: usize) { |
| self.io.consume(amt) |
| } |
| } |
| |
| impl<T: AsyncRead, U> AsyncRead for Fuse<T, U> { |
| unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [MaybeUninit<u8>]) -> bool { |
| self.io.prepare_uninitialized_buffer(buf) |
| } |
| |
| fn poll_read( |
| self: Pin<&mut Self>, |
| cx: &mut Context<'_>, |
| buf: &mut [u8], |
| ) -> Poll<Result<usize, io::Error>> { |
| self.project().io.poll_read(cx, buf) |
| } |
| } |
| |
| impl<T: AsyncBufRead, U> AsyncBufRead for Fuse<T, U> { |
| fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { |
| self.project().io.poll_fill_buf(cx) |
| } |
| |
| fn consume(self: Pin<&mut Self>, amt: usize) { |
| self.project().io.consume(amt) |
| } |
| } |
| |
| impl<T: Write, U> Write for Fuse<T, U> { |
| fn write(&mut self, src: &[u8]) -> io::Result<usize> { |
| self.io.write(src) |
| } |
| |
| fn flush(&mut self) -> io::Result<()> { |
| self.io.flush() |
| } |
| } |
| |
| impl<T: AsyncWrite, U> AsyncWrite for Fuse<T, U> { |
| fn poll_write( |
| self: Pin<&mut Self>, |
| cx: &mut Context<'_>, |
| buf: &[u8], |
| ) -> Poll<Result<usize, io::Error>> { |
| self.project().io.poll_write(cx, buf) |
| } |
| |
| fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { |
| self.project().io.poll_flush(cx) |
| } |
| |
| fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> { |
| self.project().io.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.codec.decode(buffer) |
| } |
| |
| fn decode_eof(&mut self, buffer: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> { |
| self.codec.decode_eof(buffer) |
| } |
| } |
| |
| impl<T, I, U: Encoder<I>> Encoder<I> for Fuse<T, U> { |
| type Error = U::Error; |
| |
| fn encode(&mut self, item: I, dst: &mut BytesMut) -> Result<(), Self::Error> { |
| self.codec.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. |
| /// |
| /// [`Framed`]: crate::codec::Framed |
| #[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<I>(io: T, codec: U) -> FramedParts<T, U> |
| where |
| U: Encoder<I>, |
| { |
| FramedParts { |
| io, |
| codec, |
| read_buf: BytesMut::new(), |
| write_buf: BytesMut::new(), |
| _priv: (), |
| } |
| } |
| } |