Skip to main content

compio_tls/
stream.rs

1use std::{
2    borrow::Cow,
3    io,
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use futures_util::{AsyncRead, AsyncWrite};
9
10#[derive(Debug)]
11#[allow(clippy::large_enum_variant)]
12enum TlsStreamInner<S> {
13    #[cfg(feature = "native-tls")]
14    NativeTls(crate::native::TlsStream<S>),
15    #[cfg(feature = "rustls")]
16    Rustls(futures_rustls::TlsStream<S>),
17    #[cfg(feature = "py-dynamic-openssl")]
18    PyDynamicOpenSsl(crate::py_ossl::TlsStream<S>),
19    #[cfg(not(any(
20        feature = "native-tls",
21        feature = "rustls",
22        feature = "py-dynamic-openssl",
23    )))]
24    None(std::convert::Infallible, std::marker::PhantomData<S>),
25}
26
27impl<S> TlsStreamInner<S>
28where
29    S: AsyncRead + AsyncWrite + Unpin,
30{
31    pub fn negotiated_alpn(&self) -> Option<Cow<'_, [u8]>> {
32        match self {
33            #[cfg(feature = "native-tls")]
34            Self::NativeTls(s) => s.negotiated_alpn().ok().flatten().map(Cow::from),
35            #[cfg(feature = "rustls")]
36            Self::Rustls(s) => s.get_ref().1.alpn_protocol().map(Cow::from),
37            #[cfg(feature = "py-dynamic-openssl")]
38            Self::PyDynamicOpenSsl(s) => s.negotiated_alpn().map(Cow::from),
39            #[cfg(not(any(
40                feature = "native-tls",
41                feature = "rustls",
42                feature = "py-dynamic-openssl",
43            )))]
44            Self::None(f, ..) => match *f {},
45        }
46    }
47}
48
49#[cfg(feature = "tokio")]
50impl<S> TlsStream<S> {
51    /// Whether `poll_write_vectored` takes all the buffers at once, not only
52    /// the first: rustls copies them all into its own buffer.
53    pub(crate) fn writes_vectored(&self) -> bool {
54        cfg_select! {
55            feature = "rustls" => matches!(self.0, TlsStreamInner::Rustls(_)),
56            _ => false,
57        }
58    }
59}
60
61/// A wrapper around an underlying raw stream which implements the TLS or SSL
62/// protocol.
63///
64/// A `TlsStream<S>` represents a handshake that has been completed successfully
65/// and both the server and the client are ready for receiving and sending
66/// data. Bytes read from a `TlsStream` are decrypted from `S` and bytes written
67/// to a `TlsStream` are encrypted when passing through to `S`.
68#[derive(Debug)]
69pub struct TlsStream<S>(TlsStreamInner<S>);
70
71impl<S> TlsStream<S>
72where
73    S: AsyncRead + AsyncWrite + Unpin,
74{
75    /// Returns the negotiated ALPN protocol.
76    pub fn negotiated_alpn(&self) -> Option<Cow<'_, [u8]>> {
77        self.0.negotiated_alpn()
78    }
79}
80
81#[cfg(feature = "native-tls")]
82#[doc(hidden)]
83impl<S> From<crate::native::TlsStream<S>> for TlsStream<S> {
84    fn from(value: crate::native::TlsStream<S>) -> Self {
85        Self(TlsStreamInner::NativeTls(value))
86    }
87}
88
89#[cfg(feature = "rustls")]
90#[doc(hidden)]
91impl<S> From<futures_rustls::client::TlsStream<S>> for TlsStream<S> {
92    fn from(value: futures_rustls::client::TlsStream<S>) -> Self {
93        Self(TlsStreamInner::Rustls(futures_rustls::TlsStream::Client(
94            value,
95        )))
96    }
97}
98
99#[cfg(feature = "rustls")]
100#[doc(hidden)]
101impl<S> From<futures_rustls::server::TlsStream<S>> for TlsStream<S> {
102    fn from(value: futures_rustls::server::TlsStream<S>) -> Self {
103        Self(TlsStreamInner::Rustls(futures_rustls::TlsStream::Server(
104            value,
105        )))
106    }
107}
108
109#[cfg(feature = "py-dynamic-openssl")]
110#[doc(hidden)]
111impl<S> From<crate::py_ossl::TlsStream<S>> for TlsStream<S> {
112    fn from(value: crate::py_ossl::TlsStream<S>) -> Self {
113        Self(TlsStreamInner::PyDynamicOpenSsl(value))
114    }
115}
116
117impl<S> AsyncRead for TlsStream<S>
118where
119    S: AsyncRead + AsyncWrite + Unpin,
120{
121    fn poll_read(
122        self: Pin<&mut Self>,
123        cx: &mut Context<'_>,
124        buf: &mut [u8],
125    ) -> Poll<io::Result<usize>> {
126        match &mut self.get_mut().0 {
127            #[cfg(feature = "native-tls")]
128            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_read(cx, buf),
129            #[cfg(feature = "rustls")]
130            TlsStreamInner::Rustls(s) => Pin::new(s).poll_read(cx, buf),
131            #[cfg(feature = "py-dynamic-openssl")]
132            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_read(cx, buf),
133            #[cfg(not(any(
134                feature = "native-tls",
135                feature = "rustls",
136                feature = "py-dynamic-openssl",
137            )))]
138            TlsStreamInner::None(f, ..) => match *f {},
139        }
140    }
141}
142
143impl<S> AsyncWrite for TlsStream<S>
144where
145    S: AsyncRead + AsyncWrite + Unpin,
146{
147    fn poll_write(
148        self: Pin<&mut Self>,
149        cx: &mut Context<'_>,
150        buf: &[u8],
151    ) -> Poll<io::Result<usize>> {
152        match &mut self.get_mut().0 {
153            #[cfg(feature = "native-tls")]
154            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_write(cx, buf),
155            #[cfg(feature = "rustls")]
156            TlsStreamInner::Rustls(s) => Pin::new(s).poll_write(cx, buf),
157            #[cfg(feature = "py-dynamic-openssl")]
158            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_write(cx, buf),
159            #[cfg(not(any(
160                feature = "native-tls",
161                feature = "rustls",
162                feature = "py-dynamic-openssl",
163            )))]
164            TlsStreamInner::None(f, ..) => match *f {},
165        }
166    }
167
168    fn poll_write_vectored(
169        self: Pin<&mut Self>,
170        cx: &mut Context<'_>,
171        bufs: &[io::IoSlice<'_>],
172    ) -> Poll<io::Result<usize>> {
173        match &mut self.get_mut().0 {
174            #[cfg(feature = "native-tls")]
175            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
176            #[cfg(feature = "rustls")]
177            TlsStreamInner::Rustls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
178            #[cfg(feature = "py-dynamic-openssl")]
179            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_write_vectored(cx, bufs),
180            #[cfg(not(any(
181                feature = "native-tls",
182                feature = "rustls",
183                feature = "py-dynamic-openssl",
184            )))]
185            TlsStreamInner::None(f, ..) => match *f {},
186        }
187    }
188
189    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
190        match &mut self.get_mut().0 {
191            #[cfg(feature = "native-tls")]
192            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_flush(cx),
193            #[cfg(feature = "rustls")]
194            TlsStreamInner::Rustls(s) => Pin::new(s).poll_flush(cx),
195            #[cfg(feature = "py-dynamic-openssl")]
196            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_flush(cx),
197            #[cfg(not(any(
198                feature = "native-tls",
199                feature = "rustls",
200                feature = "py-dynamic-openssl",
201            )))]
202            TlsStreamInner::None(f, ..) => match *f {},
203        }
204    }
205
206    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
207        match &mut self.get_mut().0 {
208            #[cfg(feature = "native-tls")]
209            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_close(cx),
210            #[cfg(feature = "rustls")]
211            TlsStreamInner::Rustls(s) => Pin::new(s).poll_close(cx),
212            #[cfg(feature = "py-dynamic-openssl")]
213            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_close(cx),
214            #[cfg(not(any(
215                feature = "native-tls",
216                feature = "rustls",
217                feature = "py-dynamic-openssl",
218            )))]
219            TlsStreamInner::None(f, ..) => match *f {},
220        }
221    }
222}