futures_rustls/
client.rs

1use super::*;
2use crate::common::IoSession;
3#[cfg(unix)]
4use std::os::unix::io::{AsRawFd, RawFd};
5#[cfg(windows)]
6use std::os::windows::io::{AsRawSocket, RawSocket};
7
8/// A wrapper around an underlying raw stream which implements the TLS or SSL
9/// protocol.
10#[derive(Debug)]
11pub struct TlsStream<IO> {
12    pub(crate) io: IO,
13    pub(crate) session: ClientConnection,
14    pub(crate) state: TlsState,
15
16    #[cfg(feature = "early-data")]
17    pub(crate) early_waker: Option<std::task::Waker>,
18}
19
20impl<IO> TlsStream<IO> {
21    #[inline]
22    pub fn get_ref(&self) -> (&IO, &ClientConnection) {
23        (&self.io, &self.session)
24    }
25
26    #[inline]
27    pub fn get_mut(&mut self) -> (&mut IO, &mut ClientConnection) {
28        (&mut self.io, &mut self.session)
29    }
30
31    #[inline]
32    pub fn into_inner(self) -> (IO, ClientConnection) {
33        (self.io, self.session)
34    }
35}
36
37#[cfg(unix)]
38impl<S> AsRawFd for TlsStream<S>
39where
40    S: AsRawFd,
41{
42    fn as_raw_fd(&self) -> RawFd {
43        self.get_ref().0.as_raw_fd()
44    }
45}
46
47#[cfg(windows)]
48impl<S> AsRawSocket for TlsStream<S>
49where
50    S: AsRawSocket,
51{
52    fn as_raw_socket(&self) -> RawSocket {
53        self.get_ref().0.as_raw_socket()
54    }
55}
56
57impl<IO> IoSession for TlsStream<IO> {
58    type Io = IO;
59    type Session = ClientConnection;
60
61    #[inline]
62    fn skip_handshake(&self) -> bool {
63        self.state.is_early_data()
64    }
65
66    #[inline]
67    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session) {
68        (&mut self.state, &mut self.io, &mut self.session)
69    }
70
71    #[inline]
72    fn into_io(self) -> Self::Io {
73        self.io
74    }
75}
76
77impl<IO> AsyncRead for TlsStream<IO>
78where
79    IO: AsyncRead + AsyncWrite + Unpin,
80{
81    fn poll_read(
82        self: Pin<&mut Self>,
83        cx: &mut Context<'_>,
84        buf: &mut [u8],
85    ) -> Poll<io::Result<usize>> {
86        match self.state {
87            #[cfg(feature = "early-data")]
88            TlsState::EarlyData(..) => {
89                let this = self.get_mut();
90
91                // In the EarlyData state, we have not really established a Tls connection.
92                // Before writing data through `AsyncWrite` and completing the tls handshake,
93                // we ignore read readiness and return to pending.
94                //
95                // In order to avoid event loss,
96                // we need to register a waker and wake it up after tls is connected.
97                if this
98                    .early_waker
99                    .as_ref()
100                    .filter(|waker| cx.waker().will_wake(waker))
101                    .is_none()
102                {
103                    this.early_waker = Some(cx.waker().clone());
104                }
105
106                Poll::Pending
107            }
108            TlsState::Stream | TlsState::WriteShutdown => {
109                let this = self.get_mut();
110                let mut stream =
111                    Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
112
113                match stream.as_mut_pin().poll_read(cx, buf) {
114                    Poll::Ready(Ok(n)) => {
115                        if n == 0 || stream.eof {
116                            this.state.shutdown_read();
117                        }
118
119                        Poll::Ready(Ok(n))
120                    }
121                    Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::ConnectionAborted => {
122                        this.state.shutdown_read();
123                        Poll::Ready(Err(err))
124                    }
125                    output => output,
126                }
127            }
128            TlsState::ReadShutdown | TlsState::FullyShutdown => Poll::Ready(Ok(0)),
129        }
130    }
131}
132
133impl<IO> AsyncWrite for TlsStream<IO>
134where
135    IO: AsyncRead + AsyncWrite + Unpin,
136{
137    /// Note: that it does not guarantee the final data to be sent.
138    /// To be cautious, you must manually call `flush`.
139    fn poll_write(
140        self: Pin<&mut Self>,
141        cx: &mut Context<'_>,
142        buf: &[u8],
143    ) -> Poll<io::Result<usize>> {
144        let this = self.get_mut();
145        let mut stream =
146            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
147
148        #[allow(clippy::match_single_binding)]
149        match this.state {
150            #[cfg(feature = "early-data")]
151            TlsState::EarlyData(ref mut pos, ref mut data) => {
152                use std::io::Write;
153
154                // write early data
155                if let Some(mut early_data) = stream.session.early_data() {
156                    let len = match early_data.write(buf) {
157                        Ok(n) => n,
158                        Err(err) => return Poll::Ready(Err(err)),
159                    };
160                    if len != 0 {
161                        data.extend_from_slice(&buf[..len]);
162                        return Poll::Ready(Ok(len));
163                    }
164                }
165
166                // complete handshake
167                while stream.session.is_handshaking() {
168                    ready!(stream.handshake(cx))?;
169                }
170
171                // write early data (fallback)
172                if !stream.session.is_early_data_accepted() {
173                    while *pos < data.len() {
174                        let len = ready!(stream.as_mut_pin().poll_write(cx, &data[*pos..]))?;
175                        *pos += len;
176                    }
177                }
178
179                // end
180                this.state = TlsState::Stream;
181
182                if let Some(waker) = this.early_waker.take() {
183                    waker.wake();
184                }
185
186                stream.as_mut_pin().poll_write(cx, buf)
187            }
188            _ => stream.as_mut_pin().poll_write(cx, buf),
189        }
190    }
191
192    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
193        let this = self.get_mut();
194        let mut stream =
195            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
196
197        #[cfg(feature = "early-data")]
198        {
199            if let TlsState::EarlyData(ref mut pos, ref mut data) = this.state {
200                // complete handshake
201                while stream.session.is_handshaking() {
202                    ready!(stream.handshake(cx))?;
203                }
204
205                // write early data (fallback)
206                if !stream.session.is_early_data_accepted() {
207                    while *pos < data.len() {
208                        let len = ready!(stream.as_mut_pin().poll_write(cx, &data[*pos..]))?;
209                        *pos += len;
210                    }
211                }
212
213                this.state = TlsState::Stream;
214
215                if let Some(waker) = this.early_waker.take() {
216                    waker.wake();
217                }
218            }
219        }
220
221        stream.as_mut_pin().poll_flush(cx)
222    }
223
224    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
225        #[cfg(feature = "early-data")]
226        {
227            // complete handshake
228            if matches!(self.state, TlsState::EarlyData(..)) {
229                ready!(self.as_mut().poll_flush(cx))?;
230            }
231        }
232
233        if self.state.writeable() {
234            self.session.send_close_notify();
235            self.state.shutdown_write();
236        }
237
238        let this = self.get_mut();
239        let mut stream =
240            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
241        stream.as_mut_pin().poll_close(cx)
242    }
243}
futures_rustls/client.rs

futures_rustls/
client.rs
