feat: Integrate TLS on client and server side for secure connection

ByteMaster2003 · ByteMaster2003 · commit e21852a056cc · 2025-08-19T16:06:04.000+05:30
Add required crates rustls, tokio-rustls, webpki-roots and rustls-pemfile
Implement functions to handle tls accepter and connector on server and client side
Update serve config file to take file paths for fullchain.pem and privkey.pem
Implement functions on server side to parse pem files
Keep TLS optional on server side
On Client side try to upgrade server for 10 seconds and fallback to non tls server
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -15,7 +15,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-latest, ubuntu-22.04-arm64, macos-latest, windows-latest]
+        os: [ubuntu-latest, ubuntu-latest-arm64, macos-latest, windows-latest]
 
     steps:
       - name: Checkout
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,3 +11,5 @@ rsa = { version = "0.9.8", features = ["serde", "sha2"] }
 ssh-key = "0.6.7"
 hex = "0.4.3"
 bincode = "2.0.1"
+tokio-rustls = "0.26.2"
+rustls = "0.23.31"
diff --git a/client/Cargo.toml b/client/Cargo.toml
@@ -19,3 +19,6 @@ ratatui = "0.29.0"
 color-eyre = "0.6.5"
 chrono = "0.4.41"
 tui-textarea = "0.7.0"
+tokio-rustls.workspace = true
+rustls.workspace = true
+webpki-roots = "1.0.2"
diff --git a/client/src/main.rs b/client/src/main.rs
@@ -3,7 +3,7 @@ use null_talk_client::{
     handlers::handle_client,
     types::{LogLevel, LogMessage},
     ui::run_terminal,
-    utils::configure_client,
+    utils::{configure_client, try_tls_handshake},
 };
 use std::env;
 use tokio::net::TcpStream;
@@ -17,7 +17,7 @@ async fn main() {
 
     // Create TCP connection thread
     let tcp = tokio::spawn(async move {
-        let addr = {
+        let (addr, host_name) = {
             let config_lock = data::CLIENT_CONFIG.lock().await;
             let config = match config_lock.as_ref() {
                 Some(cfg) => cfg,
@@ -27,19 +27,52 @@ async fn main() {
                 }
             };
 
-            format!("{}:{}", &config.hostname, &config.port)
+            (
+                format!("{}:{}", &config.hostname, &config.port),
+                config.hostname.clone(),
+            )
         };
 
         match TcpStream::connect(addr.clone()).await {
             Ok(stream) => {
                 LogMessage::log(
                     LogLevel::INFO,
-                    format!("Successfully connected to {}", addr.clone()),
-                    5,
+                    format!("Establishing a secure TLS connection..."),
+                    0,
                 )
                 .await;
-
-                handle_client(Box::new(stream)).await;
+                let result = try_tls_handshake(host_name, stream).await;
+                match result {
+                    Some(tls_stream) => {
+                        LogMessage::log(
+                            LogLevel::INFO,
+                            format!("Successfully connected to {}", addr),
+                            5,
+                        )
+                        .await;
+                        handle_client(tls_stream).await;
+                    }
+                    None => match TcpStream::connect(addr.clone()).await {
+                        Ok(plain_stream) => {
+                            LogMessage::log(
+                                LogLevel::INFO,
+                                format!("Successfully connected to {}, TLS not enabled", addr),
+                                5,
+                            )
+                            .await;
+                            handle_client(Box::new(plain_stream)).await;
+                        }
+                        Err(_) => {
+                            LogMessage::log(
+                                LogLevel::ERROR,
+                                format!("Failed to connect to {}", addr),
+                                0,
+                            )
+                            .await;
+                            return;
+                        }
+                    },
+                };
             }
             Err(_) => {
                 LogMessage::log(LogLevel::ERROR, format!("Failed to connect to {}", addr), 0).await;
diff --git a/client/src/utils/net.rs b/client/src/utils/net.rs
@@ -1,12 +1,22 @@
-use crate::data::CLIENT_CONFIG;
+use crate::{
+    data::CLIENT_CONFIG,
+    types::{LogLevel, LogMessage},
+};
 use common::{
-    net::{HandshakePacket, StreamReader, StreamWriter},
+    net::{AsyncStream, HandshakePacket, StreamReader, StreamWriter},
     utils::{enc as encutils, net as netutils},
 };
 use rsa::RsaPrivateKey;
+use rustls::{ClientConfig, pki_types::ServerName};
+use std::sync::Arc;
+use tokio::{
+    net::TcpStream,
+    time::{Duration, timeout},
+};
+use tokio_rustls::TlsConnector;
 
 /// ### Perform the initial handshake with the server.
-/// 
+///
 /// This function handles the entire handshake process, including
 /// sending the initial handshake packet, receiving the server's
 /// response, and completing the handshake by establishing a
@@ -83,3 +93,58 @@ pub async fn perform_handshake(
 
     Ok(session_key)
 }
+
+pub async fn create_tls_connector() -> Result<TlsConnector, Box<dyn std::error::Error + Send + Sync>>
+{
+    // Load the root certificates from the webpki-roots crate
+    let mut root_store = rustls::RootCertStore::empty();
+    root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
+
+    // Create the client configuration
+    let config = ClientConfig::builder()
+        .with_root_certificates(root_store)
+        .with_no_client_auth();
+
+    // Create the async TLS connector
+    let connector = TlsConnector::from(Arc::new(config));
+    Ok(connector)
+}
+
+pub async fn try_tls_handshake(
+    host_name: String,
+    stream: TcpStream,
+) -> Option<Box<dyn AsyncStream>> {
+    // Convert the hostname string to a ServerName
+    let domain = match ServerName::try_from(host_name) {
+        Ok(name) => name,
+        Err(_) => return None,
+    };
+
+    let tls_connector = match create_tls_connector().await {
+        Ok(connector) => connector,
+        Err(e) => {
+            LogMessage::log(
+                LogLevel::ERROR,
+                format!("Failed to create TLS connector: {}", e),
+                0,
+            )
+            .await;
+            return None;
+        }
+    };
+
+    let timeout_duration = Duration::from_secs(10); // 10-second timeout
+
+    // The handshake future
+    let handshake_future = tls_connector.connect(domain, stream);
+
+    // Apply the timeout
+    let result = timeout(timeout_duration, handshake_future).await;
+
+    let tls = match result {
+        Ok(Ok(tls)) => tls,
+        _ => return None,
+    };
+
+    Some(Box::new(tls))
+}
diff --git a/client/src/utils/tls.rs b/client/src/utils/tls.rs
diff --git a/server/Cargo.toml b/server/Cargo.toml
@@ -12,9 +12,9 @@ common = { path = "../common" }
 config.workspace = true
 hex.workspace = true
 rsa.workspace = true
-rustls = "0.23.31"
-rustls-pemfile = "2.2.0"
 serde = { workspace = true, features = ["derive"] }
 tokio = { workspace = true, features = ["full"] }
-tokio-rustls = "0.26.2"
 uuid = { version = "1.18.0", features = ["v4"] }
+rustls.workspace = true
+rustls-pemfile = "2.2.0"
+tokio-rustls.workspace = true
diff --git a/server/Config.toml b/server/Config.toml
@@ -1,7 +1,6 @@
-host = "0.0.0.0"
 port = 8080
 
-[tls]
-enabled = false
-cert_path = "/etc/letsencrypt/live/viveksahani.com/fullchain.pem"
-key_path = "/etc/letsencrypt/live/viveksahani.com/privkey.pem"
+# [tls]
+# enabled = false
+# cert_path = "/etc/letsencrypt/live/viveksahani.com/fullchain.pem"
+# key_path = "/etc/letsencrypt/live/viveksahani.com/privkey.pem"
diff --git a/server/src/config.rs b/server/src/config.rs
@@ -4,21 +4,15 @@ use std::{env, error::Error, path::PathBuf};
 
 #[derive(Debug, Deserialize, Clone)]
 pub struct TLSConfig {
-    pub enabled: bool,
-    pub cert_path: Option<String>,
-    pub key_path: Option<String>,
+    pub cert_path: String,
+    pub key_path: String,
 }
 
 /// Configuration for the server
 #[derive(Debug, Deserialize, Clone)]
 pub struct ServerConfig {
-    /// Hostname or IP address of the server
-    pub host: String,
-
-    /// Port number for the server  
+    // pub domain: String,
     pub port: u16,
-
-    /// TLS configuration
     pub tls: Option<TLSConfig>
 }
 
@@ -52,6 +46,6 @@ impl ServerConfig {
 
     /// Get the server address as a string
     pub fn get_addr(&self) -> String {
-        format!("{}:{}", self.host, self.port)
+        format!("0.0.0.0:{}", self.port)
     }
 }
diff --git a/server/src/main.rs b/server/src/main.rs
@@ -2,14 +2,13 @@ use common::net::Packet;
 use null_talk_server::{
     ServerConfig,
     handlers::{handle_client, task::start_writer_task},
+    net::create_tls_acceptor,
 };
-use rustls::pki_types::PrivateKeyDer;
-use std::{fs, io::BufReader, sync::Arc};
+use std::sync::Arc;
 use tokio::{
     net::TcpListener,
     sync::{Mutex as AsyncMutex, mpsc},
 };
-use tokio_rustls::{TlsAcceptor, rustls::ServerConfig as RustlsServerConfig};
 
 /// Main entry point for the server
 #[tokio::main]
@@ -32,78 +31,44 @@ async fn main() {
 
     // TLS check
     if let Some(tls_cfg) = &config.tls {
-        if tls_cfg.enabled {
-            // Load certs
-            let certs = {
-                let file = fs::File::open(tls_cfg.cert_path.as_ref().unwrap())
-                    .expect("Cannot open cert file");
-                let mut reader = BufReader::new(file);
-                rustls_pemfile::certs(&mut reader)
-                    .into_iter()
-                    .map(|v| v.expect("Failed to read certificate"))
-                    .collect::<Vec<_>>()
-            };
-
-            let key = {
-                let file = fs::File::open(tls_cfg.key_path.as_ref().unwrap())
-                    .expect("Cannot open key file");
-                let mut reader = BufReader::new(file);
-                let first_key = rustls_pemfile::pkcs8_private_keys(&mut reader)
-                    .next()
-                    .ok_or("No private keys found")
-                    .expect("Failed to read private key")
-                    .expect("Failed to read private key");
-
-                PrivateKeyDer::Pkcs8(first_key)
-            };
-
-            let tls_config = RustlsServerConfig::builder()
-                .with_no_client_auth()
-                .with_single_cert(certs, key)
-                .expect("Failed to create TLS config");
-
-            let acceptor = TlsAcceptor::from(Arc::new(tls_config));
-            let listener = TcpListener::bind(&server_address).await.unwrap();
-            println!("🔒 TLS Server listening on {}", &server_address);
+        let acceptor = match create_tls_acceptor(tls_cfg).await {
+            Ok(acceptor) => acceptor,
+            Err(e) => {
+                eprintln!("Failed to create TLS acceptor: {:?}", e);
+                return;
+            }
+        };
+        let listener = TcpListener::bind(&server_address).await.unwrap();
+        println!("🔒 TLS Server listening on {}", &server_address);
 
-            loop {
-                match listener.accept().await {
-                    Ok((stream, _)) => {
-                        let acceptor = acceptor.clone();
-                        let sd_clone = sender.clone();
+        loop {
+            match listener.accept().await {
+                Ok((stream, _)) => {
+                    let acceptor = acceptor.clone();
+                    let sd_clone = sender.clone();
 
-                        tokio::spawn(async move {
-                            match acceptor.accept(stream).await {
-                                Ok(tls_stream) => {
-                                    handle_client(Box::new(tls_stream), sd_clone).await
-                                }
-                                Err(e) => eprintln!("TLS handshake failed: {:?}", e),
-                            }
-                        });
-                    }
-                    Err(e) => eprintln!("Failed to accept connection: {:?}", e),
+                    tokio::spawn(async move {
+                        match acceptor.accept(stream).await {
+                            Ok(tls_stream) => handle_client(Box::new(tls_stream), sd_clone).await,
+                            Err(e) => eprintln!("TLS handshake failed: {:?}", e),
+                        }
+                    });
                 }
+                Err(e) => eprintln!("Failed to accept connection: {:?}", e),
             }
-        } else {
-            start_plain_server(&server_address, sender).await;
         }
     } else {
-        start_plain_server(&server_address, sender).await;
-    }
-}
+        let listener = TcpListener::bind(&server_address).await.unwrap();
+        println!("🚀 Plain TCP Server listening on {}", &server_address);
 
-/// Helper to start plain TCP server
-async fn start_plain_server(addr: &str, sender: Arc<AsyncMutex<mpsc::UnboundedSender<Packet>>>) {
-    let listener = TcpListener::bind(addr).await.unwrap();
-    println!("🚀 Plain TCP Server listening on {}", addr);
-
-    loop {
-        match listener.accept().await {
-            Ok((stream, _)) => {
-                let sd_clone = sender.clone();
-                tokio::spawn(async move { handle_client(Box::new(stream), sd_clone).await });
+        loop {
+            match listener.accept().await {
+                Ok((stream, _)) => {
+                    let sd_clone = sender.clone();
+                    tokio::spawn(async move { handle_client(Box::new(stream), sd_clone).await });
+                }
+                Err(e) => eprintln!("Failed to accept connection: {:?}", e),
             }
-            Err(e) => eprintln!("Failed to accept connection: {:?}", e),
         }
     }
 }
diff --git a/server/src/net.rs b/server/src/net.rs
