tokio_rustls/common/
handshake.rs

Help
1use crate::common::{Stream, TlsState};
2use rustls::{ConnectionCommon, SideData};
3use std::future::Future;
4use std::ops::{Deref, DerefMut};
5use std::pin::Pin;
6use std::task::{Context, Poll};
7use std::{io, mem};
8use tokio::io::{AsyncRead, AsyncWrite};
9
10pub(crate) trait IoSession {
11    type Io;
12    type Session;
13
14    fn skip_handshake(&self) -> bool;
15    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session);
16    fn into_io(self) -> Self::Io;
17}
18
19pub(crate) enum MidHandshake<IS: IoSession> {
20    Handshaking(IS),
21    End,
22    Error { io: IS::Io, error: io::Error },
23}
24
25impl<IS, SD> Future for MidHandshake<IS>
26where
27    IS: IoSession + Unpin,
28    IS::Io: AsyncRead + AsyncWrite + Unpin,
29    IS::Session: DerefMut + Deref<Target = ConnectionCommon<SD>> + Unpin,
30    SD: SideData,
31{
32    type Output = Result<IS, (io::Error, IS::Io)>;
33
34    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
35        let this = self.get_mut();
36
37        let mut stream = match mem::replace(this, MidHandshake::End) {
38            MidHandshake::Handshaking(stream) => stream,
39            // Starting the handshake returned an error; fail the future immediately.
40            MidHandshake::Error { io, error } => return Poll::Ready(Err((error, io))),
41            _ => panic!("unexpected polling after handshake"),
42        };
43
44        if !stream.skip_handshake() {
45            let (state, io, session) = stream.get_mut();
46            let mut tls_stream = Stream::new(io, session).set_eof(!state.readable());
47
48            macro_rules! try_poll {
49                ( $e:expr ) => {
50                    match $e {
51                        Poll::Ready(Ok(_)) => (),
52                        Poll::Ready(Err(err)) => return Poll::Ready(Err((err, stream.into_io()))),
53                        Poll::Pending => {
54                            *this = MidHandshake::Handshaking(stream);
55                            return Poll::Pending;
56                        }
57                    }
58                };
59            }
60
61            while tls_stream.session.is_handshaking() {
62                try_poll!(tls_stream.handshake(cx));
63            }
64
65            try_poll!(Pin::new(&mut tls_stream).poll_flush(cx));
66        }
67
68        Poll::Ready(Ok(stream))
69    }
70}
tokio_rustls/common/handshake.rs

tokio_rustls/common/
handshake.rs
