feat: Integrate TLS Certificate to the server and create a generic trait for tcp stream

Author: ByteMaster2003

Cargo.lock

@@ -29,8 +29,9 @@ use crate::{
     types::{LogLevel, LogMessage},
     utils::perform_handshake,
 };
+use common::net::AsyncStream;
 use std::sync::Arc;
-use tokio::{net::TcpStream, sync::Mutex};
+use tokio::sync::Mutex;
 
 /// Handles an incoming client connection.
 ///
@@ -66,8 +67,8 @@ use tokio::{net::TcpStream, sync::Mutex};
 /// [`TcpStream`]: tokio::net::TcpStream
 /// [`perform_handshake`]: crate::utils::perform_handshake
 /// [`LogMessage`]: crate::types::LogMessage
-pub async fn handle_client(stream: TcpStream) {
-    let (rd, wt) = stream.into_split();
+pub async fn handle_client(stream: Box<dyn AsyncStream>) {
+    let (rd, wt) = tokio::io::split(stream);
     let rd = Arc::new(Mutex::new(rd));
     let wt = Arc::new(Mutex::new(wt));
 

client/src/handlers/client.rs

@@ -39,7 +39,7 @@ async fn main() {
                 )
                 .await;
 
-                handle_client(stream).await;
+                handle_client(Box::new(stream)).await;
             }
             Err(_) => {
                 LogMessage::log(LogLevel::ERROR, format!("Failed to connect to {}", addr), 0).await;

client/src/main.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use bincode::{Decode, Encode};
 use tokio::{
-    net::tcp::{OwnedReadHalf, OwnedWriteHalf},
+    io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf},
     sync::Mutex,
 };
 
@@ -46,7 +46,14 @@ pub struct HandshakePacket {
     pub session_key: Option<Vec<u8>>,
 }
 
+/// Custom trait that bundles AsyncRead + AsyncWrite
+pub trait AsyncStream: AsyncRead + AsyncWrite + Unpin + Send {}
+impl<T: AsyncRead + AsyncWrite + Unpin + Send> AsyncStream for T {}
+
 /// Represents a stream reader for a client
-pub type StreamReader = Arc<Mutex<OwnedReadHalf>>;
+// pub type StreamReader<S> = Arc<Mutex<OwnedReadHalf<S>>>;
+pub type StreamReader = Arc<Mutex<ReadHalf<Box<dyn AsyncStream>>>>;
 /// Represents a stream writer for a client
-pub type StreamWriter = Arc<Mutex<OwnedWriteHalf>>;
+// pub type StreamWriter = Arc<Mutex<OwnedWriteHalf>>;
+pub type StreamWriter = Arc<Mutex<WriteHalf<Box<dyn AsyncStream>>>>;
+

common/src/net.rs

@@ -12,6 +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"] }

server/Cargo.toml

@@ -1,2 +1,7 @@
 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"

server/Config.toml

@@ -2,6 +2,13 @@ use config::{Config, File};
 use serde::Deserialize;
 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>,
+}
+
 /// Configuration for the server
 #[derive(Debug, Deserialize, Clone)]
 pub struct ServerConfig {
@@ -10,6 +17,9 @@ pub struct ServerConfig {
 
     /// Port number for the server  
     pub port: u16,
+
+    /// TLS configuration
+    pub tls: Option<TLSConfig>
 }
 
 impl ServerConfig {

server/src/config.rs

@@ -1,20 +1,23 @@
 use std::sync::Arc;
 
-use tokio::{
-    net::TcpStream,
-    sync::{Mutex as AsyncMutex, mpsc::UnboundedSender},
-};
+use tokio::sync::{Mutex as AsyncMutex, mpsc::UnboundedSender};
 
 use crate::{
-    data::CLIENTS, handlers::task::start_reader_task, types::Client, net::perform_handshake,
+    data::CLIENTS, handlers::task::start_reader_task, net::perform_handshake, types::Client,
+};
+use common::{
+    net::{AsyncStream, Packet, StreamReader, StreamWriter},
+    utils::enc::public_key_to_user_id,
 };
-use common::{net::Packet, utils::enc::public_key_to_user_id};
 
 /// Handle a new client connection
-pub async fn handle_client(stream: TcpStream, tx: Arc<AsyncMutex<UnboundedSender<Packet>>>) {
-    let (rd, wt) = stream.into_split();
-    let rd = Arc::new(AsyncMutex::new(rd));
-    let wt = Arc::new(AsyncMutex::new(wt));
+pub async fn handle_client(
+    stream: Box<dyn AsyncStream>,
+    tx: Arc<AsyncMutex<UnboundedSender<Packet>>>,
+) {
+    let (rd, wt) = tokio::io::split(stream);
+    let rd: StreamReader = Arc::new(AsyncMutex::new(rd));
+    let wt: StreamWriter = Arc::new(AsyncMutex::new(wt));
 
     let (name, session_key, public_key) = match perform_handshake(rd.clone(), wt.clone()).await {
         Ok(data) => data,

server/src/handlers/client.rs

@@ -1,7 +1,6 @@
+use crate::data;
 use common::{net::Packet, types::Message, utils::net::write_packet};
 
-use crate::data::{self, CONVERSATIONS};
-
 /// Handle a group message
 /// Send the message to every active member of the group
 pub async fn handle_group_message(packet: Packet, group_id: &str) {
@@ -35,7 +34,7 @@ pub async fn handle_group_message(packet: Packet, group_id: &str) {
 // Handle a direct message
 pub async fn handle_direct_message(packet: Packet, session_id: &str) {
     // Find the conversation
-    let dm = match CONVERSATIONS.lock().await.get(session_id) {
+    let dm = match data::CONVERSATIONS.lock().await.get(session_id) {
         Some(dm) => dm.clone(),
         None => return,
     };

server/src/handlers/msg.rs

@@ -1,14 +1,15 @@
-use std::sync::Arc;
-
 use common::net::Packet;
 use null_talk_server::{
     ServerConfig,
     handlers::{handle_client, task::start_writer_task},
 };
+use rustls::pki_types::PrivateKeyDer;
+use std::{fs, io::BufReader, 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]
@@ -23,23 +24,86 @@ async fn main() {
     println!("🔧 Configuration Loaded");
 
     let server_address = config.get_addr();
-    let listener = TcpListener::bind(&server_address).await.unwrap();
-    println!("🚀 Server listening on {}", &server_address);
 
+    // Shared channel for communication
     let (tx, rx) = mpsc::unbounded_channel::<Packet>();
     let _ = start_writer_task(rx).await;
-
     let sender: Arc<AsyncMutex<mpsc::UnboundedSender<Packet>>> = Arc::new(AsyncMutex::new(tx));
+
+    // 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);
+
+            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),
+                }
+            }
+        } else {
+            start_plain_server(&server_address, sender).await;
+        }
+    } else {
+        start_plain_server(&server_address, sender).await;
+    }
+}
+
+/// 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(stream, sd_clone).await });
+                tokio::spawn(async move { handle_client(Box::new(stream), sd_clone).await });
             }
-            Err(e) => {
-                eprintln!("Failed to accept connection: {:?}", e);
-                continue;
-            }
-        };
+            Err(e) => eprintln!("Failed to accept connection: {:?}", e),
+        }
     }
 }