network_protocol/service/
tls_client.rs

1use futures::{SinkExt, StreamExt};
2use tokio::net::TcpStream;
3use tokio_rustls::client::TlsStream;
4use tokio_util::codec::Framed;
5
6use crate::core::codec::PacketCodec;
7use crate::core::packet::Packet;
8use crate::error::Result;
9use crate::protocol::message::Message;
10use crate::transport::tls::TlsClientConfig;
11
12/// TLS secure client for connecting to TLS-enabled servers
13pub struct TlsClient {
14    framed: Framed<TlsStream<TcpStream>, PacketCodec>,
15}
16
17impl TlsClient {
18    /// Connect to a TLS server
19    pub async fn connect(addr: &str, config: TlsClientConfig) -> Result<Self> {
20        let tls_config = config.load_client_config()?;
21        let connector = tokio_rustls::TlsConnector::from(std::sync::Arc::new(tls_config));
22
23        let stream = TcpStream::connect(addr).await?;
24        let domain = config.server_name()?;
25
26        let tls_stream = connector.connect(domain, stream).await?;
27        let framed = Framed::new(tls_stream, PacketCodec);
28
29        Ok(Self { framed })
30    }
31
32    /// Send a message to the TLS server
33    pub async fn send(&mut self, message: Message) -> Result<()> {
34        let bytes = bincode::serialize(&message)?;
35        let packet = Packet {
36            version: 1,
37            payload: bytes,
38        };
39
40        self.framed.send(packet).await?;
41        Ok(())
42    }
43
44    /// Receive a message from the TLS server
45    pub async fn receive(&mut self) -> Result<Message> {
46        let packet = match self.framed.next().await {
47            Some(Ok(pkt)) => pkt,
48            Some(Err(e)) => return Err(e),
49            None => {
50                return Err(crate::error::ProtocolError::Custom(
51                    "Connection closed".to_string(),
52                ))
53            }
54        };
55
56        let message = bincode::deserialize(&packet.payload)?;
57        Ok(message)
58    }
59
60    /// Send a message and wait for a response
61    pub async fn request(&mut self, message: Message) -> Result<Message> {
62        self.send(message).await?;
63        self.receive().await
64    }
65}