network-protocol 1.2.4

//! # TLS Transport Layer
//!
//! This file is part of the Network Protocol project.
//!
//! It defines the TLS transport layer for secure network communication,
//! particularly for external untrusted connections.
//!
//! The TLS transport layer provides a secure channel for communication
//! using industry-standard TLS protocol, ensuring confidentiality,
//! integrity, and authentication of the data transmitted.
//!
//! ## Responsibilities
//! - Establish secure TLS connections
//! - Handle TLS certificates and verification
//! - Provide secure framed transport for higher protocol layers
//! - Compatible with existing packet codec infrastructure

use std::fs::File;
use std::io::{self, BufReader, Seek, Write};
use std::net::SocketAddr;
use std::path::Path;
use std::sync::Arc;

use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::crypto::{CryptoProvider, WebPkiSupportedAlgorithms};
use rustls::pki_types::pem::PemObject;
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName, UnixTime};
use rustls::{ClientConfig, DigitallySignedStruct, RootCertStore, ServerConfig};
use tokio::net::{TcpListener, TcpStream};
use tokio_rustls::client::TlsStream as ClientTlsStream;
use tokio_rustls::server::TlsStream as ServerTlsStream;
use tokio_rustls::{TlsAcceptor, TlsConnector};
use tokio_util::codec::Framed;
use tracing::{debug, error, info, instrument, warn};

use crate::core::codec::PacketCodec;
use crate::core::packet::Packet;
use crate::error::{ProtocolError, Result};
use futures::{SinkExt, StreamExt};

// Custom certificate verifiers.
//
// Both verifiers replace only the certificate *validation* step (chain building and
// hostname checks). They still verify the handshake signature with the provider's
// signature algorithms, which is what proves that the server holds the private key
// for the certificate it presented. Without that check a pinned certificate, which is
// public, could be replayed by any server.

/// Accepts only a server certificate whose SHA-256 fingerprint matches the pin.
#[derive(Debug)]
struct CertificateFingerprint {
    fingerprint: Vec<u8>,
    supported_algs: WebPkiSupportedAlgorithms,
}

impl ServerCertVerifier for CertificateFingerprint {
    fn verify_server_cert(
        &self,
        end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
        use sha2::{Digest, Sha256};

        let mut hasher = Sha256::new();
        hasher.update(end_entity);
        let hash = hasher.finalize();

        if hash.as_slice() == self.fingerprint.as_slice() {
            Ok(ServerCertVerified::assertion())
        } else {
            Err(rustls::Error::General(
                "Pinned certificate hash mismatch".into(),
            ))
        }
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls12_signature(message, cert, dss, &self.supported_algs)
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls13_signature(message, cert, dss, &self.supported_algs)
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        self.supported_algs.supported_schemes()
    }
}

/// Accepts any server certificate (insecure mode without a pin), but still requires a
/// valid handshake signature from the key in that certificate.
#[derive(Debug)]
struct AcceptAnyServerCert {
    supported_algs: WebPkiSupportedAlgorithms,
}

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

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls12_signature(message, cert, dss, &self.supported_algs)
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls13_signature(message, cert, dss, &self.supported_algs)
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        self.supported_algs.supported_schemes()
    }
}

/// The crypto provider used by every TLS config and verifier in this module.
///
/// rustls APIs without a `_with_provider` suffix fall back to the process-level default
/// provider, which panics when more than one rustls crypto backend is compiled in and
/// none has been installed. Always pass this provider explicitly instead.
fn crypto_provider() -> Arc<CryptoProvider> {
    Arc::new(rustls::crypto::ring::default_provider())
}

/// Resolves the versions and cipher suites requested with `with_tls_versions()` and
/// `with_cipher_suites()` into the protocol versions to enable and a provider whose
/// cipher suites are restricted to the requested ones.
///
/// `None` means the rustls defaults. Cipher suites are matched by their IANA identifier
/// against the suites of [`crypto_provider()`], keeping that provider's preference order.
/// Requests that leave nothing usable are rejected here rather than falling back to the
/// defaults.
fn resolve_protocol_settings(
    tls_versions: Option<&[TlsVersion]>,
    cipher_suites: Option<&[rustls::SupportedCipherSuite]>,
) -> Result<(
    Arc<CryptoProvider>,
    Vec<&'static rustls::SupportedProtocolVersion>,
)> {
    let versions: Vec<&'static rustls::SupportedProtocolVersion> = match tls_versions {
        None => rustls::DEFAULT_VERSIONS.to_vec(),
        Some(requested) => {
            let mut has_tls13 = false;
            let mut has_tls12 = false;
            for v in requested {
                match v {
                    TlsVersion::TLS12 => has_tls12 = true,
                    TlsVersion::TLS13 => has_tls13 = true,
                    TlsVersion::All => {
                        has_tls13 = true;
                        has_tls12 = true;
                    }
                }
            }
            debug!(
                "TLS versions requested: TLS1.2={}, TLS1.3={}",
                has_tls12, has_tls13
            );

            let mut versions = Vec::with_capacity(2);
            if has_tls13 {
                versions.push(&rustls::version::TLS13);
            }
            if has_tls12 {
                versions.push(&rustls::version::TLS12);
            }
            if versions.is_empty() {
                return Err(ProtocolError::TlsError(
                    "No TLS protocol versions configured: with_tls_versions() was given an empty list"
                        .into(),
                ));
            }
            versions
        }
    };

    let mut provider = CryptoProvider::clone(&crypto_provider());
    if let Some(requested) = cipher_suites {
        for suite in requested {
            if !provider
                .cipher_suites
                .iter()
                .any(|s| s.suite() == suite.suite())
            {
                warn!(suite = ?suite.suite(), "Requested cipher suite is not supported and was ignored");
            }
        }
        provider
            .cipher_suites
            .retain(|s| requested.iter().any(|r| r.suite() == s.suite()));
        if provider.cipher_suites.is_empty() {
            return Err(ProtocolError::TlsError(
                "No supported cipher suites configured: none of the suites given to with_cipher_suites() are available"
                    .into(),
            ));
        }
        debug!(
            suites = ?provider.cipher_suites.iter().map(|s| s.suite()).collect::<Vec<_>>(),
            "TLS cipher suites restricted"
        );
    }

    if !provider
        .cipher_suites
        .iter()
        .any(|s| versions.contains(&s.version()))
    {
        return Err(ProtocolError::TlsError(
            "None of the configured cipher suites can be used with the configured TLS versions"
                .into(),
        ));
    }

    Ok((Arc::new(provider), versions))
}

/// Helper function to load a private key from PKCS8 format
fn load_private_key(reader: &mut BufReader<File>) -> Result<PrivateKeyDer<'static>> {
    // Try to load PKCS8 keys
    // Seek to beginning of file first
    reader
        .seek(std::io::SeekFrom::Start(0))
        .map_err(ProtocolError::Io)?;

    // Collect every PKCS8 private key section in the PEM input
    let keys: std::result::Result<Vec<_>, _> =
        PrivatePkcs8KeyDer::pem_reader_iter(reader).collect();
    let keys =
        keys.map_err(|_| ProtocolError::TlsError("Failed to parse PKCS8 private key".into()))?;

    if !keys.is_empty() {
        return Ok(PrivateKeyDer::Pkcs8(keys[0].clone_key()));
    }

    // Note: Add support for other key formats like RSA or EC if needed

    Err(ProtocolError::TlsError(
        "No supported private key format found".into(),
    ))
}

/// TLS protocol version
pub enum TlsVersion {
    /// TLS 1.2
    TLS12,
    /// TLS 1.3
    TLS13,
    /// Both TLS 1.2 and 1.3
    All,
}

/// TLS server configuration
pub struct TlsServerConfig {
    cert_path: String,
    key_path: String,
    /// Optional path to client CA certificates for mTLS
    client_ca_path: Option<String>,
    /// Whether to require client certificates (mTLS)
    require_client_auth: bool,
    /// Allowed TLS protocol versions (None = use rustls defaults)
    tls_versions: Option<Vec<TlsVersion>>,
    /// Allowed cipher suites (None = use rustls defaults)
    cipher_suites: Option<Vec<rustls::SupportedCipherSuite>>,
    /// ALPN protocols to advertise (default: ["h2", "http/1.1"])
    alpn_protocols: Option<Vec<Vec<u8>>>,
}

impl TlsServerConfig {
    /// Create a new TLS server configuration
    pub fn new<P: AsRef<std::path::Path>>(cert_path: P, key_path: P) -> Self {
        Self {
            cert_path: cert_path.as_ref().to_string_lossy().to_string(),
            key_path: key_path.as_ref().to_string_lossy().to_string(),
            client_ca_path: None,
            require_client_auth: false,
            tls_versions: None,
            cipher_suites: None,
            alpn_protocols: Some(vec![b"h2".to_vec(), b"http/1.1".to_vec()]),
        }
    }

    /// Set allowed TLS protocol versions
    ///
    /// Only the listed versions are enabled. An empty list is rejected when the
    /// config is loaded.
    pub fn with_tls_versions(mut self, versions: Vec<TlsVersion>) -> Self {
        self.tls_versions = Some(versions);
        self
    }

    /// Set allowed cipher suites
    ///
    /// Only the listed suites are enabled, matched by suite identifier and kept in this
    /// crate's preference order. Loading the config fails if none of them is supported,
    /// or if none of them can be used with the enabled TLS versions.
    pub fn with_cipher_suites(mut self, cipher_suites: Vec<rustls::SupportedCipherSuite>) -> Self {
        self.cipher_suites = Some(cipher_suites);
        self
    }

    /// Enable mutual TLS authentication by providing a CA certificate path
    pub fn with_client_auth<S: Into<String>>(mut self, client_ca_path: S) -> Self {
        self.client_ca_path = Some(client_ca_path.into());
        self.require_client_auth = true;
        self
    }

    /// Set whether client authentication is required (true) or optional (false)
    ///
    /// Only has an effect together with `with_client_auth()`, which sets it to `true`.
    /// When optional, clients without a certificate are accepted, but a certificate that
    /// is presented must still be signed by the client CA.
    pub fn require_client_auth(mut self, required: bool) -> Self {
        self.require_client_auth = required;
        self
    }

    /// Set ALPN protocols to advertise during TLS handshake
    pub fn with_alpn_protocols(mut self, protocols: Vec<Vec<u8>>) -> Self {
        self.alpn_protocols = Some(protocols);
        self
    }

    /// Generate a self-signed certificate for development/testing purposes
    pub fn generate_self_signed<P: AsRef<Path>>(cert_path: P, key_path: P) -> io::Result<Self> {
        let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()])
            .map_err(|e| io::Error::other(format!("Certificate generation error: {e}")))?;

        // Write certificate
        let mut cert_file = File::create(&cert_path)?;
        let pem = cert.cert.pem();
        cert_file.write_all(pem.as_bytes())?;

        // Write private key
        let mut key_file = File::create(&key_path)?;
        key_file.write_all(cert.signing_key.serialize_pem().as_bytes())?;

        Ok(Self {
            cert_path: cert_path.as_ref().to_string_lossy().to_string(),
            key_path: key_path.as_ref().to_string_lossy().to_string(),
            client_ca_path: None,
            require_client_auth: false,
            tls_versions: None,
            cipher_suites: None,
            alpn_protocols: Some(vec![b"h2".to_vec(), b"http/1.1".to_vec()]),
        })
    }

    /// Load the TLS configuration from files
    pub fn load_server_config(&self) -> Result<ServerConfig> {
        // Load certificate
        let cert_file = File::open(&self.cert_path)
            .map_err(|e| ProtocolError::TlsError(format!("Failed to open cert file: {e}")))?;
        let mut cert_reader = BufReader::new(cert_file);
        let cert_chain: std::result::Result<Vec<_>, _> =
            CertificateDer::pem_reader_iter(&mut cert_reader).collect();
        let cert_chain: Vec<CertificateDer<'static>> = cert_chain
            .map_err(|_| ProtocolError::TlsError("Failed to parse certificate".into()))?;

        if cert_chain.is_empty() {
            return Err(ProtocolError::TlsError("No certificates found".into()));
        }

        // Load private key
        let key_file = File::open(&self.key_path)
            .map_err(|e| ProtocolError::TlsError(format!("Failed to open key file: {e}")))?;
        let mut key_reader = BufReader::new(key_file);
        let private_key = load_private_key(&mut key_reader)?;

        let (provider, versions) =
            resolve_protocol_settings(self.tls_versions.as_deref(), self.cipher_suites.as_deref())?;
        let config_builder = ServerConfig::builder_with_provider(provider.clone())
            .with_protocol_versions(&versions)
            .map_err(|e| {
                ProtocolError::TlsError(format!("Failed to configure TLS protocol versions: {e}"))
            })?;

        // Configure client authentication (mTLS) if a client CA was given
        let cert_builder = if let Some(client_ca_path) = &self.client_ca_path {
            // Load client CA certificates
            let client_ca_file = File::open(client_ca_path).map_err(|e| {
                ProtocolError::TlsError(format!("Failed to open client CA file: {e}"))
            })?;
            let mut client_ca_reader = BufReader::new(client_ca_file);
            let client_ca_certs: std::result::Result<Vec<_>, _> =
                CertificateDer::pem_reader_iter(&mut client_ca_reader).collect();
            let client_ca_certs: Vec<CertificateDer<'static>> = client_ca_certs.map_err(|_| {
                ProtocolError::TlsError("Failed to parse client CA certificate".into())
            })?;

            if client_ca_certs.is_empty() {
                return Err(ProtocolError::TlsError(
                    "No client CA certificates found".into(),
                ));
            }

            // Create client cert verifier
            let mut client_root_store = RootCertStore::empty();
            for cert in client_ca_certs {
                client_root_store.add(cert).map_err(|e| {
                    ProtocolError::TlsError(format!("Failed to add client CA cert: {e}"))
                })?;
            }

            // Create client authentication verifier using WebPkiClientVerifier.
            // The provider is passed explicitly: `WebPkiClientVerifier::builder` uses the
            // process-level default, which panics when it cannot be chosen automatically.
            let mut verifier_builder = rustls::server::WebPkiClientVerifier::builder_with_provider(
                Arc::new(client_root_store),
                provider,
            );
            if !self.require_client_auth {
                // Optional client auth: a client without a certificate is accepted, but a
                // certificate that is presented must still chain to the client CA.
                verifier_builder = verifier_builder.allow_unauthenticated();
            }
            let client_auth = verifier_builder.build().map_err(|e| {
                ProtocolError::TlsError(format!("Failed to build client verifier: {e}"))
            })?;

            if self.require_client_auth {
                debug!("mTLS enabled with client certificate verification required");
            } else {
                debug!("mTLS enabled with optional client certificate verification");
            }
            config_builder.with_client_cert_verifier(client_auth)
        } else {
            config_builder.with_no_client_auth()
        };

        // Build config with certificates
        let mut config = cert_builder
            .with_single_cert(cert_chain, private_key)
            .map_err(|e| ProtocolError::TlsError(format!("TLS error: {e}")))?;

        // Configure ALPN protocols if specified
        if let Some(protocols) = &self.alpn_protocols {
            config.alpn_protocols = protocols.clone();
            debug!(
                protocol_count = protocols.len(),
                "ALPN protocols configured"
            );
        }

        Ok(config)
    }

    /// Calculate SHA-256 hash for a certificate to use with pinning
    pub fn calculate_cert_hash(cert: &CertificateDer<'_>) -> Vec<u8> {
        use sha2::{Digest, Sha256};
        let mut hasher = Sha256::new();
        hasher.update(cert.as_ref());
        hasher.finalize().to_vec()
    }
}

/// TLS Client Configuration
pub struct TlsClientConfig {
    server_name: String,
    insecure: bool,
    /// Optional certificate hash to pin (SHA-256 fingerprint)
    pinned_cert_hash: Option<Vec<u8>>,
    /// Optional client certificate path for mTLS
    client_cert_path: Option<String>,
    /// Optional client key path for mTLS
    client_key_path: Option<String>,
    /// Allowed TLS protocol versions (None = use rustls defaults)
    tls_versions: Option<Vec<TlsVersion>>,
    /// Allowed cipher suites (None = use rustls defaults)
    cipher_suites: Option<Vec<rustls::SupportedCipherSuite>>,
}

impl TlsClientConfig {
    /// Create a new TLS client configuration
    pub fn new<S: Into<String>>(server_name: S) -> Self {
        Self {
            server_name: server_name.into(),
            insecure: false,
            pinned_cert_hash: None,
            client_cert_path: None,
            client_key_path: None,
            tls_versions: None,
            cipher_suites: None,
        }
    }

    /// Set allowed TLS protocol versions
    ///
    /// Only the listed versions are enabled. An empty list is rejected when the
    /// config is loaded.
    pub fn with_tls_versions(mut self, versions: Vec<TlsVersion>) -> Self {
        self.tls_versions = Some(versions);
        self
    }

    /// Set allowed cipher suites
    ///
    /// Only the listed suites are enabled, matched by suite identifier and kept in this
    /// crate's preference order. Loading the config fails if none of them is supported,
    /// or if none of them can be used with the enabled TLS versions.
    pub fn with_cipher_suites(mut self, cipher_suites: Vec<rustls::SupportedCipherSuite>) -> Self {
        self.cipher_suites = Some(cipher_suites);
        self
    }

    /// Configure client authentication for mTLS
    pub fn with_client_certificate<S: Into<String>>(mut self, cert_path: S, key_path: S) -> Self {
        self.client_cert_path = Some(cert_path.into());
        self.client_key_path = Some(key_path.into());
        self
    }

    /// Allow insecure connections (skip certificate verification)
    ///
    /// # WARNING: Security Risk
    /// This mode skips certificate validation (no CA chain or hostname checks), so any
    /// certificate is accepted unless one is pinned with `with_pinned_cert_hash()`. The
    /// server must still prove that it holds the private key for the certificate it
    /// presents: the handshake signature is always verified. This mode should ONLY be used for:
    /// - Development and testing
    /// - Debugging environments
    /// - Internal networks with certificate pinning enabled
    ///
    /// **NEVER** use this in production without explicit certificate pinning via `with_pinned_cert_hash()`.
    ///
    /// For maximum security, only use this with the "dangerous_configuration" feature enabled,
    /// which is a strong indicator this is for testing/development only.
    pub fn insecure(mut self) -> Self {
        warn!("INSECURE MODE ENABLED: Certificate verification is disabled. This should only be used for development/testing.");
        self.insecure = true;
        self
    }

    /// Pin a certificate by its SHA-256 hash/fingerprint
    ///
    /// This provides additional security by only accepting connections
    /// from servers with the exact certificate matching this hash.
    /// Can be combined with insecure mode for development environments where
    /// you want to skip standard CA verification but still verify a specific cert.
    /// The handshake signature is verified against the pinned certificate's public key,
    /// so a server that has the certificate but not its private key is rejected.
    pub fn with_pinned_cert_hash(mut self, hash: Vec<u8>) -> Self {
        if hash.len() != 32 {
            warn!(
                "Certificate hash has unexpected length: {} (expected 32 bytes for SHA-256)",
                hash.len()
            );
        }
        self.pinned_cert_hash = Some(hash);
        self
    }

    /// Load the TLS client configuration
    pub fn load_client_config(&self) -> Result<ClientConfig> {
        if self.insecure {
            self.build_insecure_client_config()
        } else {
            self.build_secure_client_config()
        }
    }

    /// Start a client config builder with the configured provider, versions and suites
    fn client_config_builder(
        &self,
    ) -> Result<(
        Arc<CryptoProvider>,
        rustls::ConfigBuilder<ClientConfig, rustls::WantsVerifier>,
    )> {
        let (provider, versions) =
            resolve_protocol_settings(self.tls_versions.as_deref(), self.cipher_suites.as_deref())?;
        let builder = ClientConfig::builder_with_provider(provider.clone())
            .with_protocol_versions(&versions)
            .map_err(|e| {
                ProtocolError::TlsError(format!("Failed to configure TLS protocol versions: {e}"))
            })?;
        Ok((provider, builder))
    }

    /// Build secure client config with system root CAs
    fn build_secure_client_config(&self) -> Result<ClientConfig> {
        let (_, builder) = self.client_config_builder()?;
        let root_store = self.load_system_root_certificates()?;
        let builder = builder.with_root_certificates(root_store);

        // Apply client auth directly
        if let (Some(client_cert_path), Some(client_key_path)) =
            (&self.client_cert_path, &self.client_key_path)
        {
            let (cert_chain, key) =
                self.load_client_credentials(client_cert_path, client_key_path)?;
            builder.with_client_auth_cert(cert_chain, key).map_err(|e| {
                ProtocolError::TlsError(format!("Failed to set client certificate: {e}"))
            })
        } else {
            Ok(builder.with_no_client_auth())
        }
    }

    /// Build insecure client config with custom verifier
    fn build_insecure_client_config(&self) -> Result<ClientConfig> {
        let (provider, builder) = self.client_config_builder()?;
        let verifier = self.create_custom_verifier(&provider);
        let custom_builder = builder
            .dangerous()
            .with_custom_certificate_verifier(verifier);

        // Apply client auth directly
        if let (Some(client_cert_path), Some(client_key_path)) =
            (&self.client_cert_path, &self.client_key_path)
        {
            let (cert_chain, key) =
                self.load_client_credentials(client_cert_path, client_key_path)?;
            custom_builder
                .with_client_auth_cert(cert_chain, key)
                .map_err(|e| {
                    ProtocolError::TlsError(format!("Failed to set client certificate: {e}"))
                })
        } else {
            Ok(custom_builder.with_no_client_auth())
        }
    }

    /// Load system root certificates
    fn load_system_root_certificates(&self) -> Result<RootCertStore> {
        let mut root_store = RootCertStore::empty();
        let native_certs = rustls_native_certs::load_native_certs();
        if let Some(e) = native_certs.errors.first() {
            return Err(ProtocolError::TlsError(format!(
                "Failed to load native certs: {e}"
            )));
        }

        for cert in native_certs.certs {
            root_store.add(cert).map_err(|e| {
                ProtocolError::TlsError(format!("Failed to add cert to root store: {e}"))
            })?;
        }

        Ok(root_store)
    }

    /// Create custom certificate verifier (pinning or accept-any)
    ///
    /// Handshake signatures are checked with the same provider's signature algorithms.
    fn create_custom_verifier(&self, provider: &CryptoProvider) -> Arc<dyn ServerCertVerifier> {
        let supported_algs = provider.signature_verification_algorithms;
        if let Some(hash) = &self.pinned_cert_hash {
            Arc::new(CertificateFingerprint {
                fingerprint: hash.clone(),
                supported_algs,
            })
        } else {
            Arc::new(AcceptAnyServerCert { supported_algs })
        }
    }

    /// Load client certificate and private key
    fn load_client_credentials(
        &self,
        cert_path: &str,
        key_path: &str,
    ) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>)> {
        // Load certificate
        let cert_file = File::open(cert_path).map_err(ProtocolError::Io)?;
        let mut cert_reader = BufReader::new(cert_file);
        let certs_result: std::result::Result<Vec<_>, _> =
            CertificateDer::pem_reader_iter(&mut cert_reader).collect();
        let certs: Vec<CertificateDer<'static>> = certs_result
            .map_err(|_| ProtocolError::TlsError("Failed to parse client certificate".into()))?;

        if certs.is_empty() {
            return Err(ProtocolError::TlsError(
                "No client certificates found".into(),
            ));
        }

        // Load private key
        let key_file = File::open(key_path).map_err(ProtocolError::Io)?;
        let mut key_reader = BufReader::new(key_file);
        let key = load_private_key(&mut key_reader)?;

        Ok((certs, key))
    }

    /// Get the server name as a rustls::ServerName
    pub fn server_name(&self) -> Result<ServerName<'_>> {
        ServerName::try_from(self.server_name.as_str())
            .map_err(|_| ProtocolError::TlsError("Invalid server name".into()))
    }

    /// Get the server name as an owned String
    pub fn server_name_string(&self) -> String {
        self.server_name.clone()
    }
}

/// Start a TLS server on the given address
#[instrument(skip(config))]
pub async fn start_server(addr: &str, config: TlsServerConfig) -> Result<()> {
    let tls_config = config.load_server_config()?;
    let acceptor = TlsAcceptor::from(Arc::new(tls_config));
    let listener = TcpListener::bind(addr).await?;

    info!(address=%addr, "TLS server listening");

    loop {
        let (stream, peer) = listener.accept().await?;
        let acceptor = acceptor.clone();

        tokio::spawn(async move {
            match acceptor.accept(stream).await {
                Ok(tls_stream) => {
                    if let Err(e) = handle_tls_connection(tls_stream, peer).await {
                        error!(%peer, error=%e, "Connection error");
                    }
                }
                Err(e) => {
                    error!(%peer, error=%e, "TLS handshake failed");
                }
            }
        });
    }
}

/// Handle a TLS connection
#[instrument(skip(tls_stream), fields(peer=%peer))]
async fn handle_tls_connection(
    tls_stream: ServerTlsStream<TcpStream>,
    peer: SocketAddr,
) -> Result<()> {
    let mut framed = Framed::new(tls_stream, PacketCodec);

    info!("TLS connection established");

    while let Some(packet) = framed.next().await {
        match packet {
            Ok(pkt) => {
                debug!(bytes = pkt.payload.len(), "Received data");
                on_packet(pkt, &mut framed).await?;
            }
            Err(e) => {
                error!(error=%e, "Protocol error");
                break;
            }
        }
    }

    info!("TLS connection closed");
    Ok(())
}

/// Handle incoming TLS packets
#[instrument(skip(framed), fields(packet_version=pkt.version, payload_size=pkt.payload.len()))]
async fn on_packet<T>(pkt: Packet, framed: &mut Framed<T, PacketCodec>) -> Result<()>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    // Echo the packet back (sample implementation)
    let response = Packet {
        version: pkt.version,
        payload: pkt.payload,
    };

    framed.send(response).await?;
    Ok(())
}

/// Connect to a TLS server
pub async fn connect(
    addr: &str,
    config: TlsClientConfig,
) -> Result<Framed<ClientTlsStream<TcpStream>, PacketCodec>> {
    let tls_config = Arc::new(config.load_client_config()?);
    let connector = TlsConnector::from(tls_config);

    let stream = TcpStream::connect(addr).await?;

    // Create ServerName from owned string to ensure 'static lifetime
    // Note: Box::leak() is used here to satisfy tokio_rustls' 'static requirement
    let server_name_str = config.server_name_string();
    let domain_static: &'static str = Box::leak(server_name_str.into_boxed_str());
    let domain = ServerName::try_from(domain_static)
        .map_err(|_| ProtocolError::TlsError("Invalid server name".into()))?;

    let tls_stream = connector
        .connect(domain, stream)
        .await
        .map_err(|e| ProtocolError::TlsError(format!("TLS connection failed: {e}")))?;

    let framed = Framed::new(tls_stream, PacketCodec);
    Ok(framed)
}