nexql-conn 0.1.3

//! Shared rustls connector for pooled and one-shot Postgres connections.

use std::fs::File;
use std::io::BufReader;
use std::path::Path;
use std::str::FromStr;

use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls::{ClientConfig, RootCertStore};
use tokio_postgres::Config;
use tokio_postgres_rustls::MakeRustlsConnect;

use crate::error::ConnError;
use crate::resolve::ConnectionParams;

/// Whether the libpq sslmode string requires TLS (not NoTls).
pub fn sslmode_needs_tls(sslmode: Option<&str>) -> bool {
    matches!(
        sslmode.map(str::to_ascii_lowercase).as_deref(),
        Some("require") | Some("verify-ca") | Some("verify-full")
    )
}

/// Whether resolved params require a TLS connector.
pub fn connection_needs_tls(params: &ConnectionParams) -> bool {
    sslmode_needs_tls(params.sslmode.as_deref())
}

/// Parse connection URL into a tokio-postgres config (shared by pool + connect_once).
///
/// tokio-postgres only understands `disable` / `prefer` / `require`; map verify-* to
/// `require` so Config parsing succeeds while rustls still verifies CAs.
pub fn pg_config_from_params(params: &ConnectionParams) -> Result<Config, ConnError> {
    let mut url = params.to_url()?;
    if let Some(mode) = params.sslmode.as_deref() {
        match mode.to_ascii_lowercase().as_str() {
            "verify-ca" | "verify-full" => {
                url = url
                    .replace("sslmode=verify-ca", "sslmode=require")
                    .replace("sslmode=verify-full", "sslmode=require");
            }
            _ => {}
        }
    }
    Config::from_str(&url).map_err(|e| ConnError::Pool(e.to_string()))
}

/// Build a rustls connector for TLS sslmodes (`require`, `verify-ca`, `verify-full`).
pub fn build_rustls_connector(params: &ConnectionParams) -> Result<MakeRustlsConnect, ConnError> {
    if !connection_needs_tls(params) {
        return Err(ConnError::Pool(
            "build_rustls_connector called for non-TLS sslmode".into(),
        ));
    }

    let root_store = load_root_store(params)?;
    let tls_config = if let (Some(cert_path), Some(key_path)) =
        (params.sslcert.as_deref(), params.sslkey.as_deref())
    {
        let certs = load_certs(Path::new(cert_path))?;
        let key = load_private_key(Path::new(key_path))?;
        ClientConfig::builder()
            .with_root_certificates(root_store)
            .with_client_auth_cert(certs, key)
            .map_err(|e| ConnError::Config(format!("client auth TLS config: {e}")))?
    } else if params.sslcert.is_some() || params.sslkey.is_some() {
        return Err(ConnError::Config(
            "sslcert and sslkey must both be set for client certificate auth".into(),
        ));
    } else {
        ClientConfig::builder()
            .with_root_certificates(root_store)
            .with_no_client_auth()
    };

    Ok(MakeRustlsConnect::new(tls_config))
}

fn load_root_store(params: &ConnectionParams) -> Result<RootCertStore, ConnError> {
    let mut root_store = RootCertStore::empty();
    if let Some(root_path) = params.sslrootcert.as_deref() {
        add_pem_roots(Path::new(root_path), &mut root_store)?;
        if root_store.is_empty() {
            return Err(ConnError::Config(format!(
                "sslrootcert {root_path} contained no valid CA certificates"
            )));
        }
        return Ok(root_store);
    }

    let native = rustls_native_certs::load_native_certs();
    for cert in native.certs {
        let _ = root_store.add(cert);
    }
    if root_store.is_empty() {
        return Err(ConnError::Config(
            "no native root certificates available for TLS".into(),
        ));
    }
    Ok(root_store)
}

fn add_pem_roots(path: &Path, store: &mut RootCertStore) -> Result<(), ConnError> {
    let file = File::open(path)
        .map_err(|e| ConnError::Config(format!("sslrootcert {}: {e}", path.display())))?;
    let mut reader = BufReader::new(file);
    for item in rustls_pemfile::certs(&mut reader) {
        let der = item.map_err(|e| ConnError::Config(format!("sslrootcert PEM: {e}")))?;
        store
            .add(der)
            .map_err(|e| ConnError::Config(format!("sslrootcert CA: {e}")))?;
    }
    Ok(())
}

fn load_certs(path: &Path) -> Result<Vec<CertificateDer<'static>>, ConnError> {
    let file = File::open(path)
        .map_err(|e| ConnError::Config(format!("sslcert {}: {e}", path.display())))?;
    let mut reader = BufReader::new(file);
    let certs: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut reader)
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| ConnError::Config(format!("sslcert PEM: {e}")))?;
    if certs.is_empty() {
        return Err(ConnError::Config(format!(
            "sslcert {} contained no certificates",
            path.display()
        )));
    }
    Ok(certs)
}

fn load_private_key(path: &Path) -> Result<PrivateKeyDer<'static>, ConnError> {
    let file = File::open(path)
        .map_err(|e| ConnError::Config(format!("sslkey {}: {e}", path.display())))?;
    let mut reader = BufReader::new(file);
    let key = rustls_pemfile::private_key(&mut reader)
        .map_err(|e| ConnError::Config(format!("sslkey PEM: {e}")))?
        .ok_or_else(|| {
            ConnError::Config(format!(
                "sslkey {} contained no private key",
                path.display()
            ))
        })?;
    Ok(key)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::resolve::params_from_url;

    #[test]
    fn sslmode_needs_tls_mapping() {
        assert!(!sslmode_needs_tls(Some("disable")));
        assert!(!sslmode_needs_tls(Some("prefer")));
        assert!(sslmode_needs_tls(Some("require")));
        assert!(sslmode_needs_tls(Some("verify-ca")));
        assert!(sslmode_needs_tls(Some("verify-full")));
    }

    #[test]
    fn build_rustls_connector_accepts_require_without_custom_ca() {
        let params = params_from_url("postgres://u@h:5432/db?sslmode=require").unwrap();
        assert!(build_rustls_connector(&params).is_ok());
    }

    #[test]
    fn build_rustls_connector_accepts_verify_full_with_native_roots() {
        let params = params_from_url("postgres://u@h:5432/db?sslmode=verify-full").unwrap();
        assert!(build_rustls_connector(&params).is_ok());
    }

    #[test]
    fn pg_config_from_params_accepts_verify_full() {
        let params = params_from_url("postgres://u@h:5432/db?sslmode=verify-full").unwrap();
        assert!(pg_config_from_params(&params).is_ok());
    }
}