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/// A wrapper around an underlying raw stream which implements the TLS or SSL
50/// protocol.
51///
52/// A `TlsStream<S>` represents a handshake that has been completed successfully
53/// and both the server and the client are ready for receiving and sending
54/// data. Bytes read from a `TlsStream` are decrypted from `S` and bytes written
55/// to a `TlsStream` are encrypted when passing through to `S`.
56#[derive(Debug)]
57pub struct TlsStream<S>(TlsStreamInner<S>);
58
59impl<S> TlsStream<S>
60where
61    S: AsyncRead + AsyncWrite + Unpin,
62{
63    /// Returns the negotiated ALPN protocol.
64    pub fn negotiated_alpn(&self) -> Option<Cow<'_, [u8]>> {
65        self.0.negotiated_alpn()
66    }
67}
68
69#[cfg(feature = "native-tls")]
70#[doc(hidden)]
71impl<S> From<crate::native::TlsStream<S>> for TlsStream<S> {
72    fn from(value: crate::native::TlsStream<S>) -> Self {
73        Self(TlsStreamInner::NativeTls(value))
74    }
75}
76
77#[cfg(feature = "rustls")]
78#[doc(hidden)]
79impl<S> From<futures_rustls::client::TlsStream<S>> for TlsStream<S> {
80    fn from(value: futures_rustls::client::TlsStream<S>) -> Self {
81        Self(TlsStreamInner::Rustls(futures_rustls::TlsStream::Client(
82            value,
83        )))
84    }
85}
86
87#[cfg(feature = "rustls")]
88#[doc(hidden)]
89impl<S> From<futures_rustls::server::TlsStream<S>> for TlsStream<S> {
90    fn from(value: futures_rustls::server::TlsStream<S>) -> Self {
91        Self(TlsStreamInner::Rustls(futures_rustls::TlsStream::Server(
92            value,
93        )))
94    }
95}
96
97#[cfg(feature = "py-dynamic-openssl")]
98#[doc(hidden)]
99impl<S> From<crate::py_ossl::TlsStream<S>> for TlsStream<S> {
100    fn from(value: crate::py_ossl::TlsStream<S>) -> Self {
101        Self(TlsStreamInner::PyDynamicOpenSsl(value))
102    }
103}
104
105impl<S> AsyncRead for TlsStream<S>
106where
107    S: AsyncRead + AsyncWrite + Unpin,
108{
109    fn poll_read(
110        self: Pin<&mut Self>,
111        cx: &mut Context<'_>,
112        buf: &mut [u8],
113    ) -> Poll<io::Result<usize>> {
114        match &mut self.get_mut().0 {
115            #[cfg(feature = "native-tls")]
116            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_read(cx, buf),
117            #[cfg(feature = "rustls")]
118            TlsStreamInner::Rustls(s) => Pin::new(s).poll_read(cx, buf),
119            #[cfg(feature = "py-dynamic-openssl")]
120            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_read(cx, buf),
121            #[cfg(not(any(
122                feature = "native-tls",
123                feature = "rustls",
124                feature = "py-dynamic-openssl",
125            )))]
126            TlsStreamInner::None(f, ..) => match *f {},
127        }
128    }
129}
130
131impl<S> AsyncWrite for TlsStream<S>
132where
133    S: AsyncRead + AsyncWrite + Unpin,
134{
135    fn poll_write(
136        self: Pin<&mut Self>,
137        cx: &mut Context<'_>,
138        buf: &[u8],
139    ) -> Poll<io::Result<usize>> {
140        match &mut self.get_mut().0 {
141            #[cfg(feature = "native-tls")]
142            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_write(cx, buf),
143            #[cfg(feature = "rustls")]
144            TlsStreamInner::Rustls(s) => Pin::new(s).poll_write(cx, buf),
145            #[cfg(feature = "py-dynamic-openssl")]
146            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_write(cx, buf),
147            #[cfg(not(any(
148                feature = "native-tls",
149                feature = "rustls",
150                feature = "py-dynamic-openssl",
151            )))]
152            TlsStreamInner::None(f, ..) => match *f {},
153        }
154    }
155
156    fn poll_write_vectored(
157        self: Pin<&mut Self>,
158        cx: &mut Context<'_>,
159        bufs: &[io::IoSlice<'_>],
160    ) -> Poll<io::Result<usize>> {
161        match &mut self.get_mut().0 {
162            #[cfg(feature = "native-tls")]
163            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
164            #[cfg(feature = "rustls")]
165            TlsStreamInner::Rustls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
166            #[cfg(feature = "py-dynamic-openssl")]
167            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_write_vectored(cx, bufs),
168            #[cfg(not(any(
169                feature = "native-tls",
170                feature = "rustls",
171                feature = "py-dynamic-openssl",
172            )))]
173            TlsStreamInner::None(f, ..) => match *f {},
174        }
175    }
176
177    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
178        match &mut self.get_mut().0 {
179            #[cfg(feature = "native-tls")]
180            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_flush(cx),
181            #[cfg(feature = "rustls")]
182            TlsStreamInner::Rustls(s) => Pin::new(s).poll_flush(cx),
183            #[cfg(feature = "py-dynamic-openssl")]
184            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_flush(cx),
185            #[cfg(not(any(
186                feature = "native-tls",
187                feature = "rustls",
188                feature = "py-dynamic-openssl",
189            )))]
190            TlsStreamInner::None(f, ..) => match *f {},
191        }
192    }
193
194    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
195        match &mut self.get_mut().0 {
196            #[cfg(feature = "native-tls")]
197            TlsStreamInner::NativeTls(s) => Pin::new(s).poll_close(cx),
198            #[cfg(feature = "rustls")]
199            TlsStreamInner::Rustls(s) => Pin::new(s).poll_close(cx),
200            #[cfg(feature = "py-dynamic-openssl")]
201            TlsStreamInner::PyDynamicOpenSsl(s) => Pin::new(s).poll_close(cx),
202            #[cfg(not(any(
203                feature = "native-tls",
204                feature = "rustls",
205                feature = "py-dynamic-openssl",
206            )))]
207            TlsStreamInner::None(f, ..) => match *f {},
208        }
209    }
210}