new_tokio_smtp/io/
socket.rs

1use std::fmt::Debug;
2use std::io as std_io;
3
4use bytes::buf::{Buf, BufMut};
5use futures::Poll;
6use tokio::io::{AsyncRead, AsyncWrite};
7use tokio::net::TcpStream;
8use tokio_tls::TlsStream;
9
10/// Abstraction over Tcp, TcpTls (and Mock)
11///
12/// Allows treating both `TcpStream` and
13/// `TlsStream<TcpStream>` the same.
14///
15/// # Features
16/// ## `mock_support`
17///
18/// if enabled this abstracts not only over `TcpStream` and
19/// `TlsStream<TcpStream` but also `Box<MockStream+Send>`
20///
21#[derive(Debug)]
22pub enum Socket {
23    Secure(TlsStream<TcpStream>),
24    Insecure(TcpStream),
25    #[cfg(feature = "mock-support")]
26    Mock(Box<dyn MockStream + Send>),
27}
28
29impl Socket {
30    /// true if it's a `TlsStream` (or if mock says so)
31    pub fn is_secure(&self) -> bool {
32        match self {
33            Socket::Secure(_) => true,
34            Socket::Insecure(_) => false,
35            #[cfg(feature = "mock-support")]
36            Socket::Mock(mock) => mock.is_secure(),
37        }
38    }
39}
40
41macro_rules! socket_mux {
42    ($self:ident, |$socket:ident| $block:block) => {{
43        match $self {
44            Socket::Secure($socket) => $block,
45            Socket::Insecure($socket) => $block,
46            #[cfg(feature = "mock-support")]
47            Socket::Mock($socket) => $block,
48        }
49    }};
50}
51
52impl std_io::Read for Socket {
53    fn read(&mut self, buf: &mut [u8]) -> Result<usize, std_io::Error> {
54        socket_mux! {self, |socket| {
55            socket.read(buf)
56        }}
57    }
58}
59
60impl std_io::Write for Socket {
61    fn write(&mut self, buf: &[u8]) -> Result<usize, std_io::Error> {
62        socket_mux! {self, |socket| {
63            socket.write(buf)
64        }}
65    }
66
67    fn flush(&mut self) -> Result<(), std_io::Error> {
68        socket_mux! {self, |socket| {
69            socket.flush()
70        }}
71    }
72}
73
74impl AsyncRead for Socket {
75    #[inline]
76    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
77        match self {
78            Socket::Secure(socket) => socket.prepare_uninitialized_buffer(buf),
79            Socket::Insecure(socket) => socket.prepare_uninitialized_buffer(buf),
80            #[cfg(feature = "mock-support")]
81            Socket::Mock(socket) => socket.prepare_uninitialized_buffer(buf),
82        }
83    }
84
85    #[inline]
86    fn poll_read(&mut self, buf: &mut [u8]) -> Poll<usize, std_io::Error> {
87        socket_mux! {self, |socket| {
88            socket.poll_read(buf)
89        }}
90    }
91
92    #[inline]
93    fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, std_io::Error>
94    where
95        Self: Sized,
96    {
97        socket_mux! {self, |socket| {
98            socket.read_buf(buf)
99        }}
100    }
101}
102
103impl AsyncWrite for Socket {
104    fn poll_write(&mut self, buf: &[u8]) -> Poll<usize, std_io::Error> {
105        socket_mux! {self, |socket| {
106            AsyncWrite::poll_write(socket, buf)
107        }}
108    }
109
110    fn poll_flush(&mut self) -> Poll<(), std_io::Error> {
111        socket_mux! {self, |socket| {
112            AsyncWrite::poll_flush(socket)
113        }}
114    }
115
116    fn shutdown(&mut self) -> Poll<(), std_io::Error> {
117        socket_mux! {self, |socket| {
118            AsyncWrite::shutdown(socket)
119        }}
120    }
121
122    fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, std_io::Error>
123    where
124        Self: Sized,
125    {
126        socket_mux! {self, |socket| {
127            AsyncWrite::write_buf(socket, buf)
128        }}
129    }
130}
131
132/// trait representing a mock stream
133pub trait MockStream: Debug + AsyncRead + AsyncWrite + 'static {
134    fn is_secure(&self) -> bool {
135        false
136    }
137    fn set_is_secure(&mut self, secure: bool);
138}
new_tokio_smtp/io/socket.rs

new_tokio_smtp/io/
socket.rs
