Skip to main content

network_protocol/transport/
tls.rs

1//! # TLS Transport Layer
2//!
3//! This file is part of the Network Protocol project.
4//!
5//! It defines the TLS transport layer for secure network communication,
6//! particularly for external untrusted connections.
7//!
8//! The TLS transport layer provides a secure channel for communication
9//! using industry-standard TLS protocol, ensuring confidentiality,
10//! integrity, and authentication of the data transmitted.
11//!
12//! ## Responsibilities
13//! - Establish secure TLS connections
14//! - Handle TLS certificates and verification
15//! - Provide secure framed transport for higher protocol layers
16//! - Compatible with existing packet codec infrastructure
17
18use std::fs::File;
19use std::io::{self, BufReader, Seek, Write};
20use std::net::SocketAddr;
21use std::path::Path;
22use std::sync::Arc;
23
24use rustls::client::danger::{ServerCertVerified, ServerCertVerifier};
25use rustls::pki_types::pem::PemObject;
26use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName, UnixTime};
27use rustls::{ClientConfig, DigitallySignedStruct, RootCertStore, ServerConfig};
28use tokio::net::{TcpListener, TcpStream};
29use tokio_rustls::client::TlsStream as ClientTlsStream;
30use tokio_rustls::server::TlsStream as ServerTlsStream;
31use tokio_rustls::{TlsAcceptor, TlsConnector};
32use tokio_util::codec::Framed;
33use tracing::{debug, error, info, instrument, warn};
34
35use crate::core::codec::PacketCodec;
36use crate::core::packet::Packet;
37use crate::error::{ProtocolError, Result};
38use futures::{SinkExt, StreamExt};
39
40// Custom certificate verifiers
41#[derive(Debug)]
42struct CertificateFingerprint {
43    fingerprint: Vec<u8>,
44}
45
46impl ServerCertVerifier for CertificateFingerprint {
47    fn verify_server_cert(
48        &self,
49        end_entity: &CertificateDer<'_>,
50        _intermediates: &[CertificateDer<'_>],
51        _server_name: &ServerName,
52        _ocsp_response: &[u8],
53        _now: UnixTime,
54    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
55        use sha2::{Digest, Sha256};
56
57        let mut hasher = Sha256::new();
58        hasher.update(end_entity);
59        let hash = hasher.finalize();
60
61        if hash.as_slice() == self.fingerprint.as_slice() {
62            Ok(ServerCertVerified::assertion())
63        } else {
64            Err(rustls::Error::General(
65                "Pinned certificate hash mismatch".into(),
66            ))
67        }
68    }
69
70    fn verify_tls12_signature(
71        &self,
72        _message: &[u8],
73        _cert: &CertificateDer<'_>,
74        _dss: &DigitallySignedStruct,
75    ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
76        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
77    }
78
79    fn verify_tls13_signature(
80        &self,
81        _message: &[u8],
82        _cert: &CertificateDer<'_>,
83        _dss: &DigitallySignedStruct,
84    ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
85        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
86    }
87
88    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
89        // Accept common signature schemes
90        vec![
91            rustls::SignatureScheme::RSA_PKCS1_SHA256,
92            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
93            rustls::SignatureScheme::ED25519,
94        ]
95    }
96}
97
98#[derive(Debug)]
99struct AcceptAnyServerCert;
100
101impl ServerCertVerifier for AcceptAnyServerCert {
102    fn verify_server_cert(
103        &self,
104        _end_entity: &CertificateDer<'_>,
105        _intermediates: &[CertificateDer<'_>],
106        _server_name: &ServerName,
107        _ocsp_response: &[u8],
108        _now: UnixTime,
109    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
110        Ok(ServerCertVerified::assertion())
111    }
112
113    fn verify_tls12_signature(
114        &self,
115        _message: &[u8],
116        _cert: &CertificateDer<'_>,
117        _dss: &DigitallySignedStruct,
118    ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
119        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
120    }
121
122    fn verify_tls13_signature(
123        &self,
124        _message: &[u8],
125        _cert: &CertificateDer<'_>,
126        _dss: &DigitallySignedStruct,
127    ) -> std::result::Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
128        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
129    }
130
131    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
132        vec![
133            rustls::SignatureScheme::RSA_PKCS1_SHA256,
134            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
135            rustls::SignatureScheme::ED25519,
136        ]
137    }
138}
139
140/// Helper function to load a private key from PKCS8 format
141fn load_private_key(reader: &mut BufReader<File>) -> Result<PrivateKeyDer<'static>> {
142    // Try to load PKCS8 keys
143    // Seek to beginning of file first
144    reader
145        .seek(std::io::SeekFrom::Start(0))
146        .map_err(ProtocolError::Io)?;
147
148    // Collect every PKCS8 private key section in the PEM input
149    let keys: std::result::Result<Vec<_>, _> =
150        PrivatePkcs8KeyDer::pem_reader_iter(reader).collect();
151    let keys =
152        keys.map_err(|_| ProtocolError::TlsError("Failed to parse PKCS8 private key".into()))?;
153
154    if !keys.is_empty() {
155        return Ok(PrivateKeyDer::Pkcs8(keys[0].clone_key()));
156    }
157
158    // Note: Add support for other key formats like RSA or EC if needed
159
160    Err(ProtocolError::TlsError(
161        "No supported private key format found".into(),
162    ))
163}
164
165/// TLS protocol version
166pub enum TlsVersion {
167    /// TLS 1.2
168    TLS12,
169    /// TLS 1.3
170    TLS13,
171    /// Both TLS 1.2 and 1.3
172    All,
173}
174
175/// TLS server configuration
176pub struct TlsServerConfig {
177    cert_path: String,
178    key_path: String,
179    /// Optional path to client CA certificates for mTLS
180    client_ca_path: Option<String>,
181    /// Whether to require client certificates (mTLS)
182    require_client_auth: bool,
183    /// Allowed TLS protocol versions (None = use rustls defaults)
184    tls_versions: Option<Vec<TlsVersion>>,
185    /// Allowed cipher suites (None = use rustls defaults)
186    cipher_suites: Option<Vec<rustls::SupportedCipherSuite>>,
187    /// ALPN protocols to advertise (default: ["h2", "http/1.1"])
188    alpn_protocols: Option<Vec<Vec<u8>>>,
189}
190
191impl TlsServerConfig {
192    /// Create a new TLS server configuration
193    pub fn new<P: AsRef<std::path::Path>>(cert_path: P, key_path: P) -> Self {
194        Self {
195            cert_path: cert_path.as_ref().to_string_lossy().to_string(),
196            key_path: key_path.as_ref().to_string_lossy().to_string(),
197            client_ca_path: None,
198            require_client_auth: false,
199            tls_versions: None,
200            cipher_suites: None,
201            alpn_protocols: Some(vec![b"h2".to_vec(), b"http/1.1".to_vec()]),
202        }
203    }
204
205    /// Set allowed TLS protocol versions
206    pub fn with_tls_versions(mut self, versions: Vec<TlsVersion>) -> Self {
207        self.tls_versions = Some(versions);
208        self
209    }
210
211    /// Set allowed cipher suites
212    pub fn with_cipher_suites(mut self, cipher_suites: Vec<rustls::SupportedCipherSuite>) -> Self {
213        self.cipher_suites = Some(cipher_suites);
214        self
215    }
216
217    /// Enable mutual TLS authentication by providing a CA certificate path
218    pub fn with_client_auth<S: Into<String>>(mut self, client_ca_path: S) -> Self {
219        self.client_ca_path = Some(client_ca_path.into());
220        self.require_client_auth = true;
221        self
222    }
223
224    /// Set whether client authentication is required (true) or optional (false)
225    pub fn require_client_auth(mut self, required: bool) -> Self {
226        self.require_client_auth = required;
227        self
228    }
229
230    /// Set ALPN protocols to advertise during TLS handshake
231    pub fn with_alpn_protocols(mut self, protocols: Vec<Vec<u8>>) -> Self {
232        self.alpn_protocols = Some(protocols);
233        self
234    }
235
236    /// Generate a self-signed certificate for development/testing purposes
237    pub fn generate_self_signed<P: AsRef<Path>>(cert_path: P, key_path: P) -> io::Result<Self> {
238        let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()])
239            .map_err(|e| io::Error::other(format!("Certificate generation error: {e}")))?;
240
241        // Write certificate
242        let mut cert_file = File::create(&cert_path)?;
243        let pem = cert.cert.pem();
244        cert_file.write_all(pem.as_bytes())?;
245
246        // Write private key
247        let mut key_file = File::create(&key_path)?;
248        key_file.write_all(cert.signing_key.serialize_pem().as_bytes())?;
249
250        Ok(Self {
251            cert_path: cert_path.as_ref().to_string_lossy().to_string(),
252            key_path: key_path.as_ref().to_string_lossy().to_string(),
253            client_ca_path: None,
254            require_client_auth: false,
255            tls_versions: None,
256            cipher_suites: None,
257            alpn_protocols: Some(vec![b"h2".to_vec(), b"http/1.1".to_vec()]),
258        })
259    }
260
261    /// Load the TLS configuration from files
262    pub fn load_server_config(&self) -> Result<ServerConfig> {
263        // Load certificate
264        let cert_file = File::open(&self.cert_path)
265            .map_err(|e| ProtocolError::TlsError(format!("Failed to open cert file: {e}")))?;
266        let mut cert_reader = BufReader::new(cert_file);
267        let cert_chain: std::result::Result<Vec<_>, _> =
268            CertificateDer::pem_reader_iter(&mut cert_reader).collect();
269        let cert_chain: Vec<CertificateDer<'static>> = cert_chain
270            .map_err(|_| ProtocolError::TlsError("Failed to parse certificate".into()))?;
271
272        if cert_chain.is_empty() {
273            return Err(ProtocolError::TlsError("No certificates found".into()));
274        }
275
276        // Load private key
277        let key_file = File::open(&self.key_path)
278            .map_err(|e| ProtocolError::TlsError(format!("Failed to open key file: {e}")))?;
279        let mut key_reader = BufReader::new(key_file);
280        let private_key = load_private_key(&mut key_reader)?;
281
282        // Validate TLS versions if specified
283        // Note: In rustls 0.22, with_safe_defaults() restricts to TLS 1.2+ (best practice)
284        if let Some(versions) = &self.tls_versions {
285            let mut has_tls13 = false;
286            let mut has_tls12 = false;
287            for v in versions {
288                match v {
289                    TlsVersion::TLS12 => has_tls12 = true,
290                    TlsVersion::TLS13 => has_tls13 = true,
291                    TlsVersion::All => {
292                        has_tls13 = true;
293                        has_tls12 = true;
294                    }
295                }
296            }
297            // Document that with_safe_defaults uses best practices
298            debug!(
299                "TLS versions requested: TLS1.2={}, TLS1.3={}",
300                has_tls12, has_tls13
301            );
302        }
303
304        // Create a server configuration with safe defaults (TLS 1.2+, modern ciphersuites)
305        let config_builder = ServerConfig::builder_with_provider(std::sync::Arc::new(
306            rustls::crypto::ring::default_provider(),
307        ))
308        .with_safe_default_protocol_versions()
309        .map_err(|_| ProtocolError::TlsError("Failed to configure TLS protocol versions".into()))?;
310
311        let cert_builder = config_builder.with_no_client_auth();
312
313        // Build config with certificates
314        let mut config = cert_builder
315            .with_single_cert(cert_chain.clone(), private_key.clone_key())
316            .map_err(|e| ProtocolError::TlsError(format!("TLS error: {e}")))?;
317
318        // Configure client authentication if required (mTLS)
319        if let Some(client_ca_path) = &self.client_ca_path {
320            // Load client CA certificates
321            let client_ca_file = File::open(client_ca_path).map_err(|e| {
322                ProtocolError::TlsError(format!("Failed to open client CA file: {e}"))
323            })?;
324            let mut client_ca_reader = BufReader::new(client_ca_file);
325            let client_ca_certs: std::result::Result<Vec<_>, _> =
326                CertificateDer::pem_reader_iter(&mut client_ca_reader).collect();
327            let client_ca_certs: Vec<CertificateDer<'static>> = client_ca_certs.map_err(|_| {
328                ProtocolError::TlsError("Failed to parse client CA certificate".into())
329            })?;
330
331            if client_ca_certs.is_empty() {
332                return Err(ProtocolError::TlsError(
333                    "No client CA certificates found".into(),
334                ));
335            }
336
337            // Create client cert verifier
338            let mut client_root_store = RootCertStore::empty();
339            for cert in client_ca_certs {
340                client_root_store.add(cert).map_err(|e| {
341                    ProtocolError::TlsError(format!("Failed to add client CA cert: {e}"))
342                })?;
343            }
344
345            // Create client authentication verifier using WebPkiClientVerifier
346            let client_auth = rustls::server::WebPkiClientVerifier::builder(std::sync::Arc::new(
347                client_root_store,
348            ))
349            .build()
350            .map_err(|e| {
351                ProtocolError::TlsError(format!("Failed to build client verifier: {e}"))
352            })?;
353
354            // Create new config builder with client auth
355            let new_builder = ServerConfig::builder_with_provider(std::sync::Arc::new(
356                rustls::crypto::ring::default_provider(),
357            ))
358            .with_safe_default_protocol_versions()
359            .map_err(|_| {
360                ProtocolError::TlsError("Failed to configure TLS protocol versions".into())
361            })?;
362            let new_cert_builder = new_builder.with_client_cert_verifier(client_auth);
363
364            // Build a new config with certificates and client auth
365            config = new_cert_builder
366                .with_single_cert(cert_chain, private_key.clone_key())
367                .map_err(|e| ProtocolError::TlsError(format!("TLS error with client auth: {e}")))?;
368
369            debug!("mTLS enabled with client certificate verification required");
370        }
371
372        // Configure ALPN protocols if specified
373        if let Some(protocols) = &self.alpn_protocols {
374            config.alpn_protocols = protocols.clone();
375            debug!(
376                protocol_count = protocols.len(),
377                "ALPN protocols configured"
378            );
379        }
380
381        Ok(config)
382    }
383
384    /// Calculate SHA-256 hash for a certificate to use with pinning
385    pub fn calculate_cert_hash(cert: &CertificateDer<'_>) -> Vec<u8> {
386        use sha2::{Digest, Sha256};
387        let mut hasher = Sha256::new();
388        hasher.update(cert.as_ref());
389        hasher.finalize().to_vec()
390    }
391}
392
393/// TLS Client Configuration
394pub struct TlsClientConfig {
395    server_name: String,
396    insecure: bool,
397    /// Optional certificate hash to pin (SHA-256 fingerprint)
398    pinned_cert_hash: Option<Vec<u8>>,
399    /// Optional client certificate path for mTLS
400    client_cert_path: Option<String>,
401    /// Optional client key path for mTLS
402    client_key_path: Option<String>,
403    /// Allowed TLS protocol versions (None = use rustls defaults)
404    tls_versions: Option<Vec<TlsVersion>>,
405    /// Allowed cipher suites (None = use rustls defaults)
406    cipher_suites: Option<Vec<rustls::SupportedCipherSuite>>,
407}
408
409impl TlsClientConfig {
410    /// Create a new TLS client configuration
411    pub fn new<S: Into<String>>(server_name: S) -> Self {
412        Self {
413            server_name: server_name.into(),
414            insecure: false,
415            pinned_cert_hash: None,
416            client_cert_path: None,
417            client_key_path: None,
418            tls_versions: None,
419            cipher_suites: None,
420        }
421    }
422
423    /// Set allowed TLS protocol versions
424    pub fn with_tls_versions(mut self, versions: Vec<TlsVersion>) -> Self {
425        self.tls_versions = Some(versions);
426        self
427    }
428
429    /// Set allowed cipher suites
430    pub fn with_cipher_suites(mut self, cipher_suites: Vec<rustls::SupportedCipherSuite>) -> Self {
431        self.cipher_suites = Some(cipher_suites);
432        self
433    }
434
435    /// Configure client authentication for mTLS
436    pub fn with_client_certificate<S: Into<String>>(mut self, cert_path: S, key_path: S) -> Self {
437        self.client_cert_path = Some(cert_path.into());
438        self.client_key_path = Some(key_path.into());
439        self
440    }
441
442    /// Allow insecure connections (skip certificate verification)
443    ///
444    /// # WARNING: Security Risk
445    /// This mode disables certificate verification entirely and should ONLY be used for:
446    /// - Development and testing
447    /// - Debugging environments
448    /// - Internal networks with certificate pinning enabled
449    ///
450    /// **NEVER** use this in production without explicit certificate pinning via `with_pinned_cert_hash()`.
451    ///
452    /// For maximum security, only use this with the "dangerous_configuration" feature enabled,
453    /// which is a strong indicator this is for testing/development only.
454    pub fn insecure(mut self) -> Self {
455        warn!("INSECURE MODE ENABLED: Certificate verification is disabled. This should only be used for development/testing.");
456        self.insecure = true;
457        self
458    }
459
460    /// Pin a certificate by its SHA-256 hash/fingerprint
461    ///
462    /// This provides additional security by only accepting connections
463    /// from servers with the exact certificate matching this hash.
464    /// Can be combined with insecure mode for development environments where
465    /// you want to skip standard CA verification but still verify a specific cert.
466    pub fn with_pinned_cert_hash(mut self, hash: Vec<u8>) -> Self {
467        if hash.len() != 32 {
468            warn!(
469                "Certificate hash has unexpected length: {} (expected 32 bytes for SHA-256)",
470                hash.len()
471            );
472        }
473        self.pinned_cert_hash = Some(hash);
474        self
475    }
476
477    /// Load the TLS client configuration
478    pub fn load_client_config(&self) -> Result<ClientConfig> {
479        self.log_tls_version_info();
480
481        if self.insecure {
482            self.build_insecure_client_config()
483        } else {
484            self.build_secure_client_config()
485        }
486    }
487
488    /// Log TLS version configuration
489    fn log_tls_version_info(&self) {
490        if let Some(versions) = &self.tls_versions {
491            let mut has_tls13 = false;
492            let mut has_tls12 = false;
493            for v in versions {
494                match v {
495                    TlsVersion::TLS12 => has_tls12 = true,
496                    TlsVersion::TLS13 => has_tls13 = true,
497                    TlsVersion::All => {
498                        has_tls13 = true;
499                        has_tls12 = true;
500                    }
501                }
502            }
503            debug!(
504                "TLS client versions requested: TLS1.2={}, TLS1.3={}",
505                has_tls12, has_tls13
506            );
507        }
508    }
509
510    /// Build secure client config with system root CAs
511    fn build_secure_client_config(&self) -> Result<ClientConfig> {
512        let root_store = self.load_system_root_certificates()?;
513        let builder = ClientConfig::builder_with_provider(std::sync::Arc::new(
514            rustls::crypto::ring::default_provider(),
515        ))
516        .with_safe_default_protocol_versions()
517        .map_err(|_| ProtocolError::TlsError("Failed to configure TLS protocol versions".into()))?
518        .with_root_certificates(root_store);
519
520        // Apply client auth directly
521        if let (Some(client_cert_path), Some(client_key_path)) =
522            (&self.client_cert_path, &self.client_key_path)
523        {
524            let (cert_chain, key) =
525                self.load_client_credentials(client_cert_path, client_key_path)?;
526            builder.with_client_auth_cert(cert_chain, key).map_err(|e| {
527                ProtocolError::TlsError(format!("Failed to set client certificate: {e}"))
528            })
529        } else {
530            Ok(builder.with_no_client_auth())
531        }
532    }
533
534    /// Build insecure client config with custom verifier
535    fn build_insecure_client_config(&self) -> Result<ClientConfig> {
536        let builder = ClientConfig::builder_with_provider(std::sync::Arc::new(
537            rustls::crypto::ring::default_provider(),
538        ))
539        .with_safe_default_protocol_versions()
540        .map_err(|_| ProtocolError::TlsError("Failed to configure TLS protocol versions".into()))?;
541        let verifier = self.create_custom_verifier();
542        let custom_builder = builder
543            .dangerous()
544            .with_custom_certificate_verifier(verifier);
545
546        // Apply client auth directly
547        if let (Some(client_cert_path), Some(client_key_path)) =
548            (&self.client_cert_path, &self.client_key_path)
549        {
550            let (cert_chain, key) =
551                self.load_client_credentials(client_cert_path, client_key_path)?;
552            custom_builder
553                .with_client_auth_cert(cert_chain, key)
554                .map_err(|e| {
555                    ProtocolError::TlsError(format!("Failed to set client certificate: {e}"))
556                })
557        } else {
558            Ok(custom_builder.with_no_client_auth())
559        }
560    }
561
562    /// Load system root certificates
563    fn load_system_root_certificates(&self) -> Result<RootCertStore> {
564        let mut root_store = RootCertStore::empty();
565        let native_certs = rustls_native_certs::load_native_certs();
566        if let Some(e) = native_certs.errors.first() {
567            return Err(ProtocolError::TlsError(format!(
568                "Failed to load native certs: {e}"
569            )));
570        }
571
572        for cert in native_certs.certs {
573            root_store.add(cert).map_err(|e| {
574                ProtocolError::TlsError(format!("Failed to add cert to root store: {e}"))
575            })?;
576        }
577
578        Ok(root_store)
579    }
580
581    /// Create custom certificate verifier (pinning or accept-any)
582    fn create_custom_verifier(&self) -> Arc<dyn ServerCertVerifier> {
583        if let Some(hash) = &self.pinned_cert_hash {
584            Arc::new(CertificateFingerprint {
585                fingerprint: hash.clone(),
586            })
587        } else {
588            Arc::new(AcceptAnyServerCert)
589        }
590    }
591
592    /// Load client certificate and private key
593    fn load_client_credentials(
594        &self,
595        cert_path: &str,
596        key_path: &str,
597    ) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>)> {
598        // Load certificate
599        let cert_file = File::open(cert_path).map_err(ProtocolError::Io)?;
600        let mut cert_reader = BufReader::new(cert_file);
601        let certs_result: std::result::Result<Vec<_>, _> =
602            CertificateDer::pem_reader_iter(&mut cert_reader).collect();
603        let certs: Vec<CertificateDer<'static>> = certs_result
604            .map_err(|_| ProtocolError::TlsError("Failed to parse client certificate".into()))?;
605
606        if certs.is_empty() {
607            return Err(ProtocolError::TlsError(
608                "No client certificates found".into(),
609            ));
610        }
611
612        // Load private key
613        let key_file = File::open(key_path).map_err(ProtocolError::Io)?;
614        let mut key_reader = BufReader::new(key_file);
615        let key = load_private_key(&mut key_reader)?;
616
617        Ok((certs, key))
618    }
619
620    /// Get the server name as a rustls::ServerName
621    pub fn server_name(&self) -> Result<ServerName<'_>> {
622        ServerName::try_from(self.server_name.as_str())
623            .map_err(|_| ProtocolError::TlsError("Invalid server name".into()))
624    }
625
626    /// Get the server name as an owned String
627    pub fn server_name_string(&self) -> String {
628        self.server_name.clone()
629    }
630}
631
632/// Start a TLS server on the given address
633#[instrument(skip(config))]
634pub async fn start_server(addr: &str, config: TlsServerConfig) -> Result<()> {
635    let tls_config = config.load_server_config()?;
636    let acceptor = TlsAcceptor::from(Arc::new(tls_config));
637    let listener = TcpListener::bind(addr).await?;
638
639    info!(address=%addr, "TLS server listening");
640
641    loop {
642        let (stream, peer) = listener.accept().await?;
643        let acceptor = acceptor.clone();
644
645        tokio::spawn(async move {
646            match acceptor.accept(stream).await {
647                Ok(tls_stream) => {
648                    if let Err(e) = handle_tls_connection(tls_stream, peer).await {
649                        error!(%peer, error=%e, "Connection error");
650                    }
651                }
652                Err(e) => {
653                    error!(%peer, error=%e, "TLS handshake failed");
654                }
655            }
656        });
657    }
658}
659
660/// Handle a TLS connection
661#[instrument(skip(tls_stream), fields(peer=%peer))]
662async fn handle_tls_connection(
663    tls_stream: ServerTlsStream<TcpStream>,
664    peer: SocketAddr,
665) -> Result<()> {
666    let mut framed = Framed::new(tls_stream, PacketCodec);
667
668    info!("TLS connection established");
669
670    while let Some(packet) = framed.next().await {
671        match packet {
672            Ok(pkt) => {
673                debug!(bytes = pkt.payload.len(), "Received data");
674                on_packet(pkt, &mut framed).await?;
675            }
676            Err(e) => {
677                error!(error=%e, "Protocol error");
678                break;
679            }
680        }
681    }
682
683    info!("TLS connection closed");
684    Ok(())
685}
686
687/// Handle incoming TLS packets
688#[instrument(skip(framed), fields(packet_version=pkt.version, payload_size=pkt.payload.len()))]
689async fn on_packet<T>(pkt: Packet, framed: &mut Framed<T, PacketCodec>) -> Result<()>
690where
691    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
692{
693    // Echo the packet back (sample implementation)
694    let response = Packet {
695        version: pkt.version,
696        payload: pkt.payload,
697    };
698
699    framed.send(response).await?;
700    Ok(())
701}
702
703/// Connect to a TLS server
704pub async fn connect(
705    addr: &str,
706    config: TlsClientConfig,
707) -> Result<Framed<ClientTlsStream<TcpStream>, PacketCodec>> {
708    let tls_config = Arc::new(config.load_client_config()?);
709    let connector = TlsConnector::from(tls_config);
710
711    let stream = TcpStream::connect(addr).await?;
712
713    // Create ServerName from owned string to ensure 'static lifetime
714    // Note: Box::leak() is used here to satisfy tokio_rustls' 'static requirement
715    let server_name_str = config.server_name_string();
716    let domain_static: &'static str = Box::leak(server_name_str.into_boxed_str());
717    let domain = ServerName::try_from(domain_static)
718        .map_err(|_| ProtocolError::TlsError("Invalid server name".into()))?;
719
720    let tls_stream = connector
721        .connect(domain, stream)
722        .await
723        .map_err(|e| ProtocolError::TlsError(format!("TLS connection failed: {e}")))?;
724
725    let framed = Framed::new(tls_stream, PacketCodec);
726    Ok(framed)
727}