use std::{
convert::TryFrom,
future::Future,
io,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use DigestAlgorithm::{Sha1, Sha256, Sha384, Sha512};
use SignatureAlgorithm::{
EcdsaSha256, EcdsaSha384, Ed25519, NoSignature, RsaSha1, RsaSha256, RsaSha384, RsaSha512,
};
use ring::digest;
use rustls::{ClientConfig, pki_types::ServerName};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio_postgres::tls::{ChannelBinding, MakeTlsConnect, TlsConnect};
use tokio_rustls::{TlsConnector, client::TlsStream};
use x509_certificate::{DigestAlgorithm, SignatureAlgorithm, X509Certificate};
#[derive(Clone)]
pub struct MakeRustlsConnect {
pub config: Arc<ClientConfig>,
}
impl MakeRustlsConnect {
pub fn new(config: impl Into<Arc<ClientConfig>>) -> Self {
Self {
config: config.into(),
}
}
}
impl<S> MakeTlsConnect<S> for MakeRustlsConnect
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
type Stream = RustlsStream<S>;
type TlsConnect = RustlsConnect;
type Error = rustls::pki_types::InvalidDnsNameError;
fn make_tls_connect(&mut self, hostname: &str) -> Result<RustlsConnect, Self::Error> {
ServerName::try_from(hostname).map(|dns_name| {
RustlsConnect(RustlsConnectData {
hostname: dns_name.to_owned(),
connector: Arc::clone(&self.config).into(),
})
})
}
}
pub struct RustlsConnect(RustlsConnectData);
struct RustlsConnectData {
hostname: ServerName<'static>,
connector: TlsConnector,
}
impl<S> TlsConnect<S> for RustlsConnect
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
type Stream = RustlsStream<S>;
type Error = io::Error;
type Future = Pin<Box<dyn Future<Output = io::Result<RustlsStream<S>>> + Send>>;
fn connect(self, stream: S) -> Self::Future {
Box::pin(async move {
self
.0
.connector
.connect(self.0.hostname, stream)
.await
.map(|s| RustlsStream(Box::pin(s)))
})
}
}
pub struct RustlsStream<S>(Pin<Box<TlsStream<S>>>);
impl<S> tokio_postgres::tls::TlsStream for RustlsStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn channel_binding(&self) -> ChannelBinding {
let (_, session) = self.0.get_ref();
match session.peer_certificates() {
Some(certs) if !certs.is_empty() => X509Certificate::from_der(&certs[0])
.ok()
.and_then(|cert| cert.signature_algorithm())
.map(|algorithm| match algorithm {
RsaSha1 | RsaSha256 | EcdsaSha256 => &digest::SHA256,
RsaSha384 | EcdsaSha384 => &digest::SHA384,
RsaSha512 => &digest::SHA512,
Ed25519 => &digest::SHA512,
NoSignature(algo) => match algo {
Sha1 | Sha256 => &digest::SHA256,
Sha384 => &digest::SHA384,
Sha512 => &digest::SHA512,
},
})
.map(|algorithm| {
let hash = digest::digest(algorithm, certs[0].as_ref());
ChannelBinding::tls_server_end_point(hash.as_ref().into())
})
.unwrap_or(ChannelBinding::none()),
_ => ChannelBinding::none(),
}
}
}
impl<S> AsyncRead for RustlsStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context,
buf: &mut ReadBuf<'_>,
) -> Poll<tokio::io::Result<()>> {
self.0.as_mut().poll_read(cx, buf)
}
}
impl<S> AsyncWrite for RustlsStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context,
buf: &[u8],
) -> Poll<tokio::io::Result<usize>> {
self.0.as_mut().poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<tokio::io::Result<()>> {
self.0.as_mut().poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<tokio::io::Result<()>> {
self.0.as_mut().poll_shutdown(cx)
}
}