|
|
|
@ -12,6 +12,9 @@ pub struct TlsStream<IO> { |
|
|
|
|
pub(crate) io: IO, |
|
|
|
|
pub(crate) session: ClientConnection, |
|
|
|
|
pub(crate) state: TlsState, |
|
|
|
|
|
|
|
|
|
#[cfg(feature = "early-data")] |
|
|
|
|
pub(crate) early_waker: Option<std::task::Waker>, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl<IO> TlsStream<IO> { |
|
|
|
@ -82,7 +85,26 @@ where |
|
|
|
|
) -> Poll<io::Result<()>> { |
|
|
|
|
match self.state { |
|
|
|
|
#[cfg(feature = "early-data")] |
|
|
|
|
TlsState::EarlyData(..) => Poll::Pending, |
|
|
|
|
TlsState::EarlyData(..) => { |
|
|
|
|
let this = self.get_mut(); |
|
|
|
|
|
|
|
|
|
// In the EarlyData state, we have not really established a Tls connection.
|
|
|
|
|
// Before writing data through `AsyncWrite` and completing the tls handshake,
|
|
|
|
|
// we ignore read readiness and return to pending.
|
|
|
|
|
//
|
|
|
|
|
// In order to avoid event loss,
|
|
|
|
|
// we need to register a waker and wake it up after tls is connected.
|
|
|
|
|
if this |
|
|
|
|
.early_waker |
|
|
|
|
.as_ref() |
|
|
|
|
.filter(|waker| cx.waker().will_wake(waker)) |
|
|
|
|
.is_none() |
|
|
|
|
{ |
|
|
|
|
this.early_waker = Some(cx.waker().clone()); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
Poll::Pending |
|
|
|
|
} |
|
|
|
|
TlsState::Stream | TlsState::WriteShutdown => { |
|
|
|
|
let this = self.get_mut(); |
|
|
|
|
let mut stream = |
|
|
|
@ -134,9 +156,6 @@ where |
|
|
|
|
if let Some(mut early_data) = stream.session.early_data() { |
|
|
|
|
let len = match early_data.write(buf) { |
|
|
|
|
Ok(n) => n, |
|
|
|
|
Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => { |
|
|
|
|
return Poll::Pending |
|
|
|
|
} |
|
|
|
|
Err(err) => return Poll::Ready(Err(err)), |
|
|
|
|
}; |
|
|
|
|
if len != 0 { |
|
|
|
@ -160,6 +179,11 @@ where |
|
|
|
|
|
|
|
|
|
// end
|
|
|
|
|
this.state = TlsState::Stream; |
|
|
|
|
|
|
|
|
|
if let Some(waker) = this.early_waker.take() { |
|
|
|
|
waker.wake(); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
stream.as_mut_pin().poll_write(cx, buf) |
|
|
|
|
} |
|
|
|
|
_ => stream.as_mut_pin().poll_write(cx, buf), |
|
|
|
@ -188,6 +212,10 @@ where |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
this.state = TlsState::Stream; |
|
|
|
|
|
|
|
|
|
if let Some(waker) = this.early_waker.take() { |
|
|
|
|
waker.wake(); |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
@ -195,19 +223,19 @@ where |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
|
|
|
|
if self.state.writeable() { |
|
|
|
|
self.session.send_close_notify(); |
|
|
|
|
self.state.shutdown_write(); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
#[cfg(feature = "early-data")] |
|
|
|
|
{ |
|
|
|
|
// we skip the handshake
|
|
|
|
|
if let TlsState::EarlyData(..) = self.state { |
|
|
|
|
return Pin::new(&mut self.io).poll_shutdown(cx); |
|
|
|
|
// complete handshake
|
|
|
|
|
if matches!(self.state, TlsState::EarlyData(..)) { |
|
|
|
|
ready!(self.as_mut().poll_flush(cx))?; |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
if self.state.writeable() { |
|
|
|
|
self.session.send_close_notify(); |
|
|
|
|
self.state.shutdown_write(); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
let this = self.get_mut(); |
|
|
|
|
let mut stream = |
|
|
|
|
Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable()); |
|
|
|
|