libp2p_webrtc/tokio/
stream.rs

1// Copyright 2023 Protocol Labs.
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 std::{
22    pin::Pin,
23    sync::Arc,
24    task::{Context, Poll},
25};
26
27use futures::prelude::*;
28use libp2p_webrtc_utils::MAX_MSG_LEN;
29use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
30use webrtc::data::data_channel::{DataChannel, PollDataChannel};
31
32/// A substream on top of a WebRTC data channel.
33///
34/// To be a proper libp2p substream, we need to implement [`AsyncRead`] and [`AsyncWrite`] as well
35/// as support a half-closed state which we do by framing messages in a protobuf envelope.
36pub struct Stream {
37    inner: libp2p_webrtc_utils::Stream<Compat<PollDataChannel>>,
38}
39
40pub(crate) type DropListener = libp2p_webrtc_utils::DropListener<Compat<PollDataChannel>>;
41
42impl Stream {
43    /// Returns a new `Substream` and a listener, which will notify the receiver when/if the
44    /// substream is dropped.
45    pub(crate) fn new(data_channel: Arc<DataChannel>) -> (Self, DropListener) {
46        let mut data_channel = PollDataChannel::new(data_channel).compat();
47        data_channel.get_mut().set_read_buf_capacity(MAX_MSG_LEN);
48
49        let (inner, drop_listener) = libp2p_webrtc_utils::Stream::new(data_channel);
50
51        (Self { inner }, drop_listener)
52    }
53}
54impl AsyncRead for Stream {
55    fn poll_read(
56        self: Pin<&mut Self>,
57        cx: &mut Context<'_>,
58        buf: &mut [u8],
59    ) -> Poll<std::io::Result<usize>> {
60        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
61    }
62}
63
64impl AsyncWrite for Stream {
65    fn poll_write(
66        self: Pin<&mut Self>,
67        cx: &mut Context<'_>,
68        buf: &[u8],
69    ) -> Poll<std::io::Result<usize>> {
70        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
71    }
72
73    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
74        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
75    }
76
77    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
78        Pin::new(&mut self.get_mut().inner).poll_close(cx)
79    }
80}