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 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 `AsyncRead::framed` adapter. pub struct Framed { #[pin] inner: FramedRead2>>, } } pin_project! { pub(crate) struct Fuse { #[pin] pub(crate) io: T, pub(crate) codec: U, } } /// Abstracts over `FramedRead2` being either `FramedRead2>>` or /// `FramedRead2>` 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, &mut Self::Codec>; } impl ProjectFuse for Fuse { type Io = T; type Codec = U; fn project(self: Pin<&mut Self>) -> Fuse, &mut Self::Codec> { let self_ = self.project(); Fuse { io: self_.io, codec: self_.codec, } } } impl Framed 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 { Framed { inner: framed_read2(framed_write2(Fuse { io: inner, codec })), } } } impl Framed { /// 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) -> Framed { 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 /// `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().io } /// 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().io } /// 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().codec } /// 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().codec } /// Returns a reference to the read buffer. pub fn read_buffer(&self) -> &BytesMut { self.inner.buffer() } /// 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().io } /// 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 { 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 Stream for Framed where T: AsyncRead, U: Decoder, { type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.poll_next(cx) } } impl Sink for Framed where T: AsyncWrite, U: Encoder, U::Error: From, { type Error = U::Error; fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 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> { self.project().inner.get_pin_mut().poll_flush(cx) } fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().inner.get_pin_mut().poll_close(cx) } } impl fmt::Debug for Framed 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 Read for Fuse { fn read(&mut self, dst: &mut [u8]) -> io::Result { self.io.read(dst) } } impl BufRead for Fuse { fn fill_buf(&mut self) -> io::Result<&[u8]> { self.io.fill_buf() } fn consume(&mut self, amt: usize) { self.io.consume(amt) } } impl AsyncRead for Fuse { unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [MaybeUninit]) -> bool { self.io.prepare_uninitialized_buffer(buf) } fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8], ) -> Poll> { self.project().io.poll_read(cx, buf) } } impl AsyncBufRead for Fuse { fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().io.poll_fill_buf(cx) } fn consume(self: Pin<&mut Self>, amt: usize) { self.project().io.consume(amt) } } impl Write for Fuse { fn write(&mut self, src: &[u8]) -> io::Result { self.io.write(src) } fn flush(&mut self) -> io::Result<()> { self.io.flush() } } impl AsyncWrite for Fuse { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.project().io.poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().io.poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project().io.poll_shutdown(cx) } } impl Decoder for Fuse { type Item = U::Item; type Error = U::Error; fn decode(&mut self, buffer: &mut BytesMut) -> Result, Self::Error> { self.codec.decode(buffer) } fn decode_eof(&mut self, buffer: &mut BytesMut) -> Result, Self::Error> { self.codec.decode_eof(buffer) } } impl Encoder for Fuse { type Item = U::Item; type Error = U::Error; fn encode(&mut self, item: Self::Item, 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. #[derive(Debug)] pub struct FramedParts { /// 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 FramedParts { /// Create a new, default, `FramedParts` pub fn new(io: T, codec: U) -> FramedParts { FramedParts { io, codec, read_buf: BytesMut::new(), write_buf: BytesMut::new(), _priv: (), } } }