network_protocol/service/
tls_daemon.rs

1use futures::{SinkExt, StreamExt};
2use std::sync::Arc;
3use std::time::Duration;
4use tokio::net::{TcpListener, TcpStream};
5use tokio::sync::{mpsc, Mutex};
6use tokio::task::JoinSet;
7use tokio_rustls::server::TlsStream;
8use tokio_util::codec::Framed;
9use tracing::{debug, error, info, instrument, warn};
10
11use crate::core::codec::PacketCodec;
12use crate::core::packet::Packet;
13use crate::protocol::dispatcher::Dispatcher;
14use crate::protocol::message::Message;
15// Secure connection not needed since TLS handles encryption
16use crate::error::Result;
17use crate::service::daemon::close_connections;
18use crate::transport::tls::TlsServerConfig;
19use crate::utils::timeout::HANDSHAKE_TIMEOUT;
20
21/// How long shutdown waits for open connections before closing them.
22const SHUTDOWN_GRACE: Duration = Duration::from_secs(10);
23
24/// Start a secure TLS server and listen for connections
25///
26/// Runs until the process receives Ctrl-C, then shuts down gracefully.
27#[instrument(skip(tls_config))]
28pub async fn start(addr: &str, tls_config: TlsServerConfig) -> Result<()> {
29    let (shutdown_tx, shutdown_rx) = mpsc::channel::<()>(1);
30
31    // Forward Ctrl-C to the shutdown channel. The task owns the only sender, so the
32    // channel stays open (and the server keeps running) until the signal arrives.
33    tokio::spawn(async move {
34        if let Ok(()) = tokio::signal::ctrl_c().await {
35            info!("Received shutdown signal, initiating graceful shutdown");
36            let _ = shutdown_tx.send(()).await;
37        }
38    });
39
40    start_with_shutdown(addr, tls_config, shutdown_rx).await
41}
42
43/// Start a secure TLS server with an external shutdown channel
44///
45/// The server shuts down gracefully when a `()` is sent on the channel. Dropping every
46/// sender without sending does not stop the server.
47#[instrument(skip(tls_config, shutdown_rx))]
48pub async fn start_with_shutdown(
49    addr: &str,
50    tls_config: TlsServerConfig,
51    mut shutdown_rx: mpsc::Receiver<()>,
52) -> Result<()> {
53    let config = tls_config.load_server_config()?;
54    let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(config));
55
56    let listener = TcpListener::bind(addr).await?;
57    info!(address=%addr, "TLS daemon listening");
58
59    // 🔁 Shared dispatcher
60    let dispatcher = Arc::new({
61        let d = Dispatcher::new();
62        let _ = d.register("PING", |_| Ok(Message::Pong));
63        let _ = d.register("ECHO", |msg| Ok(msg.clone()));
64        d
65    });
66
67    // Track active connections
68    let active_connections = Arc::new(Mutex::new(0u32));
69
70    // Once every sender is gone, `recv()` returns `None` immediately; stop polling it
71    // then instead of treating the closed channel as a shutdown request.
72    let mut shutdown_open = true;
73
74    // Connection tasks, so shutdown can wait for them and then close the rest
75    let mut connections = JoinSet::new();
76
77    // Server main loop with graceful shutdown
78    loop {
79        tokio::select! {
80            // Check for shutdown signal
81            signal = shutdown_rx.recv(), if shutdown_open => {
82                if signal.is_none() {
83                    shutdown_open = false;
84                    continue;
85                }
86                info!(connections = connections.len(), "Shutting down server. Waiting for connections to close...");
87                close_connections(&mut connections, SHUTDOWN_GRACE).await;
88                return Ok(());
89            }
90
91            // Reap finished connection tasks
92            Some(_) = connections.join_next(), if !connections.is_empty() => {}
93
94            // Accept new connections
95            accept_result = listener.accept() => {
96                match accept_result {
97                    Ok((stream, peer)) => {
98                        info!(%peer, "New connection accepted");
99                        let dispatcher = dispatcher.clone();
100                        let acceptor = acceptor.clone();
101                        let active_connections = active_connections.clone();
102
103                        // Increment active connections counter
104                        {
105                            let mut count = active_connections.lock().await;
106                            *count += 1;
107                        }
108
109                        connections.spawn(async move {
110                            // Bound the handshake so a client that stalls cannot hold a
111                            // connection slot open indefinitely.
112                            let handshake = tokio::time::timeout(HANDSHAKE_TIMEOUT, acceptor.accept(stream)).await;
113                            match handshake {
114                                Ok(Ok(tls_stream)) => {
115                                    if let Err(e) = handle_tls_connection(tls_stream, dispatcher, peer, active_connections).await {
116                                        error!(%peer, error=%e, "Connection error");
117                                    }
118                                },
119                                Ok(Err(e)) => {
120                                    error!(%peer, error=%e, "TLS handshake failed");
121                                    // Decrement connections on handshake failure
122                                    let mut count = active_connections.lock().await;
123                                    *count -= 1;
124                                }
125                                Err(_) => {
126                                    warn!(%peer, "TLS handshake timed out");
127                                    let mut count = active_connections.lock().await;
128                                    *count -= 1;
129                                }
130                            }
131                        });
132                    }
133                    Err(e) => {
134                        error!(error=%e, "Error accepting connection");
135                    }
136                }
137            }
138        }
139    }
140}
141
142/// Handle a TLS connection
143#[instrument(skip(tls_stream, dispatcher, active_connections), fields(peer=%peer))]
144async fn handle_tls_connection(
145    tls_stream: TlsStream<TcpStream>,
146    dispatcher: Arc<Dispatcher>,
147    peer: std::net::SocketAddr,
148    active_connections: Arc<Mutex<u32>>,
149) -> Result<()> {
150    let mut framed = Framed::new(tls_stream, PacketCodec);
151
152    info!("TLS connection established");
153
154    // Unlike regular daemon, we don't need a separate handshake
155    // TLS already provides the encryption layer
156
157    // Message loop
158    loop {
159        let packet = match framed.next().await {
160            Some(Ok(pkt)) => pkt,
161            Some(Err(e)) => {
162                error!(error=%e, "Protocol error");
163                break;
164            }
165            None => break,
166        };
167
168        // Deserialize the message
169        let msg = match bincode::deserialize::<Message>(&packet.payload) {
170            Ok(m) => m,
171            Err(e) => {
172                error!(error=%e, "Deserialization error");
173                continue;
174            }
175        };
176
177        debug!(message=?msg, "Received message");
178
179        // Process with dispatcher
180        match dispatcher.dispatch(&msg) {
181            Ok(reply) => {
182                let reply_bytes = match bincode::serialize(&reply) {
183                    Ok(bytes) => bytes,
184                    Err(e) => {
185                        error!(error=%e, "Serialization error");
186                        continue;
187                    }
188                };
189
190                let reply_packet = Packet {
191                    version: packet.version,
192                    payload: reply_bytes,
193                };
194
195                if let Err(e) = framed.send(reply_packet).await {
196                    error!(error=%e, "Send error");
197                    break;
198                }
199            }
200            Err(e) => {
201                error!(error=%e, "Dispatch error");
202                break;
203            }
204        }
205    }
206
207    info!("Connection closed");
208
209    // Decrement connection counter on disconnect
210    {
211        let mut count = active_connections.lock().await;
212        *count -= 1;
213    }
214
215    Ok(())
216}
network_protocol/service/tls_daemon.rs

network_protocol/service/
tls_daemon.rs
