      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200

    mod error;
    mod framed_read;
    mod framed_write;

    pub use self::error::{RecvError, SendError, UserError};

    use self::framed_read::FramedRead;
    use self::framed_write::FramedWrite;

    use crate::frame::{self, Data, Frame};

    use bytes::Buf;
    use futures_core::Stream;
    use futures_sink::Sink;
    use std::pin::Pin;
    use std::task::{Context, Poll};
    use tokio::io::{AsyncRead, AsyncWrite};
    use tokio_util::codec::length_delimited;

    use std::io;

    #[derive(Debug)]
    pub struct Codec<T, B> {
        inner: FramedRead<FramedWrite<T, B>>,
    }

    impl<T, B> Codec<T, B>
    where
        T: AsyncRead + AsyncWrite + Unpin,
        B: Buf,
    {
        /// Returns a new `Codec` with the default max frame size
        #[inline]
        pub fn new(io: T) -> Self {
            Self::with_max_recv_frame_size(io, frame::DEFAULT_MAX_FRAME_SIZE as usize)
        }

        /// Returns a new `Codec` with the given maximum frame size
        pub fn with_max_recv_frame_size(io: T, max_frame_size: usize) -> Self {
            // Wrap with writer
            let framed_write = FramedWrite::new(io);

            // Delimit the frames
            let delimited = length_delimited::Builder::new()
                .big_endian()
                .length_field_length(3)
                .length_adjustment(9)
                .num_skip(0) // Don't skip the header
                .new_read(framed_write);

            let mut inner = FramedRead::new(delimited);

            // Use FramedRead's method since it checks the value is within range.
            inner.set_max_frame_size(max_frame_size);

            Codec { inner }
        }
    }

    impl<T, B> Codec<T, B> {
        /// Updates the max received frame size.
        ///
        /// 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 size in effect
        /// before calling this function, then the frame will be allowed.
        #[inline]
        pub fn set_max_recv_frame_size(&mut self, val: usize) {
            self.inner.set_max_frame_size(val)
        }

        /// Returns the current max received frame size setting.
        ///
        /// This is the largest size this codec will accept from the wire. Larger
        /// frames will be rejected.
        #[cfg(feature = "unstable")]
        #[inline]
        pub fn max_recv_frame_size(&self) -> usize {
            self.inner.max_frame_size()
        }

        /// Returns the max frame size that can be sent to the peer.
        pub fn max_send_frame_size(&self) -> usize {
            self.inner.get_ref().max_frame_size()
        }

        /// Set the peer's max frame size.
        pub fn set_max_send_frame_size(&mut self, val: usize) {
            self.framed_write().set_max_frame_size(val)
        }

        /// Set the peer's header table size size.
        pub fn set_send_header_table_size(&mut self, val: usize) {
            self.framed_write().set_header_table_size(val)
        }

        /// Set the max header list size that can be received.
        pub fn set_max_recv_header_list_size(&mut self, val: usize) {
            self.inner.set_max_header_list_size(val);
        }

        /// Get a reference to the inner stream.
        #[cfg(feature = "unstable")]
        pub fn get_ref(&self) -> &T {
            self.inner.get_ref().get_ref()
        }

        /// Get a mutable reference to the inner stream.
        pub fn get_mut(&mut self) -> &mut T {
            self.inner.get_mut().get_mut()
        }

        /// Takes the data payload value that was fully written to the socket
        pub(crate) fn take_last_data_frame(&mut self) -> Option<Data<B>> {
            self.framed_write().take_last_data_frame()
        }

        fn framed_write(&mut self) -> &mut FramedWrite<T, B> {
            self.inner.get_mut()
        }
    }

    impl<T, B> Codec<T, B>
    where
        T: AsyncWrite + Unpin,
        B: Buf,
    {
        /// Returns `Ready` when the codec can buffer a frame
        pub fn poll_ready(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
            self.framed_write().poll_ready(cx)
        }

        /// Buffer a frame.
        ///
        /// `poll_ready` must be called first to ensure that a frame may be
        /// accepted.
        ///
        /// TODO: Rename this to avoid conflicts with Sink::buffer
        pub fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
            self.framed_write().buffer(item)
        }

        /// Flush buffered data to the wire
        pub fn flush(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
            self.framed_write().flush(cx)
        }

        /// Shutdown the send half
        pub fn shutdown(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
            self.framed_write().shutdown(cx)
        }
    }

    impl<T, B> Stream for Codec<T, B>
    where
        T: AsyncRead + Unpin,
    {
        type Item = Result<Frame, RecvError>;

        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Pin::new(&mut self.inner).poll_next(cx)
        }
    }

    impl<T, B> Sink<Frame<B>> for Codec<T, B>
    where
        T: AsyncWrite + Unpin,
        B: Buf,
    {
        type Error = SendError;

        fn start_send(mut self: Pin<&mut Self>, item: Frame<B>) -> Result<(), Self::Error> {
            Codec::buffer(&mut self, item)?;
            Ok(())
        }
        /// Returns `Ready` when the codec can buffer a frame
        fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.framed_write().poll_ready(cx).map_err(Into::into)
        }

        /// Flush buffered data to the wire
        fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.framed_write().flush(cx).map_err(Into::into)
        }

        fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            ready!(self.shutdown(cx))?;
            Poll::Ready(Ok(()))
        }
    }

    // TODO: remove (or improve) this
    impl<T> From<T> for Codec<T, bytes::Bytes>
    where
        T: AsyncRead + AsyncWrite + Unpin,
    {
        fn from(src: T) -> Self {
            Self::new(src)
        }
    }