compio_tls/
maybe.rs

1use std::{
2    borrow::Cow,
3    io,
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use futures_util::{AsyncRead, AsyncWrite};
9
10use crate::TlsStream;
11
12#[derive(Debug)]
13#[allow(clippy::large_enum_variant)]
14enum MaybeTlsStreamInner<S> {
15    Plain(S),
16    Tls(TlsStream<S>),
17}
18
19/// A futures-based stream that can be either plain TCP or TLS-encrypted.
20#[derive(Debug)]
21pub struct MaybeTlsStream<S>(MaybeTlsStreamInner<S>);
22
23impl<S> MaybeTlsStream<S> {
24    /// Create an unencrypted stream.
25    pub fn new_plain(stream: S) -> Self {
26        Self(MaybeTlsStreamInner::Plain(stream))
27    }
28
29    /// Create a TLS-encrypted stream.
30    pub fn new_tls(stream: TlsStream<S>) -> Self {
31        Self(MaybeTlsStreamInner::Tls(stream))
32    }
33
34    /// Whether the stream is TLS-encrypted.
35    pub fn is_tls(&self) -> bool {
36        matches!(self.0, MaybeTlsStreamInner::Tls(_))
37    }
38
39    /// Whether `poll_write_vectored` takes all the buffers at once; unknown,
40    /// so no, for a plain stream.
41    #[cfg(feature = "tokio")]
42    pub(crate) fn writes_vectored(&self) -> bool {
43        match &self.0 {
44            MaybeTlsStreamInner::Plain(_) => false,
45            MaybeTlsStreamInner::Tls(s) => s.writes_vectored(),
46        }
47    }
48}
49
50impl<S> MaybeTlsStream<S>
51where
52    S: AsyncRead + AsyncWrite + Unpin,
53{
54    /// Returns the negotiated ALPN protocol.
55    pub fn negotiated_alpn(&self) -> Option<Cow<'_, [u8]>> {
56        match &self.0 {
57            MaybeTlsStreamInner::Plain(_) => None,
58            MaybeTlsStreamInner::Tls(s) => s.negotiated_alpn(),
59        }
60    }
61}
62
63impl<S> AsyncRead for MaybeTlsStream<S>
64where
65    S: AsyncRead + AsyncWrite + Unpin,
66{
67    fn poll_read(
68        self: Pin<&mut Self>,
69        cx: &mut Context<'_>,
70        buf: &mut [u8],
71    ) -> Poll<io::Result<usize>> {
72        match &mut self.get_mut().0 {
73            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_read(cx, buf),
74            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
75        }
76    }
77}
78
79impl<S> AsyncWrite for MaybeTlsStream<S>
80where
81    S: AsyncRead + AsyncWrite + Unpin,
82{
83    fn poll_write(
84        self: Pin<&mut Self>,
85        cx: &mut Context<'_>,
86        buf: &[u8],
87    ) -> Poll<io::Result<usize>> {
88        match &mut self.get_mut().0 {
89            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_write(cx, buf),
90            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
91        }
92    }
93
94    fn poll_write_vectored(
95        self: Pin<&mut Self>,
96        cx: &mut Context<'_>,
97        bufs: &[io::IoSlice<'_>],
98    ) -> Poll<io::Result<usize>> {
99        match &mut self.get_mut().0 {
100            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
101            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
102        }
103    }
104
105    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
106        match &mut self.get_mut().0 {
107            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_flush(cx),
108            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_flush(cx),
109        }
110    }
111
112    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
113        match &mut self.get_mut().0 {
114            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_close(cx),
115            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_close(cx),
116        }
117    }
118}
compio_tls/maybe.rs

compio_tls/
maybe.rs
