Add encryption and decryption using AES-GCM and X25519 for secure client communication

Cargo.lock

@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "addr2line"
@@ -17,6 +17,41 @@ version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
 
+[[package]]
+name = "aead"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0"
+dependencies = [
+ "crypto-common",
+ "generic-array",
+]
+
+[[package]]
+name = "aes"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
+dependencies = [
+ "cfg-if",
+ "cipher",
+ "cpufeatures",
+]
+
+[[package]]
+name = "aes-gcm"
+version = "0.10.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1"
+dependencies = [
+ "aead",
+ "aes",
+ "cipher",
+ "ctr",
+ "ghash",
+ "subtle",
+]
+
 [[package]]
 name = "aho-corasick"
 version = "1.1.3"
@@ -96,12 +131,24 @@ dependencies = [
  "windows-targets 0.52.6",
 ]
 
+[[package]]
+name = "base64"
+version = "0.21.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567"
+
 [[package]]
 name = "bitflags"
 version = "2.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de"
 
+[[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
 [[package]]
 name = "bytes"
 version = "1.8.0"
@@ -118,12 +165,27 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 name = "chatserver"
 version = "0.1.0"
 dependencies = [
+ "aes-gcm",
+ "base64",
  "colog",
+ "crypto",
  "log",
+ "rand",
+ "rand_core",
  "tokio",
  "x25519-dalek",
 ]
 
+[[package]]
+name = "cipher"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
+dependencies = [
+ "crypto-common",
+ "inout",
+]
+
 [[package]]
 name = "colog"
 version = "1.3.0"
@@ -160,6 +222,35 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "crypto"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf1e6e5492f8f0830c37f301f6349e0dac8b2466e4fe89eef90e9eef906cd046"
+dependencies = [
+ "crypto-common",
+]
+
+[[package]]
+name = "crypto-common"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3"
+dependencies = [
+ "generic-array",
+ "rand_core",
+ "typenum",
+]
+
+[[package]]
+name = "ctr"
+version = "0.9.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835"
+dependencies = [
+ "cipher",
+]
+
 [[package]]
 name = "curve25519-dalek"
 version = "4.1.3"
@@ -215,6 +306,37 @@ version = "0.2.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d"
 
+[[package]]
+name = "generic-array"
+version = "0.14.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
+dependencies = [
+ "typenum",
+ "version_check",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi",
+]
+
+[[package]]
+name = "ghash"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1"
+dependencies = [
+ "opaque-debug",
+ "polyval",
+]
+
 [[package]]
 name = "gimli"
 version = "0.31.1"
@@ -233,6 +355,15 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
 
+[[package]]
+name = "inout"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5"
+dependencies = [
+ "generic-array",
+]
+
 [[package]]
 name = "is_terminal_polyfill"
 version = "1.70.1"
@@ -303,6 +434,12 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "opaque-debug"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.3"
@@ -332,6 +469,27 @@ version = "0.2.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "915a1e146535de9163f3987b8944ed8cf49a18bb0056bcebcdcece385cece4ff"
 
+[[package]]
+name = "polyval"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "opaque-debug",
+ "universal-hash",
+]
+
+[[package]]
+name = "ppv-lite86"
+version = "0.2.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04"
+dependencies = [
+ "zerocopy",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.92"
@@ -350,11 +508,35 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "rand"
+version = "0.8.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core",
+]
+
 [[package]]
 name = "rand_core"
 version = "0.6.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+dependencies = [
+ "getrandom",
+]
 
 [[package]]
 name = "redox_syscall"
@@ -512,18 +694,40 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "typenum"
+version = "1.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
 
+[[package]]
+name = "universal-hash"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea"
+dependencies = [
+ "crypto-common",
+ "subtle",
+]
+
 [[package]]
 name = "utf8parse"
 version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
 
+[[package]]
+name = "version_check"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
+
 [[package]]
 name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"
@@ -690,6 +894,27 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "zerocopy"
+version = "0.7.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
+dependencies = [
+ "byteorder",
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.7.35"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "zeroize"
 version = "1.8.1"

Cargo.toml

@@ -7,4 +7,9 @@ edition = "2021"
 colog = "1.3.0"
 log = "0.4.22"
 tokio = { version = "1.41.1", features = ["full"] }
-x25519-dalek = "2.0.1"
+x25519-dalek = "2.0.0-rc.3"
+aes-gcm = "0.10.3"
+rand = "0.8.5"
+rand_core = "0.6.4"
+crypto = "0.5.1"
+base64 = "0.21"

src/main.rs

@@ -1,56 +1,163 @@
-use log::{debug, error, info, trace, warn};
+use aes_gcm::{
-use std::env::args;
+    aead::{Aead, KeyInit, OsRng},
-use std::net::{TcpListener, TcpStream};
+    Aes256Gcm, Key, Nonce,
-
+};
-#[derive(Debug)]
+use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
-struct Settings {
+use colog;
-    host: String,
+use log::{error, info};
-    port: String,
+use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
-}
+use tokio::net::{TcpListener, TcpStream};
-
+use tokio::sync::broadcast;
-impl Settings {
+use x25519_dalek::{EphemeralSecret, PublicKey};
-    fn new(args: &[String]) -> Result<Settings, &'static str> {
-        if args.len() < 4 {
-            return Err("not enough arguments");
-        }
-
-        let port = args[2].clone();
-        let host = args[4].clone();
-
-        Ok(Settings { host, port })
-    }
-
-    fn get_full_host(&self) -> String {
-        format!("{}:{}", self.host, self.port)
-    }
-}
-
-async fn handle_client(mut stream: TcpStream) {
-    info!("Connected to {}", stream.peer_addr().unwrap());
-    stream.set_nodelay(true).unwrap();
-
-    loop {
-        
-    }
-}
 
 #[tokio::main]
 async fn main() {
+    // Initialize the logger
     colog::init();
-    info!("Starting...");
+
-    let args: Vec<String> = args().collect();
+    // Bind a TCP listener to accept incoming connections
-    let settings = Settings::new(&args).unwrap();
+    let listener = TcpListener::bind("0.0.0.0:8080").await.unwrap();
-    info!("Server Address: {}:{}", settings.host, settings.port);
+    info!("Server running on port 8080");
-    info!("Starting to listen to connections...");
+
-    let listener = TcpListener::bind(Settings::get_full_host(&settings)).unwrap();
+    // Create a broadcast channel for sharing messages
-    match listener.accept() {
+    let (tx, _) = broadcast::channel(100);
-        Ok((socket, _)) => {
+
-            tokio::spawn(async move {
+    loop {
-                handle_client(socket).await;
+        // Accept a new client
-            });
+        let (socket, addr) = listener.accept().await.unwrap();
-        }
+        info!("Client connected: {}", addr);
-        Err(e) => {
+
-            error!("Something went wrong {}", e);
+        let tx = tx.clone();
-        }
+        let rx = tx.subscribe();
+
+        // Handle the client in a new task
+        tokio::spawn(async move {
+            if let Err(e) = handle_client(socket, tx, rx).await {
+                error!("Error handling client: {}", e);
+            }
+        });
     }
 }
+
+async fn handle_client(
+    socket: TcpStream,
+    tx: broadcast::Sender<String>,
+    mut rx: broadcast::Receiver<String>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let (reader, mut writer) = socket.into_split();
+    let mut reader = BufReader::new(reader);
+    let mut line = String::new();
+
+    let server_secret = EphemeralSecret::random_from_rng(OsRng);
+    let server_public = PublicKey::from(&server_secret);
+
+    // Send the server's public key to the client
+    writer.write_all(server_public.as_bytes()).await?;
+
+    // Receive the client's public key
+    let mut client_public_bytes = [0u8; 32];
+    reader.read_exact(&mut client_public_bytes).await?;
+    let client_public = PublicKey::from(client_public_bytes);
+
+    // Compute the shared secret
+    let shared_secret = server_secret.diffie_hellman(&client_public);
+
+    let key = Key::<Aes256Gcm>::from_slice(shared_secret.as_bytes());
+
+    let cipher_reader = Aes256Gcm::new(&key);
+    let cipher_writer = Aes256Gcm::new(&key);
+    let nonce_reader = Nonce::from_slice(b"unique nonce"); // 96-bits; fixed nonce
+    let nonce_writer = nonce_reader.clone();
+
+    // Read task for receiving messages from the client
+    let read_task = tokio::spawn(async move {
+        loop {
+            line.clear();
+            match reader.read_line(&mut line).await {
+                Ok(bytes_read) => {
+                    if bytes_read == 0 {
+                        info!("Client disconnected");
+                        break;
+                    }
+
+                    let decoded = match BASE64.decode(line.trim().as_bytes()) {
+                        Ok(decoded) => decoded,
+                        Err(e) => {
+                            error!("Base64 decode error: {:?}", e);
+                            continue;
+                        }
+                    };
+
+                    let decrypted = match cipher_reader.decrypt(&nonce_reader, decoded.as_ref()) {
+                        Ok(decrypted) => decrypted,
+                        Err(e) => {
+                            error!("Decryption error: {:?}", e);
+                            continue;
+                        }
+                    };
+
+                    let message = match String::from_utf8(decrypted) {
+                        Ok(msg) => msg,
+                        Err(e) => {
+                            error!("UTF-8 conversion error: {:?}", e);
+                            continue;
+                        }
+                    };
+
+                    info!("Received message: {}", message.trim());
+
+                    if message.trim() == "/quit" {
+                        info!("Client requested to quit");
+                        break;
+                    }
+
+                    // Broadcast the message to all clients
+                    match tx.send(message) {
+                        Ok(_) => info!("Message broadcast successfully"),
+                        Err(e) => {
+                            error!("Failed to broadcast message: {:?}", e);
+                            break;
+                        }
+                    }
+                }
+                Err(e) => {
+                    error!("Error reading from client: {:?}", e);
+                    break;
+                }
+            }
+        }
+    });
+
+    // Write task for sending messages to the client
+    let write_task = tokio::spawn(async move {
+        while let Ok(msg) = rx.recv().await {
+            if !msg.is_empty() {
+                // Encrypt the message
+                let encrypted = match cipher_writer.encrypt(&nonce_writer, msg.as_bytes()) {
+                    Ok(encrypted) => encrypted,
+                    Err(e) => {
+                        error!("Encryption error: {:?}", e);
+                        continue;
+                    }
+                };
+
+                // Base64 encode the encrypted message
+                let encoded = BASE64.encode(&encrypted);
+
+                if let Err(e) = writer.write_all((encoded + "\n").as_bytes()).await {
+                    error!("Failed to send message: {:?}", e);
+                    break;
+                }
+            }
+        }
+    });
+
+    // Wait for both tasks to complete
+    tokio::select! {
+        _ = read_task => (),
+        _ = write_task => (),
+    }
+
+    info!("Client handling completed");
+    Ok(())
+}