network_protocol/service/
tls_client.rs

1use futures::{SinkExt, StreamExt};
2use std::sync::Arc;
3use tokio::net::TcpStream;
4use tokio_rustls::client::TlsStream;
5use tokio_util::codec::Framed;
6use tracing::{debug, instrument};
7
8use crate::core::codec::PacketCodec;
9use crate::core::packet::Packet;
10use crate::error::{ProtocolError, Result};
11use crate::protocol::message::Message;
12use crate::transport::session_cache::SessionCache;
13use crate::transport::tls::TlsClientConfig;
14use crate::utils::timeout::HANDSHAKE_TIMEOUT;
15
16/// TLS secure client for connecting to TLS-enabled servers
17///
18/// Supports optional session resumption for improved reconnection performance.
19/// With session resumption, reconnections can skip the full TLS handshake,
20/// reducing latency by ~50-70%.
21pub struct TlsClient {
22    framed: Framed<TlsStream<TcpStream>, PacketCodec>,
23    /// Optional session cache for resumption support
24    session_cache: Option<Arc<SessionCache>>,
25    /// Session identifier (used with session cache)
26    session_id: Option<String>,
27}
28
29impl TlsClient {
30    /// Connect to a TLS server
31    #[instrument(skip(config))]
32    pub async fn connect(addr: &str, config: TlsClientConfig) -> Result<Self> {
33        Self::connect_with_session(addr, config, None).await
34    }
35
36    /// Connect to a TLS server with session resumption support
37    ///
38    /// With a session cache, the server's session tickets are stored in it and offered
39    /// on later connections that use the same cache, so a reconnect to the same server
40    /// name resumes the session instead of running a full handshake. Without one, every
41    /// connection does a full handshake.
42    ///
43    /// # Arguments
44    /// * `addr` - Server address to connect to
45    /// * `config` - TLS configuration
46    /// * `session_cache` - Optional session cache for resumption
47    ///
48    /// # Example
49    /// ```ignore
50    /// let cache = SessionCache::new(100, Duration::from_secs(3600));
51    /// let client = TlsClient::connect_with_session(
52    ///     "127.0.0.1:8443",
53    ///     config,
54    ///     Some(Arc::new(cache))
55    /// ).await?;
56    /// ```
57    pub async fn connect_with_session(
58        addr: &str,
59        config: TlsClientConfig,
60        session_cache: Option<Arc<SessionCache>>,
61    ) -> Result<Self> {
62        let tls_config = match &session_cache {
63            // The cache supplies one shared config per TLS configuration, so later
64            // connections see the same verifier and session store and can resume.
65            Some(cache) => {
66                debug!("Session resumption enabled");
67                cache.tls_client_config(&config)?
68            }
69            None => {
70                let mut tls_config = config.load_client_config()?;
71                tls_config.resumption = rustls::client::Resumption::disabled();
72                Arc::new(tls_config)
73            }
74        };
75        let connector = tokio_rustls::TlsConnector::from(tls_config);
76
77        let domain = config.owned_server_name()?;
78
79        let stream = TcpStream::connect(addr).await?;
80
81        let tls_stream = tokio::time::timeout(HANDSHAKE_TIMEOUT, connector.connect(domain, stream))
82            .await
83            .map_err(|_| ProtocolError::Timeout)??;
84        let framed = Framed::new(tls_stream, PacketCodec);
85
86        let session_id = format!(
87            "{}_{}",
88            addr,
89            std::time::SystemTime::now()
90                .duration_since(std::time::UNIX_EPOCH)
91                .unwrap_or_default()
92                .as_secs()
93        );
94
95        Ok(Self {
96            framed,
97            session_cache,
98            session_id: Some(session_id),
99        })
100    }
101
102    /// Send a message to the TLS server
103    pub async fn send(&mut self, message: Message) -> Result<()> {
104        let bytes = bincode::serialize(&message)?;
105        let packet = Packet {
106            version: 1,
107            payload: bytes,
108        };
109
110        self.framed.send(packet).await?;
111        Ok(())
112    }
113
114    /// Receive a message from the TLS server
115    pub async fn receive(&mut self) -> Result<Message> {
116        let packet = match self.framed.next().await {
117            Some(Ok(pkt)) => pkt,
118            Some(Err(e)) => return Err(e),
119            None => {
120                return Err(crate::error::ProtocolError::Custom(
121                    "Connection closed".to_string(),
122                ))
123            }
124        };
125
126        let message = bincode::deserialize(&packet.payload)?;
127        Ok(message)
128    }
129
130    /// Send a message and wait for a response
131    pub async fn request(&mut self, message: Message) -> Result<Message> {
132        self.send(message).await?;
133        self.receive().await
134    }
135
136    /// Get the session cache if configured
137    pub fn session_cache(&self) -> Option<&SessionCache> {
138        self.session_cache.as_deref()
139    }
140
141    /// Get the session identifier
142    pub fn session_id(&self) -> Option<&str> {
143        self.session_id.as_deref()
144    }
145}
network_protocol/service/tls_client.rs

network_protocol/service/
tls_client.rs
