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
40impl<S> MaybeTlsStream<S>
41where
42    S: AsyncRead + AsyncWrite + Unpin,
43{
44    /// Returns the negotiated ALPN protocol.
45    pub fn negotiated_alpn(&self) -> Option<Cow<'_, [u8]>> {
46        match &self.0 {
47            MaybeTlsStreamInner::Plain(_) => None,
48            MaybeTlsStreamInner::Tls(s) => s.negotiated_alpn(),
49        }
50    }
51}
52
53impl<S> AsyncRead for MaybeTlsStream<S>
54where
55    S: AsyncRead + AsyncWrite + Unpin,
56{
57    fn poll_read(
58        self: Pin<&mut Self>,
59        cx: &mut Context<'_>,
60        buf: &mut [u8],
61    ) -> Poll<io::Result<usize>> {
62        match &mut self.get_mut().0 {
63            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_read(cx, buf),
64            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
65        }
66    }
67}
68
69impl<S> AsyncWrite for MaybeTlsStream<S>
70where
71    S: AsyncRead + AsyncWrite + Unpin,
72{
73    fn poll_write(
74        self: Pin<&mut Self>,
75        cx: &mut Context<'_>,
76        buf: &[u8],
77    ) -> Poll<io::Result<usize>> {
78        match &mut self.get_mut().0 {
79            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_write(cx, buf),
80            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
81        }
82    }
83
84    fn poll_write_vectored(
85        self: Pin<&mut Self>,
86        cx: &mut Context<'_>,
87        bufs: &[io::IoSlice<'_>],
88    ) -> Poll<io::Result<usize>> {
89        match &mut self.get_mut().0 {
90            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
91            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
92        }
93    }
94
95    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
96        match &mut self.get_mut().0 {
97            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_flush(cx),
98            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_flush(cx),
99        }
100    }
101
102    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
103        match &mut self.get_mut().0 {
104            MaybeTlsStreamInner::Plain(stream) => Pin::new(stream).poll_close(cx),
105            MaybeTlsStreamInner::Tls(stream) => Pin::new(stream).poll_close(cx),
106        }
107    }
108}
compio_tls/maybe.rs

compio_tls/
maybe.rs
