use std::future::Future;
use std::io;
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
pub use rustls;
use rustls::server::AcceptedAlert;
use rustls::{ClientConfig, ClientConnection, CommonState, ServerConfig, ServerConnection};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
macro_rules! ready {
( $e:expr ) => {
match $e {
std::task::Poll::Ready(t) => t,
std::task::Poll::Pending => return std::task::Poll::Pending,
}
};
}
pub mod client;
mod common;
use common::{MidHandshake, TlsState};
pub mod server;
#[derive(Clone)]
pub struct TlsConnector {
inner: Arc<ClientConfig>,
#[cfg(feature = "early-data")]
early_data: bool,
}
#[derive(Clone)]
pub struct TlsAcceptor {
inner: Arc<ServerConfig>,
}
impl From<Arc<ClientConfig>> for TlsConnector {
fn from(inner: Arc<ClientConfig>) -> TlsConnector {
TlsConnector {
inner,
#[cfg(feature = "early-data")]
early_data: false,
}
}
}
impl From<Arc<ServerConfig>> for TlsAcceptor {
fn from(inner: Arc<ServerConfig>) -> TlsAcceptor {
TlsAcceptor { inner }
}
}
impl TlsConnector {
#[cfg(feature = "early-data")]
pub fn early_data(mut self, flag: bool) -> TlsConnector {
self.early_data = flag;
self
}
#[inline]
pub fn connect<IO>(&self, domain: pki_types::ServerName<'static>, stream: IO) -> Connect<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
{
self.connect_with(domain, stream, |_| ())
}
pub fn connect_with<IO, F>(
&self,
domain: pki_types::ServerName<'static>,
stream: IO,
f: F,
) -> Connect<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
F: FnOnce(&mut ClientConnection),
{
let mut session = match ClientConnection::new(self.inner.clone(), domain) {
Ok(session) => session,
Err(error) => {
return Connect(MidHandshake::Error {
io: stream,
error: io::Error::new(io::ErrorKind::Other, error),
});
}
};
f(&mut session);
Connect(MidHandshake::Handshaking(client::TlsStream {
io: stream,
#[cfg(not(feature = "early-data"))]
state: TlsState::Stream,
#[cfg(feature = "early-data")]
state: if self.early_data && session.early_data().is_some() {
TlsState::EarlyData(0, Vec::new())
} else {
TlsState::Stream
},
#[cfg(feature = "early-data")]
early_waker: None,
session,
}))
}
}
impl TlsAcceptor {
#[inline]
pub fn accept<IO>(&self, stream: IO) -> Accept<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
{
self.accept_with(stream, |_| ())
}
pub fn accept_with<IO, F>(&self, stream: IO, f: F) -> Accept<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
F: FnOnce(&mut ServerConnection),
{
let mut session = match ServerConnection::new(self.inner.clone()) {
Ok(session) => session,
Err(error) => {
return Accept(MidHandshake::Error {
io: stream,
error: io::Error::new(io::ErrorKind::Other, error),
});
}
};
f(&mut session);
Accept(MidHandshake::Handshaking(server::TlsStream {
session,
io: stream,
state: TlsState::Stream,
}))
}
}
pub struct LazyConfigAcceptor<IO> {
acceptor: rustls::server::Acceptor,
io: Option<IO>,
alert: Option<(rustls::Error, AcceptedAlert)>,
}
impl<IO> LazyConfigAcceptor<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
{
#[inline]
pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self {
Self {
acceptor,
io: Some(io),
alert: None,
}
}
pub fn take_io(&mut self) -> Option<IO> {
self.io.take()
}
}
impl<IO> Future for LazyConfigAcceptor<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
{
type Output = Result<StartHandshake<IO>, io::Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
loop {
let io = match this.io.as_mut() {
Some(io) => io,
None => {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::Other,
"acceptor cannot be polled after acceptance",
)))
}
};
if let Some((err, mut alert)) = this.alert.take() {
match alert.write(&mut common::SyncWriteAdapter { io, cx }) {
Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
this.alert = Some((err, alert));
return Poll::Pending;
}
Ok(0) | Err(_) => {
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err)))
}
Ok(_) => {
this.alert = Some((err, alert));
continue;
}
};
}
let mut reader = common::SyncReadAdapter { io, cx };
match this.acceptor.read_tls(&mut reader) {
Ok(0) => return Err(io::ErrorKind::UnexpectedEof.into()).into(),
Ok(_) => {}
Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Poll::Pending,
Err(e) => return Err(e).into(),
}
match this.acceptor.accept() {
Ok(Some(accepted)) => {
let io = this.io.take().unwrap();
return Poll::Ready(Ok(StartHandshake { accepted, io }));
}
Ok(None) => {}
Err((err, alert)) => {
this.alert = Some((err, alert));
}
}
}
}
}
pub struct StartHandshake<IO> {
accepted: rustls::server::Accepted,
io: IO,
}
impl<IO> StartHandshake<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
{
pub fn client_hello(&self) -> rustls::server::ClientHello<'_> {
self.accepted.client_hello()
}
pub fn into_stream(self, config: Arc<ServerConfig>) -> Accept<IO> {
self.into_stream_with(config, |_| ())
}
pub fn into_stream_with<F>(self, config: Arc<ServerConfig>, f: F) -> Accept<IO>
where
F: FnOnce(&mut ServerConnection),
{
let mut conn = match self.accepted.into_connection(config) {
Ok(conn) => conn,
Err((error, alert)) => {
return Accept(MidHandshake::SendAlert {
io: self.io,
alert,
error: io::Error::new(io::ErrorKind::InvalidData, error),
});
}
};
f(&mut conn);
Accept(MidHandshake::Handshaking(server::TlsStream {
session: conn,
io: self.io,
state: TlsState::Stream,
}))
}
}
pub struct Connect<IO>(MidHandshake<client::TlsStream<IO>>);
pub struct Accept<IO>(MidHandshake<server::TlsStream<IO>>);
pub struct FallibleConnect<IO>(MidHandshake<client::TlsStream<IO>>);
pub struct FallibleAccept<IO>(MidHandshake<server::TlsStream<IO>>);
impl<IO> Connect<IO> {
#[inline]
pub fn into_fallible(self) -> FallibleConnect<IO> {
FallibleConnect(self.0)
}
pub fn get_ref(&self) -> Option<&IO> {
match &self.0 {
MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
MidHandshake::SendAlert { io, .. } => Some(io),
MidHandshake::Error { io, .. } => Some(io),
MidHandshake::End => None,
}
}
pub fn get_mut(&mut self) -> Option<&mut IO> {
match &mut self.0 {
MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
MidHandshake::SendAlert { io, .. } => Some(io),
MidHandshake::Error { io, .. } => Some(io),
MidHandshake::End => None,
}
}
}
impl<IO> Accept<IO> {
#[inline]
pub fn into_fallible(self) -> FallibleAccept<IO> {
FallibleAccept(self.0)
}
pub fn get_ref(&self) -> Option<&IO> {
match &self.0 {
MidHandshake::Handshaking(sess) => Some(sess.get_ref().0),
MidHandshake::SendAlert { io, .. } => Some(io),
MidHandshake::Error { io, .. } => Some(io),
MidHandshake::End => None,
}
}
pub fn get_mut(&mut self) -> Option<&mut IO> {
match &mut self.0 {
MidHandshake::Handshaking(sess) => Some(sess.get_mut().0),
MidHandshake::SendAlert { io, .. } => Some(io),
MidHandshake::Error { io, .. } => Some(io),
MidHandshake::End => None,
}
}
}
impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Connect<IO> {
type Output = io::Result<client::TlsStream<IO>>;
#[inline]
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
}
}
impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Accept<IO> {
type Output = io::Result<server::TlsStream<IO>>;
#[inline]
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.0).poll(cx).map_err(|(err, _)| err)
}
}
impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleConnect<IO> {
type Output = Result<client::TlsStream<IO>, (io::Error, IO)>;
#[inline]
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.0).poll(cx)
}
}
impl<IO: AsyncRead + AsyncWrite + Unpin> Future for FallibleAccept<IO> {
type Output = Result<server::TlsStream<IO>, (io::Error, IO)>;
#[inline]
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.0).poll(cx)
}
}
#[allow(clippy::large_enum_variant)] #[derive(Debug)]
pub enum TlsStream<T> {
Client(client::TlsStream<T>),
Server(server::TlsStream<T>),
}
impl<T> TlsStream<T> {
pub fn get_ref(&self) -> (&T, &CommonState) {
use TlsStream::*;
match self {
Client(io) => {
let (io, session) = io.get_ref();
(io, session)
}
Server(io) => {
let (io, session) = io.get_ref();
(io, session)
}
}
}
pub fn get_mut(&mut self) -> (&mut T, &mut CommonState) {
use TlsStream::*;
match self {
Client(io) => {
let (io, session) = io.get_mut();
(io, &mut *session)
}
Server(io) => {
let (io, session) = io.get_mut();
(io, &mut *session)
}
}
}
}
impl<T> From<client::TlsStream<T>> for TlsStream<T> {
fn from(s: client::TlsStream<T>) -> Self {
Self::Client(s)
}
}
impl<T> From<server::TlsStream<T>> for TlsStream<T> {
fn from(s: server::TlsStream<T>) -> Self {
Self::Server(s)
}
}
#[cfg(unix)]
impl<S> AsRawFd for TlsStream<S>
where
S: AsRawFd,
{
fn as_raw_fd(&self) -> RawFd {
self.get_ref().0.as_raw_fd()
}
}
#[cfg(windows)]
impl<S> AsRawSocket for TlsStream<S>
where
S: AsRawSocket,
{
fn as_raw_socket(&self) -> RawSocket {
self.get_ref().0.as_raw_socket()
}
}
impl<T> AsyncRead for TlsStream<T>
where
T: AsyncRead + AsyncWrite + Unpin,
{
#[inline]
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
match self.get_mut() {
TlsStream::Client(x) => Pin::new(x).poll_read(cx, buf),
TlsStream::Server(x) => Pin::new(x).poll_read(cx, buf),
}
}
}
impl<T> AsyncWrite for TlsStream<T>
where
T: AsyncRead + AsyncWrite + Unpin,
{
#[inline]
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
match self.get_mut() {
TlsStream::Client(x) => Pin::new(x).poll_write(cx, buf),
TlsStream::Server(x) => Pin::new(x).poll_write(cx, buf),
}
}
#[inline]
fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[io::IoSlice<'_>],
) -> Poll<io::Result<usize>> {
match self.get_mut() {
TlsStream::Client(x) => Pin::new(x).poll_write_vectored(cx, bufs),
TlsStream::Server(x) => Pin::new(x).poll_write_vectored(cx, bufs),
}
}
#[inline]
fn is_write_vectored(&self) -> bool {
match self {
TlsStream::Client(x) => x.is_write_vectored(),
TlsStream::Server(x) => x.is_write_vectored(),
}
}
#[inline]
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match self.get_mut() {
TlsStream::Client(x) => Pin::new(x).poll_flush(cx),
TlsStream::Server(x) => Pin::new(x).poll_flush(cx),
}
}
#[inline]
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match self.get_mut() {
TlsStream::Client(x) => Pin::new(x).poll_shutdown(cx),
TlsStream::Server(x) => Pin::new(x).poll_shutdown(cx),
}
}
}