smartproxy/rust/crates/rustproxy-passthrough/src/tls_handler.rs

use std::io::BufReader;
use std::sync::Arc;

use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls::ServerConfig;
use tokio::net::TcpStream;
use tokio_rustls::{TlsAcceptor, TlsConnector, server::TlsStream as ServerTlsStream};
use tracing::debug;

/// Ensure the default crypto provider is installed.
fn ensure_crypto_provider() {
    let _ = rustls::crypto::ring::default_provider().install_default();
}

/// Build a TLS acceptor from PEM-encoded cert and key data.
pub fn build_tls_acceptor(cert_pem: &str, key_pem: &str) -> Result<TlsAcceptor, Box<dyn std::error::Error + Send + Sync>> {
    build_tls_acceptor_with_config(cert_pem, key_pem, None)
}

/// Build a TLS acceptor with optional RouteTls configuration for version/cipher tuning.
pub fn build_tls_acceptor_with_config(
    cert_pem: &str,
    key_pem: &str,
    tls_config: Option<&rustproxy_config::RouteTls>,
) -> Result<TlsAcceptor, Box<dyn std::error::Error + Send + Sync>> {
    ensure_crypto_provider();
    let certs = load_certs(cert_pem)?;
    let key = load_private_key(key_pem)?;

    let mut config = if let Some(route_tls) = tls_config {
        // Apply TLS version restrictions
        let versions = resolve_tls_versions(route_tls.versions.as_deref());
        let builder = ServerConfig::builder_with_protocol_versions(&versions);
        builder
            .with_no_client_auth()
            .with_single_cert(certs, key)?
    } else {
        ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(certs, key)?
    };

    // Apply session timeout if configured
    if let Some(route_tls) = tls_config {
        if let Some(timeout_secs) = route_tls.session_timeout {
            config.session_storage = rustls::server::ServerSessionMemoryCache::new(
                256, // max sessions
            );
            debug!("TLS session timeout configured: {}s", timeout_secs);
        }
    }

    Ok(TlsAcceptor::from(Arc::new(config)))
}

/// Resolve TLS version strings to rustls SupportedProtocolVersion.
fn resolve_tls_versions(versions: Option<&[String]>) -> Vec<&'static rustls::SupportedProtocolVersion> {
    let versions = match versions {
        Some(v) if !v.is_empty() => v,
        _ => return vec![&rustls::version::TLS12, &rustls::version::TLS13],
    };

    let mut result = Vec::new();
    for v in versions {
        match v.as_str() {
            "TLSv1.2" | "TLS1.2" | "1.2" | "TLSv12" => {
                if !result.contains(&&rustls::version::TLS12) {
                    result.push(&rustls::version::TLS12);
                }
            }
            "TLSv1.3" | "TLS1.3" | "1.3" | "TLSv13" => {
                if !result.contains(&&rustls::version::TLS13) {
                    result.push(&rustls::version::TLS13);
                }
            }
            other => {
                debug!("Unknown TLS version '{}', ignoring", other);
            }
        }
    }

    if result.is_empty() {
        // Fallback to both if no valid versions specified
        vec![&rustls::version::TLS12, &rustls::version::TLS13]
    } else {
        result
    }
}

/// Accept a TLS connection from a client stream.
pub async fn accept_tls(
    stream: TcpStream,
    acceptor: &TlsAcceptor,
) -> Result<ServerTlsStream<TcpStream>, Box<dyn std::error::Error + Send + Sync>> {
    let tls_stream = acceptor.accept(stream).await?;
    debug!("TLS handshake completed");
    Ok(tls_stream)
}

/// Connect to a backend with TLS (for terminate-and-reencrypt mode).
pub async fn connect_tls(
    host: &str,
    port: u16,
) -> Result<tokio_rustls::client::TlsStream<TcpStream>, Box<dyn std::error::Error + Send + Sync>> {
    ensure_crypto_provider();
    let config = rustls::ClientConfig::builder()
        .dangerous()
        .with_custom_certificate_verifier(Arc::new(InsecureVerifier))
        .with_no_client_auth();

    let connector = TlsConnector::from(Arc::new(config));

    let stream = TcpStream::connect(format!("{}:{}", host, port)).await?;
    stream.set_nodelay(true)?;

    let server_name = rustls::pki_types::ServerName::try_from(host.to_string())?;
    let tls_stream = connector.connect(server_name, stream).await?;
    debug!("Backend TLS connection established to {}:{}", host, port);
    Ok(tls_stream)
}

/// Load certificates from PEM string.
fn load_certs(pem: &str) -> Result<Vec<CertificateDer<'static>>, Box<dyn std::error::Error + Send + Sync>> {
    let mut reader = BufReader::new(pem.as_bytes());
    let certs: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut reader)
        .collect::<Result<Vec<_>, _>>()?;
    if certs.is_empty() {
        return Err("No certificates found in PEM data".into());
    }
    Ok(certs)
}

/// Load private key from PEM string.
fn load_private_key(pem: &str) -> Result<PrivateKeyDer<'static>, Box<dyn std::error::Error + Send + Sync>> {
    let mut reader = BufReader::new(pem.as_bytes());
    // Try PKCS8 first, then RSA, then EC
    let key = rustls_pemfile::private_key(&mut reader)?
        .ok_or("No private key found in PEM data")?;
    Ok(key)
}

/// Insecure certificate verifier for backend connections (terminate-and-reencrypt).
/// In internal networks, backends may use self-signed certs.
#[derive(Debug)]
struct InsecureVerifier;

impl rustls::client::danger::ServerCertVerifier for InsecureVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &rustls::pki_types::ServerName<'_>,
        _ocsp_response: &[u8],
        _now: rustls::pki_types::UnixTime,
    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::danger::ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        vec![
            rustls::SignatureScheme::RSA_PKCS1_SHA256,
            rustls::SignatureScheme::RSA_PKCS1_SHA384,
            rustls::SignatureScheme::RSA_PKCS1_SHA512,
            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
            rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
            rustls::SignatureScheme::ED25519,
            rustls::SignatureScheme::RSA_PSS_SHA256,
            rustls::SignatureScheme::RSA_PSS_SHA384,
            rustls::SignatureScheme::RSA_PSS_SHA512,
        ]
    }
}