use futures_util::future;
use std::io::{self, BufReader, Cursor, Read, Write};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::SystemTime;
use rustls::{
client::{ServerCertVerified, ServerCertVerifier, WebPkiVerifier},
CertificateError, ClientConfig, ClientConnection, Error as TlsError, OwnedTrustAnchor,
RootCertStore, ServerName,
};
use crate::error::Error;
use crate::io::ReadBuf;
use crate::net::tls::util::StdSocket;
use crate::net::tls::TlsConfig;
use crate::net::Socket;
pub struct RustlsSocket<S: Socket> {
inner: StdSocket<S>,
state: ClientConnection,
close_notify_sent: bool,
}
impl<S: Socket> RustlsSocket<S> {
fn poll_complete_io(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
loop {
match self.state.complete_io(&mut self.inner) {
Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
futures_util::ready!(self.inner.poll_ready(cx))?;
}
ready => return Poll::Ready(ready.map(|_| ())),
}
}
}
async fn complete_io(&mut self) -> io::Result<()> {
future::poll_fn(|cx| self.poll_complete_io(cx)).await
}
}
impl<S: Socket> Socket for RustlsSocket<S> {
fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> {
self.state.reader().read(buf.init_mut())
}
fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> {
match self.state.writer().write(buf) {
Ok(0) => Err(io::ErrorKind::WouldBlock.into()),
other => other,
}
}
fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.poll_complete_io(cx)
}
fn poll_write_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.poll_complete_io(cx)
}
fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.poll_complete_io(cx)
}
fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
if !self.close_notify_sent {
self.state.send_close_notify();
self.close_notify_sent = true;
}
futures_util::ready!(self.poll_complete_io(cx))?;
self.inner.socket.poll_shutdown(cx)
}
}
pub async fn handshake<S>(socket: S, tls_config: TlsConfig<'_>) -> Result<RustlsSocket<S>, Error>
where
S: Socket,
{
let config = ClientConfig::builder().with_safe_defaults();
let user_auth = match (tls_config.client_cert_path, tls_config.client_key_path) {
(Some(cert_path), Some(key_path)) => {
let cert_chain = certs_from_pem(cert_path.data().await?)?;
let key_der = private_key_from_pem(key_path.data().await?)?;
Some((cert_chain, key_der))
}
(None, None) => None,
(_, _) => {
return Err(Error::Configuration(
"user auth key and certs must be given together".into(),
))
}
};
let config = if tls_config.accept_invalid_certs {
if let Some(user_auth) = user_auth {
config
.with_custom_certificate_verifier(Arc::new(DummyTlsVerifier))
.with_client_auth_cert(user_auth.0, user_auth.1)
.map_err(Error::tls)?
} else {
config
.with_custom_certificate_verifier(Arc::new(DummyTlsVerifier))
.with_no_client_auth()
}
} else {
let mut cert_store = RootCertStore::empty();
cert_store.add_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.iter().map(|ta| {
OwnedTrustAnchor::from_subject_spki_name_constraints(
ta.subject,
ta.spki,
ta.name_constraints,
)
}));
if let Some(ca) = tls_config.root_cert_path {
let data = ca.data().await?;
let mut cursor = Cursor::new(data);
for cert in rustls_pemfile::certs(&mut cursor)
.map_err(|_| Error::Tls(format!("Invalid certificate {ca}").into()))?
{
cert_store
.add(&rustls::Certificate(cert))
.map_err(|err| Error::Tls(err.into()))?;
}
}
if tls_config.accept_invalid_hostnames {
let verifier = WebPkiVerifier::new(cert_store, None);
if let Some(user_auth) = user_auth {
config
.with_custom_certificate_verifier(Arc::new(NoHostnameTlsVerifier { verifier }))
.with_client_auth_cert(user_auth.0, user_auth.1)
.map_err(Error::tls)?
} else {
config
.with_custom_certificate_verifier(Arc::new(NoHostnameTlsVerifier { verifier }))
.with_no_client_auth()
}
} else if let Some(user_auth) = user_auth {
config
.with_root_certificates(cert_store)
.with_client_auth_cert(user_auth.0, user_auth.1)
.map_err(Error::tls)?
} else {
config
.with_root_certificates(cert_store)
.with_no_client_auth()
}
};
let host = rustls::ServerName::try_from(tls_config.hostname).map_err(Error::tls)?;
let mut socket = RustlsSocket {
inner: StdSocket::new(socket),
state: ClientConnection::new(Arc::new(config), host).map_err(Error::tls)?,
close_notify_sent: false,
};
socket.complete_io().await?;
Ok(socket)
}
fn certs_from_pem(pem: Vec<u8>) -> Result<Vec<rustls::Certificate>, Error> {
let cur = Cursor::new(pem);
let mut reader = BufReader::new(cur);
rustls_pemfile::certs(&mut reader)?
.into_iter()
.map(|v| Ok(rustls::Certificate(v)))
.collect()
}
fn private_key_from_pem(pem: Vec<u8>) -> Result<rustls::PrivateKey, Error> {
let cur = Cursor::new(pem);
let mut reader = BufReader::new(cur);
loop {
match rustls_pemfile::read_one(&mut reader)? {
Some(
rustls_pemfile::Item::RSAKey(key)
| rustls_pemfile::Item::PKCS8Key(key)
| rustls_pemfile::Item::ECKey(key),
) => return Ok(rustls::PrivateKey(key)),
None => break,
_ => {}
}
}
Err(Error::Configuration("no keys found pem file".into()))
}
struct DummyTlsVerifier;
impl ServerCertVerifier for DummyTlsVerifier {
fn verify_server_cert(
&self,
_end_entity: &rustls::Certificate,
_intermediates: &[rustls::Certificate],
_server_name: &ServerName,
_scts: &mut dyn Iterator<Item = &[u8]>,
_ocsp_response: &[u8],
_now: SystemTime,
) -> Result<ServerCertVerified, TlsError> {
Ok(ServerCertVerified::assertion())
}
}
pub struct NoHostnameTlsVerifier {
verifier: WebPkiVerifier,
}
impl ServerCertVerifier for NoHostnameTlsVerifier {
fn verify_server_cert(
&self,
end_entity: &rustls::Certificate,
intermediates: &[rustls::Certificate],
server_name: &ServerName,
scts: &mut dyn Iterator<Item = &[u8]>,
ocsp_response: &[u8],
now: SystemTime,
) -> Result<ServerCertVerified, TlsError> {
match self.verifier.verify_server_cert(
end_entity,
intermediates,
server_name,
scts,
ocsp_response,
now,
) {
Err(TlsError::InvalidCertificate(reason))
if reason == CertificateError::NotValidForName =>
{
Ok(ServerCertVerified::assertion())
}
res => res,
}
}
}