netxserver/server/
maybe_stream.rs

1use std::pin::Pin;
2use std::task::{Context, Poll};
3use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
4use tokio::net::TcpStream;
5#[cfg(all(feature = "use_openssl", not(feature = "use_rustls")))]
6use tokio_openssl::SslStream;
7
8#[cfg(all(feature = "use_rustls", not(feature = "use_openssl")))]
9use tokio_rustls::server::TlsStream;
10
11/// Enum representing a stream that can be either plain TCP or TLS/SSL.
12#[derive(Debug)]
13pub enum MaybeStream {
14    Plain(TcpStream),
15    #[cfg(all(feature = "use_openssl", not(feature = "use_rustls")))]
16    ServerSsl(SslStream<TcpStream>),
17    #[cfg(all(feature = "use_rustls", not(feature = "use_openssl")))]
18    ServerTls(TlsStream<TcpStream>),
19}
20
21impl AsyncRead for MaybeStream {
22    /// Polls for reading from the stream.
23    ///
24    /// # Arguments
25    ///
26    /// * `cx` - The context of the current task.
27    /// * `buf` - The buffer to read data into.
28    ///
29    /// # Returns
30    ///
31    /// A `Poll` indicating the result of the read operation.
32    #[inline]
33    fn poll_read(
34        self: Pin<&mut Self>,
35        cx: &mut Context<'_>,
36        buf: &mut ReadBuf<'_>,
37    ) -> Poll<std::io::Result<()>> {
38        match self.get_mut() {
39            MaybeStream::Plain(ref mut s) => Pin::new(s).poll_read(cx, buf),
40            #[cfg(all(feature = "use_openssl", not(feature = "use_rustls")))]
41            MaybeStream::ServerSsl(ref mut s) => Pin::new(s).poll_read(cx, buf),
42            #[cfg(all(feature = "use_rustls", not(feature = "use_openssl")))]
43            MaybeStream::ServerTls(ref mut s) => Pin::new(s).poll_read(cx, buf),
44        }
45    }
46}
47
48impl AsyncWrite for MaybeStream {
49    /// Polls for writing to the stream.
50    ///
51    /// # Arguments
52    ///
53    /// * `cx` - The context of the current task.
54    /// * `buf` - The buffer containing data to write.
55    ///
56    /// # Returns
57    ///
58    /// A `Poll` indicating the result of the write operation.
59    #[inline]
60    fn poll_write(
61        self: Pin<&mut Self>,
62        cx: &mut Context<'_>,
63        buf: &[u8],
64    ) -> Poll<Result<usize, std::io::Error>> {
65        match self.get_mut() {
66            MaybeStream::Plain(ref mut s) => Pin::new(s).poll_write(cx, buf),
67            #[cfg(all(feature = "use_openssl", not(feature = "use_rustls")))]
68            MaybeStream::ServerSsl(ref mut s) => Pin::new(s).poll_write(cx, buf),
69            #[cfg(all(feature = "use_rustls", not(feature = "use_openssl")))]
70            MaybeStream::ServerTls(ref mut s) => Pin::new(s).poll_write(cx, buf),
71        }
72    }
73
74    /// Polls for flushing the stream.
75    ///
76    /// # Arguments
77    ///
78    /// * `cx` - The context of the current task.
79    ///
80    /// # Returns
81    ///
82    /// A `Poll` indicating the result of the flush operation.
83    #[inline]
84    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
85        match self.get_mut() {
86            MaybeStream::Plain(ref mut s) => Pin::new(s).poll_flush(cx),
87            #[cfg(all(feature = "use_openssl", not(feature = "use_rustls")))]
88            MaybeStream::ServerSsl(ref mut s) => Pin::new(s).poll_flush(cx),
89            #[cfg(all(feature = "use_rustls", not(feature = "use_openssl")))]
90            MaybeStream::ServerTls(ref mut s) => Pin::new(s).poll_flush(cx),
91        }
92    }
93
94    /// Polls for shutting down the stream.
95    ///
96    /// # Arguments
97    ///
98    /// * `cx` - The context of the current task.
99    ///
100    /// # Returns
101    ///
102    /// A `Poll` indicating the result of the shutdown operation.
103    #[inline]
104    fn poll_shutdown(
105        self: Pin<&mut Self>,
106        cx: &mut Context<'_>,
107    ) -> Poll<Result<(), std::io::Error>> {
108        match self.get_mut() {
109            MaybeStream::Plain(ref mut s) => Pin::new(s).poll_shutdown(cx),
110            #[cfg(all(feature = "use_openssl", not(feature = "use_rustls")))]
111            MaybeStream::ServerSsl(ref mut s) => Pin::new(s).poll_shutdown(cx),
112            #[cfg(all(feature = "use_rustls", not(feature = "use_openssl")))]
113            MaybeStream::ServerTls(ref mut s) => Pin::new(s).poll_shutdown(cx),
114        }
115    }
116}
netxserver/server/maybe_stream.rs

netxserver/server/
maybe_stream.rs
