litep2p/transport/websocket/
substream.rs

1// Copyright 2023 litep2p developers
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21use crate::{protocol::Permit, BandwidthSink};
22
23use tokio::io::{AsyncRead, AsyncWrite};
24use tokio_util::compat::Compat;
25
26use std::{
27    io,
28    pin::Pin,
29    task::{Context, Poll},
30};
31
32/// Substream that holds the inner substream provided by the transport
33/// and a permit which keeps the connection open.
34#[derive(Debug)]
35pub struct Substream {
36    /// Underlying socket.
37    io: Compat<crate::yamux::Stream>,
38
39    /// Bandwidth sink.
40    bandwidth_sink: BandwidthSink,
41
42    /// Connection permit.
43    _permit: Permit,
44}
45
46impl Substream {
47    /// Create new [`Substream`].
48    pub fn new(
49        io: Compat<crate::yamux::Stream>,
50        bandwidth_sink: BandwidthSink,
51        _permit: Permit,
52    ) -> Self {
53        Self {
54            io,
55            bandwidth_sink,
56            _permit,
57        }
58    }
59}
60
61impl AsyncRead for Substream {
62    fn poll_read(
63        mut self: Pin<&mut Self>,
64        cx: &mut Context<'_>,
65        buf: &mut tokio::io::ReadBuf<'_>,
66    ) -> Poll<io::Result<()>> {
67        match futures::ready!(Pin::new(&mut self.io).poll_read(cx, buf)) {
68            Err(error) => Poll::Ready(Err(error)),
69            Ok(res) => {
70                self.bandwidth_sink.increase_inbound(buf.filled().len());
71                Poll::Ready(Ok(res))
72            }
73        }
74    }
75}
76
77impl AsyncWrite for Substream {
78    fn poll_write(
79        mut self: Pin<&mut Self>,
80        cx: &mut Context<'_>,
81        buf: &[u8],
82    ) -> Poll<Result<usize, io::Error>> {
83        match futures::ready!(Pin::new(&mut self.io).poll_write(cx, buf)) {
84            Err(error) => Poll::Ready(Err(error)),
85            Ok(nwritten) => {
86                self.bandwidth_sink.increase_outbound(nwritten);
87                Poll::Ready(Ok(nwritten))
88            }
89        }
90    }
91
92    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
93        Pin::new(&mut self.io).poll_flush(cx)
94    }
95
96    fn poll_shutdown(
97        mut self: Pin<&mut Self>,
98        cx: &mut Context<'_>,
99    ) -> Poll<Result<(), io::Error>> {
100        Pin::new(&mut self.io).poll_shutdown(cx)
101    }
102}
litep2p/transport/websocket/substream.rs

litep2p/transport/websocket/
substream.rs
