tokio-postgres-generic-rustls 0.1.1

//! # tokio-postgres-generic-rustls
//! _An impelementation of TLS based on rustls for tokio-postgres_
//!
//! This crate allows users to select a crypto backend, or bring their own, rather than relying on
//! primitives provided by `ring` directly. This is done through the use of x509-cert for
//! certificate parsing rather than X509-certificate, while also adding an abstraction for
//! computing digests.
//!
//! By default, tokio-postgres-generic-rustls does not provide a digest implementation, but one or
//! more are provided behind crate features.
//!
//! | Feature      | Impelementation    |
//! | ------------ | ------------------ |
//! | `aws-lc-rs`  | `AwsLcRsDigest`    |
//! | `graviola`   | `GraviolaDigest`    |
//! | `ring`       | `RingDigest`       |
//! | `rustcrypto` | `RustcryptoDigest` |
//!
//! ## Usage
//! Using this crate is fairly straightforward. First, select your digest impelementation via crate
//! features (or provide your own), then construct rustls connector for tokio-postgres with your
//! rustls client configuration.
//!
//! The following example demonstrates providing a custom digest backend.
//!
//! ```rust
//! use tokio_postgres_generic_rustls::{DigestImplementation, DigestAlgorithm, MakeRustlsConnect};
//!
//! #[derive(Clone)]
//! struct DemoDigest;
//!
//! impl DigestImplementation for DemoDigest {
//!     fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
//!         todo!("digest it")
//!     }
//! }
//!
//! let cert_store = rustls::RootCertStore::empty();
//!
//! let config = rustls::ClientConfig::builder()
//!     .with_root_certificates(cert_store)
//!     .with_no_client_auth();
//!
//! let tls = MakeRustlsConnect::new(config, DemoDigest);
//!
//! let connect_future = tokio_postgres::connect("postgres://username:password@localhost:5432/db", tls);
//!
//! // connect_future.await;
//! ```
//!
//! ## License
//! This project is licensed under either of
//!
//! - Apache License, Version 2.0, (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
//! - MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT)
//!
//! at your option.

#![deny(unsafe_code)]
#![deny(missing_docs)]

use std::{future::Future, pin::Pin, sync::Arc};

use rustls::{pki_types::ServerName, ClientConfig};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_postgres::tls::{ChannelBinding, MakeTlsConnect, TlsConnect};
use tokio_rustls::{client::TlsStream, Connect, TlsConnector};
use x509_cert::{
    der::{
        oid::db::rfc5912::{
            ECDSA_WITH_SHA_256, ECDSA_WITH_SHA_384, SHA_1_WITH_RSA_ENCRYPTION,
            SHA_256_WITH_RSA_ENCRYPTION, SHA_384_WITH_RSA_ENCRYPTION, SHA_512_WITH_RSA_ENCRYPTION,
        },
        Decode,
    },
    spki::ObjectIdentifier,
    Certificate,
};

/// Trait used to provide a custom digest backend to tokio_postgres_generic_rustls
///
/// This trait is implementated for three types by default, each behind it's own feature flag. The
/// provided backends are aws-lc-rs, ring, and rustcrypto.
pub trait DigestImplementation {
    /// Hash the provided bytes with the provided algorithm
    ///
    /// ```rust
    /// use tokio_postgres_generic_rustls::{DigestImplementation, DigestAlgorithm};
    ///
    /// struct CustomAwsLcRsDigest;
    ///
    /// impl DigestImplementation for CustomAwsLcRsDigest {
    ///     fn digest(&self, digest_algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
    ///         let digest_alg = match digest_algorithm {
    ///             // Note: SHA1 is upgraded to SHA256 as per https://datatracker.ietf.org/doc/html/rfc5929#section-4.1
    ///             DigestAlgorithm::Sha1 | DigestAlgorithm::Sha256 => &aws_lc_rs::digest::SHA256,
    ///             DigestAlgorithm::Sha384 => &aws_lc_rs::digest::SHA384,
    ///             DigestAlgorithm::Sha512 => &aws_lc_rs::digest::SHA512,
    ///         };
    ///
    ///         aws_lc_rs::digest::digest(digest_alg, bytes).as_ref().into()
    ///     }
    /// }
    /// ```
    fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8>;
}

#[cfg(feature = "aws-lc-rs")]
pub use aws_lc_rs_backend::AwsLcRsDigest;

#[cfg(feature = "graviola")]
pub use graviola_backend::GraviolaDigest;

#[cfg(feature = "ring")]
pub use ring_backend::RingDigest;

#[cfg(feature = "rustcrypto")]
pub use rustcrypto_backend::RustcryptoDigest;

#[cfg(feature = "aws-lc-rs")]
mod aws_lc_rs_backend {
    use super::{DigestAlgorithm, DigestImplementation};

    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
    /// A digest backend provided by aws-lc-rs
    ///
    /// Usage:
    /// ```rust
    /// # let cert_store = rustls::RootCertStore::empty();
    /// #
    /// # let rustls_config = rustls::ClientConfig::builder()
    /// #     .with_root_certificates(cert_store)
    /// #     .with_no_client_auth();
    /// #
    /// use tokio_postgres_generic_rustls::{MakeRustlsConnect, AwsLcRsDigest};
    ///
    /// let tls = MakeRustlsConnect::new(rustls_config, AwsLcRsDigest);
    ///
    /// let connect_future = tokio_postgres::connect("postgres://username:password@localhost:5432/db", tls);
    /// ```
    pub struct AwsLcRsDigest;

    impl DigestImplementation for AwsLcRsDigest {
        fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
            let digest_alg = match algorithm {
                // Note: SHA1 is upgraded to SHA256 as per https://datatracker.ietf.org/doc/html/rfc5929#section-4.1
                DigestAlgorithm::Sha1 | DigestAlgorithm::Sha256 => &aws_lc_rs::digest::SHA256,
                DigestAlgorithm::Sha384 => &aws_lc_rs::digest::SHA384,
                DigestAlgorithm::Sha512 => &aws_lc_rs::digest::SHA512,
            };

            aws_lc_rs::digest::digest(digest_alg, bytes).as_ref().into()
        }
    }
}

#[cfg(feature = "graviola")]
mod graviola_backend {
    use super::{DigestAlgorithm, DigestImplementation};

    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
    /// A digest backend provided by aws-lc-rs
    ///
    /// Usage:
    /// ```rust
    /// # let cert_store = rustls::RootCertStore::empty();
    /// #
    /// # let rustls_config = rustls::ClientConfig::builder()
    /// #     .with_root_certificates(cert_store)
    /// #     .with_no_client_auth();
    /// #
    /// use tokio_postgres_generic_rustls::{MakeRustlsConnect, GraviolaDigest};
    ///
    /// let tls = MakeRustlsConnect::new(rustls_config, GraviolaDigest);
    ///
    /// let connect_future = tokio_postgres::connect("postgres://username:password@localhost:5432/db", tls);
    /// ```
    pub struct GraviolaDigest;

    impl DigestImplementation for GraviolaDigest {
        fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
            use graviola::hashing::Hash;

            match algorithm {
                // Note: SHA1 is upgraded to SHA256 as per https://datatracker.ietf.org/doc/html/rfc5929#section-4.1
                DigestAlgorithm::Sha1 | DigestAlgorithm::Sha256 => {
                    graviola::hashing::Sha256::hash(bytes).as_ref().to_vec()
                }
                DigestAlgorithm::Sha384 => graviola::hashing::Sha384::hash(bytes).as_ref().to_vec(),
                DigestAlgorithm::Sha512 => graviola::hashing::Sha512::hash(bytes).as_ref().to_vec(),
            }
        }
    }
}

#[cfg(feature = "ring")]
mod ring_backend {
    use super::{DigestAlgorithm, DigestImplementation};

    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
    /// a digest backend provided by ring
    ///
    /// Usage:
    /// ```rust
    /// # let cert_store = rustls::RootCertStore::empty();
    /// #
    /// # let rustls_config = rustls::ClientConfig::builder()
    /// #     .with_root_certificates(cert_store)
    /// #     .with_no_client_auth();
    /// #
    /// use tokio_postgres_generic_rustls::{MakeRustlsConnect, RingDigest};
    ///
    /// let tls = MakeRustlsConnect::new(rustls_config, RingDigest);
    ///
    /// let connect_future = tokio_postgres::connect("postgres://username:password@localhost:5432/db", tls);
    /// ```
    pub struct RingDigest;

    impl DigestImplementation for RingDigest {
        fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
            let digest_alg = match algorithm {
                // Note: SHA1 is upgraded to SHA256 as per https://datatracker.ietf.org/doc/html/rfc5929#section-4.1
                DigestAlgorithm::Sha1 | DigestAlgorithm::Sha256 => &ring::digest::SHA256,
                DigestAlgorithm::Sha384 => &ring::digest::SHA384,
                DigestAlgorithm::Sha512 => &ring::digest::SHA512,
            };

            ring::digest::digest(digest_alg, bytes).as_ref().into()
        }
    }
}

#[cfg(feature = "rustcrypto")]
mod rustcrypto_backend {
    use super::{DigestAlgorithm, DigestImplementation};
    use sha2::{Digest, Sha256, Sha384, Sha512};

    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
    /// a digest backend provided by rustcrypto
    ///
    /// Usage:
    /// ```rust
    /// # let cert_store = rustls::RootCertStore::empty();
    /// #
    /// # let rustls_config = rustls::ClientConfig::builder()
    /// #     .with_root_certificates(cert_store)
    /// #     .with_no_client_auth();
    /// #
    /// use tokio_postgres_generic_rustls::{MakeRustlsConnect, RustcryptoDigest};
    ///
    /// let tls = MakeRustlsConnect::new(rustls_config, RustcryptoDigest);
    ///
    /// let connect_future = tokio_postgres::connect("postgres://username:password@localhost:5432/db", tls);
    pub struct RustcryptoDigest;

    impl DigestImplementation for RustcryptoDigest {
        fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
            match algorithm {
                // Note: SHA1 is upgraded to SHA256 as per https://datatracker.ietf.org/doc/html/rfc5929#section-4.1
                DigestAlgorithm::Sha1 | DigestAlgorithm::Sha256 => {
                    Sha256::digest(bytes).as_slice().into()
                }
                DigestAlgorithm::Sha384 => Sha384::digest(bytes).as_slice().into(),
                DigestAlgorithm::Sha512 => Sha512::digest(bytes).as_slice().into(),
            }
        }
    }
}

#[derive(Clone)]
/// The primary interface for consumers of this crate
///
/// This type can be provided to tokio-postgres' `connect` method in order to add TLS to a postgres
/// connection.
/// ```rust
/// # use tokio_postgres_generic_rustls::{DigestImplementation, DigestAlgorithm};
/// #
/// # #[derive(Clone)]
/// # struct DemoDigest;
/// #
/// # impl DigestImplementation for DemoDigest {
/// #     fn digest(&self, algorithm: DigestAlgorithm, bytes: &[u8]) -> Vec<u8> {
/// #         todo!("digest it")
/// #     }
/// # }
/// #
/// # let cert_store = rustls::RootCertStore::empty();
/// #
/// # let rustls_config = rustls::ClientConfig::builder()
/// #     .with_root_certificates(cert_store)
/// #     .with_no_client_auth();
/// #
/// use tokio_postgres_generic_rustls::MakeRustlsConnect;
///
/// let tls = MakeRustlsConnect::new(rustls_config, DemoDigest);
///
/// let connect_future = tokio_postgres::connect("postgres://username:password@localhost:5432/db", tls);
/// ```
pub struct MakeRustlsConnect<D> {
    config: Arc<ClientConfig>,
    digest_impl: D,
}

impl<D> MakeRustlsConnect<D>
where
    D: DigestImplementation,
{
    /// Create a new MakeRustlsConnect instance
    pub fn new(config: ClientConfig, digest_impl: D) -> Self {
        Self {
            config: Arc::new(config),
            digest_impl,
        }
    }
}

impl<D, S> MakeTlsConnect<S> for MakeRustlsConnect<D>
where
    D: DigestImplementation + Clone + Unpin,
    S: AsyncRead + AsyncWrite + Unpin + Send,
{
    type Stream = RustlsStream<D, S>;
    type TlsConnect = RustlsConnect<D>;
    type Error = rustls::pki_types::InvalidDnsNameError;

    fn make_tls_connect(&mut self, domain: &str) -> Result<Self::TlsConnect, Self::Error> {
        ServerName::try_from(domain).map(|dns_name| RustlsConnect {
            dns_name: dns_name.to_owned(),
            connector: Arc::clone(&self.config).into(),
            digest_impl: self.digest_impl.clone(),
        })
    }
}

#[doc(hidden)]
pub struct RustlsConnect<D> {
    dns_name: ServerName<'static>,
    connector: TlsConnector,
    digest_impl: D,
}

#[doc(hidden)]
pub struct ConnectFuture<D, S> {
    connect: Connect<S>,
    digest_impl: D,
}

impl<D, S> Future for ConnectFuture<D, S>
where
    D: DigestImplementation + Clone + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    type Output = std::io::Result<RustlsStream<D, S>>;

    fn poll(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        let this = self.get_mut();

        let res = std::task::ready!(Pin::new(&mut this.connect).poll(cx));

        std::task::Poll::Ready(res.map(|io| RustlsStream {
            io,
            digest_impl: this.digest_impl.clone(),
        }))
    }
}

impl<D, S> TlsConnect<S> for RustlsConnect<D>
where
    D: DigestImplementation + Clone + Unpin,
    S: AsyncRead + AsyncWrite + Unpin + Send,
{
    type Stream = RustlsStream<D, S>;
    type Error = std::io::Error;
    type Future = ConnectFuture<D, S>;

    fn connect(self, stream: S) -> Self::Future {
        ConnectFuture {
            connect: self.connector.connect(self.dns_name, stream),
            digest_impl: self.digest_impl.clone(),
        }
    }
}

enum SignatureAlgorithm {
    // 1.2.840.113549.1.1.5
    Sha1Rsa,

    // 1.2.840.113549.1.1.11
    Sha256Rsa,

    // 1.2.840.113549.1.1.12
    Sha384Rsa,

    // 1.2.840.113549.1.1.13
    Sha512Rsa,

    // 1.2.840.10045.4.3.2
    EcdsaSha256,

    // 1.2.840.10045.4.3.3
    EcdsaSha384,
}

impl SignatureAlgorithm {
    fn try_from_identifier(oid: &ObjectIdentifier) -> Option<Self> {
        if oid == &SHA_1_WITH_RSA_ENCRYPTION {
            Some(Self::Sha1Rsa)
        } else if oid == &SHA_256_WITH_RSA_ENCRYPTION {
            Some(Self::Sha256Rsa)
        } else if oid == &SHA_384_WITH_RSA_ENCRYPTION {
            Some(Self::Sha384Rsa)
        } else if oid == &SHA_512_WITH_RSA_ENCRYPTION {
            Some(Self::Sha512Rsa)
        } else if oid == &ECDSA_WITH_SHA_256 {
            Some(Self::EcdsaSha256)
        } else if oid == &ECDSA_WITH_SHA_384 {
            Some(Self::EcdsaSha384)
        } else {
            None
        }
    }

    fn digest_algorithm(self) -> DigestAlgorithm {
        match self {
            Self::Sha1Rsa => DigestAlgorithm::Sha1,
            Self::Sha256Rsa | Self::EcdsaSha256 => DigestAlgorithm::Sha256,
            Self::Sha384Rsa | Self::EcdsaSha384 => DigestAlgorithm::Sha384,
            Self::Sha512Rsa => DigestAlgorithm::Sha512,
        }
    }
}

/// Digest algorithms that can be used in tls-server-end-point channel bindings.
///
/// This type is only useful in the context of defining a custom digest impelementation, and
/// otherwise can be ignored.
pub enum DigestAlgorithm {
    /// The provided certificate requests Sha1 hasing. as per
    /// <https://datatracker.ietf.org/doc/html/rfc5929#section-4.1> sha256 hashing should be used
    /// instead
    Sha1,

    /// The provided certificate requests Sha256 hasing
    Sha256,

    /// The provided certificate requests sha284 hashing
    Sha384,

    /// The provided certificate requests Sha512 hashing
    Sha512,
}

#[doc(hidden)]
pub struct RustlsStream<D, S> {
    io: TlsStream<S>,
    digest_impl: D,
}

impl<D, S> tokio_postgres::tls::TlsStream for RustlsStream<D, S>
where
    D: DigestImplementation + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn channel_binding(&self) -> tokio_postgres::tls::ChannelBinding {
        let (_, session) = self.io.get_ref();

        match session.peer_certificates() {
            Some(certs) if !certs.is_empty() => Certificate::from_der(&certs[0])
                .ok()
                .and_then(|cert| {
                    SignatureAlgorithm::try_from_identifier(&cert.signature_algorithm.oid)
                })
                .map(|signature_algorithm| {
                    let digest = self
                        .digest_impl
                        .digest(signature_algorithm.digest_algorithm(), &certs[0]);

                    ChannelBinding::tls_server_end_point(digest)
                })
                .unwrap_or_else(ChannelBinding::none),
            _ => ChannelBinding::none(),
        }
    }
}

impl<D, S> AsyncRead for RustlsStream<D, S>
where
    D: Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        Pin::new(&mut self.get_mut().io).poll_read(cx, buf)
    }
}

impl<D, S> AsyncWrite for RustlsStream<D, S>
where
    D: Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        Pin::new(&mut self.get_mut().io).poll_write(cx, buf)
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        Pin::new(&mut self.get_mut().io).poll_flush(cx)
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), std::io::Error>> {
        Pin::new(&mut self.get_mut().io).poll_shutdown(cx)
    }

    fn is_write_vectored(&self) -> bool {
        self.io.is_write_vectored()
    }

    fn poll_write_vectored(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        bufs: &[std::io::IoSlice<'_>],
    ) -> std::task::Poll<Result<usize, std::io::Error>> {
        Pin::new(&mut self.get_mut().io).poll_write_vectored(cx, bufs)
    }
}