Skip to main content

volans_swarm/
substream.rs

1use volans_core::{Negotiated, muxing::SubstreamBox};
2use either::Either;
3use futures::{AsyncRead, AsyncWrite};
4
5use std::{
6    fmt,
7    hash::{Hash, Hasher},
8    io,
9    pin::Pin,
10    sync::Arc,
11    task::{Context, Poll},
12};
13
14#[derive(Debug, Clone)]
15pub(crate) struct ActiveStreamCounter(Arc<()>);
16
17impl ActiveStreamCounter {
18    pub(crate) fn new() -> Self {
19        Self(Arc::new(()))
20    }
21
22    pub(crate) fn no_active_streams(&self) -> bool {
23        Arc::strong_count(&self.0) == 1
24    }
25}
26
27#[derive(Debug)]
28pub struct Substream {
29    stream: Negotiated<SubstreamBox>,
30    counter: Option<ActiveStreamCounter>,
31}
32
33impl Substream {
34    pub(crate) fn new(stream: Negotiated<SubstreamBox>, counter: ActiveStreamCounter) -> Self {
35        Self {
36            stream,
37            counter: Some(counter),
38        }
39    }
40
41    pub fn ignore_for_keep_alive(&mut self) {
42        self.counter.take();
43    }
44}
45
46impl AsyncRead for Substream {
47    fn poll_read(
48        self: Pin<&mut Self>,
49        cx: &mut Context<'_>,
50        buf: &mut [u8],
51    ) -> Poll<io::Result<usize>> {
52        Pin::new(&mut self.get_mut().stream).poll_read(cx, buf)
53    }
54
55    fn poll_read_vectored(
56        self: Pin<&mut Self>,
57        cx: &mut Context<'_>,
58        bufs: &mut [io::IoSliceMut<'_>],
59    ) -> Poll<io::Result<usize>> {
60        Pin::new(&mut self.get_mut().stream).poll_read_vectored(cx, bufs)
61    }
62}
63
64impl AsyncWrite for Substream {
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().stream).poll_write(cx, buf)
71    }
72
73    fn poll_write_vectored(
74        self: Pin<&mut Self>,
75        cx: &mut Context<'_>,
76        bufs: &[io::IoSlice<'_>],
77    ) -> Poll<io::Result<usize>> {
78        Pin::new(&mut self.get_mut().stream).poll_write_vectored(cx, bufs)
79    }
80
81    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
82        Pin::new(&mut self.get_mut().stream).poll_flush(cx)
83    }
84
85    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
86        Pin::new(&mut self.get_mut().stream).poll_close(cx)
87    }
88}
89
90#[derive(Clone, Eq)]
91pub struct StreamProtocol {
92    inner: Either<&'static str, Arc<str>>,
93}
94
95impl StreamProtocol {
96    pub const fn new(s: &'static str) -> Self {
97        match s.as_bytes() {
98            [b'/', ..] => {}
99            _ => panic!("Protocols should start with a /"),
100        }
101
102        StreamProtocol {
103            inner: Either::Left(s),
104        }
105    }
106    pub fn try_from_owned(protocol: String) -> Result<Self, InvalidProtocol> {
107        if !protocol.starts_with('/') {
108            return Err(InvalidProtocol::missing_forward_slash());
109        }
110
111        Ok(StreamProtocol {
112            // FIXME: Can we somehow reuse the
113            // allocation from the owned string?
114            inner: Either::Right(Arc::from(protocol)),
115        })
116    }
117}
118
119impl AsRef<str> for StreamProtocol {
120    fn as_ref(&self) -> &str {
121        either::for_both!(&self.inner, s => s)
122    }
123}
124
125impl fmt::Debug for StreamProtocol {
126    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
127        either::for_both!(&self.inner, s => s.fmt(f))
128    }
129}
130
131impl fmt::Display for StreamProtocol {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        self.inner.fmt(f)
134    }
135}
136
137impl PartialEq<&str> for StreamProtocol {
138    fn eq(&self, other: &&str) -> bool {
139        self.as_ref() == *other
140    }
141}
142
143impl PartialEq<StreamProtocol> for &str {
144    fn eq(&self, other: &StreamProtocol) -> bool {
145        *self == other.as_ref()
146    }
147}
148
149impl PartialEq for StreamProtocol {
150    fn eq(&self, other: &Self) -> bool {
151        self.as_ref() == other.as_ref()
152    }
153}
154
155impl Hash for StreamProtocol {
156    fn hash<H: Hasher>(&self, state: &mut H) {
157        self.as_ref().hash(state)
158    }
159}
160
161#[derive(Debug)]
162pub struct InvalidProtocol {
163    _private: (),
164}
165
166impl InvalidProtocol {
167    pub(crate) fn missing_forward_slash() -> Self {
168        InvalidProtocol { _private: () }
169    }
170}