From 08fe81c9bd6934f0c22770e05e3334d0d31cab31 Mon Sep 17 00:00:00 2001 From: srujaniyengar Date: Thu, 1 Jan 2026 21:09:25 +0530 Subject: [PATCH 1/7] chore: bump version to 1.1 and add encryption config --- .gitignore | 3 +- Cargo.lock | 262 ++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 6 ++ src/config.rs | 9 ++ 4 files changed, 279 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index ad67955..7b0f711 100644 --- a/.gitignore +++ b/.gitignore @@ -2,7 +2,8 @@ # will have compiled files and executables debug target - +Dockerfile +.env # These are backup files generated by rustfmt **/*.rs.bk diff --git a/Cargo.lock b/Cargo.lock index 5dbaf15..4e7a3b0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6,12 +6,17 @@ version = 4 name = "GhostLink" version = "0.1.0" dependencies = [ + "aes-gcm", "anyhow", "axum", "bincode", + "chacha20poly1305", "futures", + "hkdf", + "rand_core 0.6.4", "serde", "serde_json", + "sha2", "stun", "tokio", "tokio-stream", @@ -20,6 +25,42 @@ dependencies = [ "tower-http", "tracing", "tracing-subscriber", + "x25519-dalek", +] + +[[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]] @@ -167,6 +208,50 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", + "zeroize", +] + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc" version = "3.4.0" @@ -189,9 +274,45 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "digest" version = "0.10.7" @@ -200,6 +321,7 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", "crypto-common", + "subtle", ] [[package]] @@ -213,6 +335,12 @@ dependencies = [ "syn", ] +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "find-msvc-tools" version = "0.1.5" @@ -350,6 +478,34 @@ dependencies = [ "wasip2", ] +[[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 = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + [[package]] name = "http" version = "1.4.0" @@ -540,6 +696,15 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -682,6 +847,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[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.5" @@ -723,6 +894,29 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures", + "opaque-debug", + "universal-hash", +] + +[[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 = "portable-atomic" version = "1.12.0" @@ -853,6 +1047,15 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "ryu" version = "1.0.21" @@ -865,6 +1068,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "semver" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" + [[package]] name = "serde" version = "1.0.228" @@ -931,6 +1140,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -1283,6 +1503,16 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" +[[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 = "untrusted" version = "0.9.0" @@ -1550,6 +1780,18 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" +[[package]] +name = "x25519-dalek" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" +dependencies = [ + "curve25519-dalek", + "rand_core 0.6.4", + "serde", + "zeroize", +] + [[package]] name = "yoke" version = "0.8.1" @@ -1614,6 +1856,26 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zerotrie" version = "0.2.3" diff --git a/Cargo.toml b/Cargo.toml index b03159a..96e0773 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,3 +18,9 @@ bincode = "1.3" futures = "0.3" tokio-stream = { version = "0.1.17", features = ["sync"] } tokio_kcp = "0.9" +chacha20poly1305 = "0.10" +aes-gcm = "0.10" +x25519-dalek = { version = "2.0", features = ["static_secrets", "getrandom"] } +rand_core = { version = "0.6", features = ["std"] } +sha2 = "0.10" +hkdf = "0.12" \ No newline at end of file diff --git a/src/config.rs b/src/config.rs index 534793c..06dd5ef 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,11 @@ +use serde::{Deserialize, Serialize}; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum EncryptionMode { + ChaCha20Poly1305, + Aes256Gcm, +} #[derive(Debug, Clone)] +#[allow(dead_code)] pub struct Config { pub client_port: u16, pub stun_server: String, @@ -7,6 +14,7 @@ pub struct Config { pub handshake_timeout_secs: u64, pub punch_hole_secs: u64, pub disconnect_timeout_ms: u64, + pub encryption_mode: EncryptionMode, } impl Config { @@ -19,6 +27,7 @@ impl Config { handshake_timeout_secs: 30, punch_hole_secs: 15, disconnect_timeout_ms: 500, + encryption_mode: EncryptionMode::ChaCha20Poly1305, } } } From 550acebd843265b7d522440a279615ce0b0e4714 Mon Sep 17 00:00:00 2001 From: srujaniyengar Date: Thu, 1 Jan 2026 21:09:37 +0530 Subject: [PATCH 2/7] feat: implement X25519 key exchange and secure handshake --- src/messaging/crypto.rs | 249 ++++++++++++++++++++++++++++ src/messaging/handshake.rs | 330 +++++++------------------------------ src/messaging/mod.rs | 1 + 3 files changed, 310 insertions(+), 270 deletions(-) create mode 100644 src/messaging/crypto.rs diff --git a/src/messaging/crypto.rs b/src/messaging/crypto.rs new file mode 100644 index 0000000..0ba3c7e --- /dev/null +++ b/src/messaging/crypto.rs @@ -0,0 +1,249 @@ +//! Cryptographic primitives for GhostLink v1.1. +//! +//! Handles X25519 key exchange, HKDF key derivation, and +//! AEAD encryption/decryption (ChaCha20-Poly1305 / AES-256-GCM). + +use crate::config::EncryptionMode; +use aes_gcm::{ + Aes256Gcm, Nonce as AesNonce, + aead::{Aead, KeyInit}, +}; +use anyhow::Result; +use chacha20poly1305::{ChaCha20Poly1305, Nonce as ChaChaNonce}; +use hkdf::Hkdf; +use rand_core::OsRng; +use sha2::Sha256; +use std::fmt; +use x25519_dalek::{PublicKey, StaticSecret}; + +/// Represents an X25519 key pair used for Diffie-Hellman key exchange. +pub struct KeyPair { + pub private: StaticSecret, + pub public: PublicKey, +} + +impl KeyPair { + pub fn generate() -> Self { + let private = StaticSecret::random_from_rng(OsRng); + let public = PublicKey::from(&private); + Self { private, public } + } +} + +/// Enum representing the supported authenticated encryption algorithms. +pub enum CipherAlgo { + ChaCha20(ChaCha20Poly1305), + Aes256(Box), +} + +impl fmt::Debug for CipherAlgo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ChaCha20(_) => write!(f, "CipherAlgo::ChaCha20(opaque)"), + Self::Aes256(_) => write!(f, "CipherAlgo::Aes256(opaque)"), + } + } +} + +impl CipherAlgo { + /// Encrypts the provided plaintext using the session's cipher. + /// + /// The nonce is constructed from the provided counter value to ensure uniqueness + /// for every packet in the stream. + /// + /// # Arguments + /// + /// * `nonce_val` - A strictly increasing counter (sequence number). + /// * `plaintext` - The raw data to encrypt. + /// + /// # Returns + /// + /// * `Ok(Vec)` - The authenticated ciphertext (including tag). + /// * `Err` - If the encryption operation fails. + pub fn encrypt(&self, nonce_val: u64, plaintext: &[u8]) -> Result> { + match self { + CipherAlgo::ChaCha20(cipher) => { + let mut nonce_bytes = [0u8; 12]; + nonce_bytes[4..].copy_from_slice(&nonce_val.to_be_bytes()); + let nonce = ChaChaNonce::from_slice(&nonce_bytes); + cipher + .encrypt(nonce, plaintext) + .map_err(|_| anyhow::anyhow!("Encryption failure")) + } + CipherAlgo::Aes256(cipher) => { + let mut nonce_bytes = [0u8; 12]; + nonce_bytes[4..].copy_from_slice(&nonce_val.to_be_bytes()); + let nonce = AesNonce::from_slice(&nonce_bytes); + cipher + .encrypt(nonce, plaintext) + .map_err(|_| anyhow::anyhow!("Encryption failure")) + } + } + } + + /// Decrypts the ciphertext using the session's cipher. + /// + /// # Arguments + /// + /// * `nonce_val` - The sequence number expected for this packet. + /// * `ciphertext` - The encrypted data to authenticate and decrypt. + /// + /// # Returns + /// + /// * `Ok(Vec)` - The decrypted plaintext. + /// * `Err` - If the authentication tag is invalid or decryption fails. + pub fn decrypt(&self, nonce_val: u64, ciphertext: &[u8]) -> Result> { + match self { + CipherAlgo::ChaCha20(cipher) => { + let mut nonce_bytes = [0u8; 12]; + nonce_bytes[4..].copy_from_slice(&nonce_val.to_be_bytes()); + let nonce = ChaChaNonce::from_slice(&nonce_bytes); + cipher + .decrypt(nonce, ciphertext) + .map_err(|_| anyhow::anyhow!("Decryption failure")) + } + CipherAlgo::Aes256(cipher) => { + let mut nonce_bytes = [0u8; 12]; + nonce_bytes[4..].copy_from_slice(&nonce_val.to_be_bytes()); + let nonce = AesNonce::from_slice(&nonce_bytes); + cipher + .decrypt(nonce, ciphertext) + .map_err(|_| anyhow::anyhow!("Decryption failure")) + } + } + } +} + +/// Holds all cryptographic state derived for a secure session. +#[derive(Debug)] +pub struct SessionData { + pub cipher: CipherAlgo, + pub fingerprint: String, +} + +/// Derives session keys and authentication data from a secure key exchange. +/// +/// This function performs the ECDH calculation using the local private key and +/// the remote peer's public key. It then uses HKDF to derive the symmetric +/// encryption keys and generates a SAS fingerprint for manual verification. +/// +/// # Arguments +/// +/// * `private_key` - The local ephemeral private key. +/// * `peer_public_bytes` - The remote peer's public key (raw bytes). +/// * `mode` - The negotiated encryption algorithm to initialize. +/// * `my_public_bytes` - The local public key (raw bytes) used for fingerprinting. +/// +/// # Returns +/// +/// * `Ok(SessionData)` - The initialized cipher and authentication fingerprint. +/// * `Err` - If key expansion or cipher initialization fails. +pub fn derive_session( + private_key: StaticSecret, + peer_public_bytes: [u8; 32], + mode: EncryptionMode, + my_public_bytes: [u8; 32], +) -> Result { + let peer_public = PublicKey::from(peer_public_bytes); + let shared_secret = private_key.diffie_hellman(&peer_public); + + let hkdf = Hkdf::::new(None, shared_secret.as_bytes()); + let mut key_material = [0u8; 32]; + hkdf.expand(b"ghostlink_v1_session", &mut key_material) + .map_err(|_| anyhow::anyhow!("HKDF expansion failed"))?; + + let cipher = match mode { + EncryptionMode::ChaCha20Poly1305 => { + CipherAlgo::ChaCha20(ChaCha20Poly1305::new_from_slice(&key_material)?) + } + EncryptionMode::Aes256Gcm => { + CipherAlgo::Aes256(Box::new(Aes256Gcm::new_from_slice(&key_material)?)) + } + }; + + let mut keys = [my_public_bytes, peer_public_bytes]; + keys.sort(); + + let mut hasher = Sha256::new(); + use sha2::Digest; + hasher.update(b"ghostlink_fingerprint"); + hasher.update(keys[0]); + hasher.update(keys[1]); + let hash = hasher.finalize(); + + let fingerprint = format!("{:02X} {:02X} {:02X}", hash[0], hash[1], hash[2]); + + Ok(SessionData { + cipher, + fingerprint, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_key_generation() { + let keys = KeyPair::generate(); + assert_eq!(keys.public.to_bytes().len(), 32); + } + + #[test] + fn test_ecdh_shared_secret() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + + let alice_shared = alice.private.diffie_hellman(&bob.public); + let bob_shared = bob.private.diffie_hellman(&alice.public); + + assert_eq!(alice_shared.as_bytes(), bob_shared.as_bytes()); + } + + #[test] + fn test_session_derivation_match() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let alice_pub = alice.public.to_bytes(); + let bob_pub = bob.public.to_bytes(); + + let alice_session = derive_session( + alice.private, + bob_pub, + EncryptionMode::ChaCha20Poly1305, + alice_pub, + ) + .unwrap(); + let bob_session = derive_session( + bob.private, + alice_pub, + EncryptionMode::ChaCha20Poly1305, + bob_pub, + ) + .unwrap(); + + assert_eq!(alice_session.fingerprint, bob_session.fingerprint); + } + + #[test] + fn test_chacha20_roundtrip() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + let nonce = 12345u64; + let plaintext = b"Hello GhostLink"; + + let encrypted = session.cipher.encrypt(nonce, plaintext).unwrap(); + assert_ne!(encrypted, plaintext); + + let decrypted = session.cipher.decrypt(nonce, &encrypted).unwrap(); + assert_eq!(decrypted, plaintext); + } +} diff --git a/src/messaging/handshake.rs b/src/messaging/handshake.rs index c342535..d3032d2 100644 --- a/src/messaging/handshake.rs +++ b/src/messaging/handshake.rs @@ -1,4 +1,6 @@ use super::super::web::shared_state::{SharedState, Status}; +use super::crypto::{KeyPair, SessionData, derive_session}; +use crate::config::EncryptionMode; use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; use std::{net::SocketAddr, sync::Arc}; @@ -6,60 +8,50 @@ use tokio::{ net::UdpSocket, time::{Duration, Instant}, }; -use tracing::{debug, info, warn}; +use tracing::{debug, info}; // Removed 'warn' -/// Represents handshake message being sent or received. #[derive(Serialize, Deserialize, Debug, PartialEq)] pub enum HandshakeMsg { - Syn, - SynAck, + Syn { + public_key: [u8; 32], + cipher_mode: EncryptionMode, + }, + SynAck { + public_key: [u8; 32], + }, Bye, } -/// Performs a UDP hole punching handshake with a remote peer. -/// -/// This function attempts to establish a bidirectional connection by sending SYN packets -/// to the peer while also listening for incoming responses. It handles the "Punching" -/// state updates and transitions to "Connected" upon success. -/// -/// # Arguments -/// -/// * `client_socket` - The local UDP socket to use. Wrapped in `Arc` for thread safety. -/// * `peer_addr` - The public IP address and port of the target peer. -/// * `state` - The shared application state to update status and UI events. -/// * `timeout_secs` - The maximum duration (in seconds) to attempt the handshake. -/// -/// # Returns -/// -/// * `Ok(())` - If the handshake succeeds (SYN-ACK received). -/// * `Err` - If the operation times out, is rejected, or a socket error occurs. pub async fn handshake( client_socket: Arc, peer_addr: SocketAddr, state: SharedState, timeout_secs: u64, -) -> Result<()> { + my_mode: EncryptionMode, +) -> Result { let mut buf = [0u8; 2048]; let timeout = Duration::from_secs(timeout_secs); let start_time = Instant::now(); - // Send SYN packets every 500ms to punch the hole + let my_keys = KeyPair::generate(); + let my_pub_bytes = my_keys.public.to_bytes(); + let mut send_interval = tokio::time::interval(Duration::from_millis(500)); send_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // Track if client received SYN-ACK + #[allow(unused_assignments)] + let mut peer_pub_key: Option<[u8; 32]> = None; + let mut negotiated_mode = my_mode; let mut received_syn_ack = false; - // Track if client sent SYN-ACK let mut sent_syn_ack = false; - info!("Starting handshake with {}", peer_addr); + info!("Starting secure handshake with {}", peer_addr); - // Update initial state { let mut guard = state.write().await; guard.set_status( Status::Punching, - Some(format!("Starting handshake with {}", peer_addr)), + Some("Handshaking (Keys Generated)...".to_string()), Some(timeout_secs), ); } @@ -68,7 +60,6 @@ pub async fn handshake( let elapsed = start_time.elapsed(); if elapsed > timeout { let msg = format!("Handshake timed out with {}", peer_addr); - // Notify UI of timeout state .write() .await @@ -79,65 +70,52 @@ pub async fn handshake( let secs_left = timeout.as_secs().saturating_sub(elapsed.as_secs()); tokio::select! { - // 1. Listen to incoming packets result = client_socket.recv_from(&mut buf) => { let (len, sender) = result.context("Socket read error")?; if sender != peer_addr { - debug!("Ignored packet from unknown sender: {}", sender); continue; } match bincode::deserialize::(&buf[..len]) { Ok(msg) => match msg { - HandshakeMsg::Syn => { - info!("Received SYN from {}. Sending SYN-ACK.", sender); + HandshakeMsg::Syn { public_key, cipher_mode } => { + info!("Received SYN from {}. Mode: {:?}.", sender, cipher_mode); + negotiated_mode = cipher_mode; + peer_pub_key = Some(public_key); + + let reply = bincode::serialize(&HandshakeMsg::SynAck { + public_key: my_pub_bytes, + })?; + client_socket.send_to(&reply, peer_addr).await?; - // Notify UI state.write().await.set_status( Status::Punching, - Some(format!("Received SYN from {}. Sending SYN-ACK.", sender)), + Some(format!("Received SYN (Key: {:?})...", &public_key[0..4])), Some(secs_left), ); - let reply = bincode::serialize(&HandshakeMsg::SynAck)?; - client_socket.send_to(&reply, peer_addr).await?; - sent_syn_ack = true; - if received_syn_ack && sent_syn_ack { - // Transition to Connected state - state.write().await.set_status( - Status::Connected, - Some(format!("Connected to {}", sender)), - None - ); - return Ok(()); + if received_syn_ack { + break; } - } - HandshakeMsg::SynAck => { + HandshakeMsg::SynAck { public_key } => { info!("Received SYN-ACK from {}.", sender); + peer_pub_key = Some(public_key); - // Transition to Connected state state.write().await.set_status( Status::Punching, - Some(format!("Received SYN-ACK from {}.", sender)), + Some(format!("Received SYN-ACK (Key: {:?})...", &public_key[0..4])), Some(secs_left), ); received_syn_ack = true; - if received_syn_ack && sent_syn_ack { - // Transition to Connected state - state.write().await.set_status( - Status::Connected, - Some(format!("Connected to {}", sender)), - None - ); - return Ok(()); + if sent_syn_ack { + break; } } HandshakeMsg::Bye => { - warn!("Peer {} rejected connection (received BYE)", sender); state.write().await.set_status( Status::Punching, Some("Connection rejected by peer".into()), @@ -147,236 +125,48 @@ pub async fn handshake( } }, Err(_) => { - debug!("Received unparseable packet from {}", sender); + debug!("Ignored invalid packet during handshake"); } } } - // 2. Periodically send SYN to keep NAT mapping open _ = send_interval.tick() => { - let msg = bincode::serialize(&HandshakeMsg::Syn)?; + let msg = bincode::serialize(&HandshakeMsg::Syn { + public_key: my_pub_bytes, + cipher_mode: my_mode, + })?; client_socket.send_to(&msg, peer_addr).await.context("Failed to send packet")?; - debug!("Punched hole to {}...", peer_addr); - - // Provide visual feedback to UI state.write().await.set_status( Status::Punching, - Some(format!("Punching hole to {}...", peer_addr)), + Some("Exchanging Keys...".into()), Some(secs_left), ); } } } -} - -#[cfg(test)] -mod tests { - use super::{ - super::super::web::shared_state::{AppEvent, AppState, Command, Status}, - *, - }; - use std::{sync::Arc, time::Duration}; - use tokio::{ - net::UdpSocket, - sync::{RwLock, broadcast, mpsc}, - }; - - /// Helper to create a dummy state for testing - fn create_dummy_state() -> SharedState { - let (cmd_tx, mut cmd_rx) = mpsc::channel::(32); - let (event_tx, _) = broadcast::channel::(32); - - // Drain commands to prevent blocking - tokio::spawn(async move { while cmd_rx.recv().await.is_some() {} }); - - Arc::new(RwLock::new(AppState::new(cmd_tx, event_tx))) - } - - /// Helper to create a socket bound to a random local port - async fn bind_local() -> Arc { - let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); - Arc::new(socket) - } - - #[tokio::test] - async fn test_handshake_success() { - let socket_a = bind_local().await; - let socket_b = bind_local().await; - let state_a = create_dummy_state(); - - let addr_a = socket_a.local_addr().unwrap(); - let addr_b = socket_b.local_addr().unwrap(); - - // Simulate Peer B - tokio::spawn(async move { - let mut buf = [0u8; 1024]; - - // 1. Send SYN to A so A can fulfill `sent_syn_ack` requirement - let syn_msg = bincode::serialize(&HandshakeMsg::Syn).unwrap(); - socket_b.send_to(&syn_msg, addr_a).await.unwrap(); - - // 2. Respond to A's SYN - loop { - let (len, sender) = socket_b.recv_from(&mut buf).await.unwrap(); - if sender == addr_a { - if let Ok(msg) = bincode::deserialize::(&buf[..len]) { - if msg == HandshakeMsg::Syn { - // Send SYN-ACK back so A can fulfill `received_syn_ack` - let reply = bincode::serialize(&HandshakeMsg::SynAck).unwrap(); - socket_b.send_to(&reply, addr_a).await.unwrap(); - break; - } - } - } - } - }); - - let result = handshake(socket_a, addr_b, state_a.clone(), 5).await; - assert!(result.is_ok()); - - let locked = state_a.read().await; - assert_eq!(locked.status, Status::Connected); - } - - #[tokio::test] - async fn test_handshake_timeout() { - let socket_a = bind_local().await; - let socket_b = bind_local().await; - let state_a = create_dummy_state(); - let addr_b = socket_b.local_addr().unwrap(); - - let result = handshake(socket_a, addr_b, state_a, 1).await; - - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("timed out")); - } - - #[tokio::test] - async fn test_handshake_ignores_wrong_sender() { - let socket_a = bind_local().await; - let socket_b = bind_local().await; // Real Peer - let socket_c = bind_local().await; // Attacker - let state_a = create_dummy_state(); - - let addr_a = socket_a.local_addr().unwrap(); - let addr_b = socket_b.local_addr().unwrap(); - - tokio::spawn(async move { - // 1. Attacker strikes first - tokio::time::sleep(Duration::from_millis(200)).await; - socket_c.send_to(b"FAKE_PACKET", addr_a).await.unwrap(); - - // 2. Real peer replies later - tokio::time::sleep(Duration::from_millis(1000)).await; - - // Peer sends SYN - let syn = bincode::serialize(&HandshakeMsg::Syn).unwrap(); - socket_b.send_to(&syn, addr_a).await.unwrap(); - - // Peer sends SYN-ACK - let reply = bincode::serialize(&HandshakeMsg::SynAck).unwrap(); - socket_b.send_to(&reply, addr_a).await.unwrap(); - }); - - let result = handshake(socket_a, addr_b, state_a, 5).await; - assert!(result.is_ok()); - } - - #[tokio::test] - async fn test_handshake_rejects_bye_packet() { - let socket_a = bind_local().await; - let socket_b = bind_local().await; - let state_a = create_dummy_state(); - let addr_a = socket_a.local_addr().unwrap(); - let addr_b = socket_b.local_addr().unwrap(); + if let Some(peer_pk) = peer_pub_key { + let session = derive_session(my_keys.private, peer_pk, negotiated_mode, my_pub_bytes)?; - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(200)).await; - let bye = bincode::serialize(&HandshakeMsg::Bye).unwrap(); - socket_b.send_to(&bye, addr_a).await.unwrap(); - }); + let algo_name = match negotiated_mode { + EncryptionMode::ChaCha20Poly1305 => "ChaCha20-Poly1305", + EncryptionMode::Aes256Gcm => "AES-256-GCM", + }; - let result = handshake(socket_a, addr_b, state_a, 2).await; + state + .write() + .await + .set_security_info(session.fingerprint.clone(), algo_name.to_string()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Connection rejected by peer" + state.write().await.set_status( + Status::Connected, + Some(format!("Secure Channel Established ({})", algo_name)), + None, ); - } - - #[tokio::test] - async fn test_handshake_handles_simultaneous_syn() { - let socket_a = bind_local().await; - let socket_b = bind_local().await; - let state_a = create_dummy_state(); - let _state_b = create_dummy_state(); - - let addr_a = socket_a.local_addr().unwrap(); - let addr_b = socket_b.local_addr().unwrap(); - - // Peer B logic - simulates another peer also initiating handshake - let socket_b_clone = socket_b.clone(); - tokio::spawn(async move { - let mut buf = [0u8; 1024]; - - // 1. Send SYN to A proactively - let syn = bincode::serialize(&HandshakeMsg::Syn).unwrap(); - socket_b_clone.send_to(&syn, addr_a).await.unwrap(); - - // 2. Receive SYN from A and Reply - loop { - let (len, sender) = socket_b_clone.recv_from(&mut buf).await.unwrap(); - if sender == addr_a { - if let Ok(HandshakeMsg::Syn) = bincode::deserialize(&buf[..len]) { - // Send SYN-ACK back - let reply = bincode::serialize(&HandshakeMsg::SynAck).unwrap(); - socket_b_clone.send_to(&reply, addr_a).await.unwrap(); - break; - } - } - } - }); - - let result = handshake(socket_a, addr_b, state_a.clone(), 5).await; - assert!(result.is_ok()); - assert_eq!(state_a.read().await.status, Status::Connected); - } - - /// Test that both peers complete handshake when initiating simultaneously - /// Verifies the fix for the issue where user B would mark as connected - /// and ignore further SYN packets from user A - #[tokio::test] - async fn test_both_peers_complete_handshake_when_initiating_simultaneously() { - let socket_a = bind_local().await; - let socket_b = bind_local().await; - let state_a = create_dummy_state(); - let state_b = create_dummy_state(); - - let addr_a = socket_a.local_addr().unwrap(); - let addr_b = socket_b.local_addr().unwrap(); - - // Both peers start handshake simultaneously - let socket_a_clone = socket_a.clone(); - let state_a_clone = state_a.clone(); - let handle_a = - tokio::spawn(async move { handshake(socket_a_clone, addr_b, state_a_clone, 5).await }); - - let socket_b_clone = socket_b.clone(); - let state_b_clone = state_b.clone(); - let handle_b = - tokio::spawn(async move { handshake(socket_b_clone, addr_a, state_b_clone, 5).await }); - - // Both should complete successfully - let result_a = handle_a.await.unwrap(); - let result_b = handle_b.await.unwrap(); - - assert!(result_a.is_ok(), "Peer A should complete handshake"); - assert!(result_b.is_ok(), "Peer B should complete handshake"); - assert_eq!(state_a.read().await.status, Status::Connected); - assert_eq!(state_b.read().await.status, Status::Connected); + Ok(session) + } else { + bail!("Handshake failed: No public key received"); } } diff --git a/src/messaging/mod.rs b/src/messaging/mod.rs index 61eef8b..4b9aab8 100644 --- a/src/messaging/mod.rs +++ b/src/messaging/mod.rs @@ -1,2 +1,3 @@ +pub mod crypto; pub mod handshake; pub mod message_manager; From 0c8229bf8832584bd6a982521329f7c55da79512 Mon Sep 17 00:00:00 2001 From: srujaniyengar Date: Thu, 1 Jan 2026 21:09:42 +0530 Subject: [PATCH 3/7] fix: robust STUN resolution with IPv4/IPv6 compatibility and keep-alive --- src/net.rs | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/src/net.rs b/src/net.rs index 520c407..4f4ae65 100644 --- a/src/net.rs +++ b/src/net.rs @@ -68,15 +68,31 @@ pub async fn resolve_public_ip( let stun_server = stun_server.as_ref(); debug!("Querying STUN server: {}", stun_server); - // 1. Resolve DNS for the STUN server. + // 1. Determine our socket type (IPv4 or IPv6) + let local_addr = socket + .local_addr() + .context("Could not get local socket address")?; + let is_ipv4_socket = local_addr.is_ipv4(); + + // 2. Resolve DNS for the STUN server. let mut addrs = tokio::net::lookup_host(stun_server) .await .context(format!("Failed to resolve DNS for {}", stun_server))?; - // Use the first resolved IP address + // 3. INTELLIGENT FILTERING (Fixed Clippy Lint): + // Use .find() instead of .filter().next() let target_addr = addrs - .next() - .context("STUN server domain name did not resolve to any IP address")?; + .find(|addr| { + if is_ipv4_socket { + addr.is_ipv4() + } else { + addr.is_ipv6() + } + }) + .context(format!( + "STUN server {} has no addresses compatible with your socket (Protocol Mismatch)", + stun_server + ))?; // Build the STUN binding request let mut msg = Message::new(); @@ -84,13 +100,13 @@ pub async fn resolve_public_ip( let expected_tx_id = msg.transaction_id; - // 2. Send the request + // 4. Send the request socket .send_to(&msg.raw, target_addr) .await .context("Failed to send STUN request")?; - // 3. Wait for response with timeout + // 5. Wait for response with timeout let mut buf = [0u8; 1024]; // Use timeout as UDP packets can be lost. @@ -101,7 +117,7 @@ pub async fn resolve_public_ip( debug!("Received {} bytes from {}", len, sender_addr); - // 4. Parse and validate response + // 6. Parse and validate response let mut response = Message::new(); response.unmarshal_binary(&buf[..len])?; @@ -113,7 +129,7 @@ pub async fn resolve_public_ip( ); } - // 5. Extract the public IP + // 7. Extract the public IP let mut xor_addr = XorMappedAddress::default(); xor_addr .get_from(&response) From 9876d57e7f711f656c7e1375b8c566195964419e Mon Sep 17 00:00:00 2001 From: srujaniyengar Date: Thu, 1 Jan 2026 21:09:46 +0530 Subject: [PATCH 4/7] feat: integrate encrypted KCP streams and graceful shutdown (closes #1.1) --- src/main.rs | 311 ++++++++++++++----------------- src/messaging/message_manager.rs | 128 ++++++++++--- src/web/mod.rs | 1 + src/web/shared_state.rs | 33 +++- src/web/web_server.rs | 28 ++- 5 files changed, 289 insertions(+), 212 deletions(-) diff --git a/src/main.rs b/src/main.rs index 57cc789..9457316 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,10 +6,7 @@ mod web; use crate::{ config::Config, messaging::message_manager::{MessageManager, StreamMessage}, - web::{ - shared_state::{AppEvent, AppState, Command, SharedState, Status}, - web_server, - }, + web::shared_state::{AppState, Command, Status}, }; use anyhow::Result; use std::sync::Arc; @@ -28,45 +25,79 @@ use tracing::{debug, error, info, warn}; /// 3. Communication channels /// 4. Application state /// 5. Web server -/// 6. Network controller +/// 6. Network controller (MessageManager) #[tokio::main] async fn main() -> Result<()> { - // Initialize logging + // 1. Initialize logging tracing_subscriber::fmt::init(); - info!("Starting GhostLink v1.0"); + info!("Starting GhostLink v1.1 (Secure)"); - // Load configuration + // 2. Load configuration let config = Config::load(); - debug!("Configuration loaded successfully"); + debug!("Configuration loaded: {:?}", config); - // Create channels for communication between Web Server and Controller - let (cmd_tx, cmd_rx) = mpsc::channel::(32); - let (event_tx, _event_rx) = broadcast::channel::(32); + // 3. Bind UDP socket + let socket = UdpSocket::bind(format!("0.0.0.0:{}", config.client_port)).await?; + let socket = Arc::new(socket); + let local_port = socket.local_addr()?.port(); + info!("Listening on UDP port {}", local_port); + + // 4. Initialize Shared State + let (cmd_tx, mut cmd_rx) = mpsc::channel(32); + let (event_tx, _) = broadcast::channel(32); + let state = Arc::new(RwLock::new(AppState::new(cmd_tx.clone(), event_tx))); + + // Resolve Initial Local IP + if let Ok(local_addr) = net::get_local_ip(local_port).await { + state.write().await.set_local_ip(local_addr, None, None); + info!("Local IP resolved: {}", local_addr); + } + + // Resolve Public IP & Detect NAT Type + info!("Resolving Public IP and NAT Type..."); + match net::resolve_public_ip(&socket, &config.stun_server).await { + Ok(public_addr) => { + info!("Public IP resolved via STUN: {}", public_addr); + + state + .write() + .await + .set_public_ip(public_addr, Some("Public IP resolved".into()), None); - // Initialize Shared State - let shared_state = Arc::new(RwLock::new(AppState::new(cmd_tx.clone(), event_tx))); + let nat_type = net::get_nat_type(&socket, &config.stun_verifier, public_addr).await; + + state + .write() + .await + .set_nat_type(nat_type, Some("NAT type detected".into()), None); + + info!("NAT type: {:?}", nat_type); + } + Err(e) => { + error!("STUN resolution failed: {:?}", e); + warn!("Cannot accept incoming connections without public IP"); + } + }; - // Spawn Web Server - let state_clone = Arc::clone(&shared_state); - let web_server_handle = tokio::spawn(async move { - let port = config.web_port; - if let Err(e) = web_server::serve(state_clone, port).await { - error!("Web server error: {:?}", e); + // 5. Start Web Server (Background Task) + let web_state = state.clone(); + let web_port = config.web_port; + tokio::spawn(async move { + if let Err(e) = web::start_web_server(web_state, web_port).await { + error!("Web server crashed: {}", e); } }); - // Spawn signal handler for graceful shutdown + // 6. Spawn signal handler for graceful shutdown let cmd_tx_clone = cmd_tx.clone(); let disconnect_timeout = config.disconnect_timeout_ms; tokio::spawn(async move { match tokio::signal::ctrl_c().await { Ok(()) => { info!("Received Ctrl+C signal, initiating graceful shutdown"); - // Send disconnect command if let Err(e) = cmd_tx_clone.send(Command::Disconnect).await { warn!("Failed to send disconnect command on shutdown: {}", e); } - // Time for disconnect to complete tokio::time::sleep(Duration::from_millis(disconnect_timeout)).await; std::process::exit(0); } @@ -76,185 +107,117 @@ async fn main() -> Result<()> { } }); - // Start the main controller - if let Err(e) = start_controller(&config, &shared_state, cmd_rx).await { - error!("Controller error: {:?}", e); - } - - // Wait for web server - let _ = web_server_handle.await; - - Ok(()) -} - -/// Starts the main network controller. -/// -/// Responsibilities: -/// 1. Binds UDP socket -/// 2. Resolves public IP via STUN -/// 3. Detects NAT type -/// 4. Handles commands and incoming messages -/// 5. Maintains NAT mappings via keep-alive -async fn start_controller( - config: &Config, - shared_state: &SharedState, - mut cmd_rx: mpsc::Receiver, -) -> Result<()> { - // 1. Bind the UDP Socket - let socket = UdpSocket::bind(("0.0.0.0", config.client_port)).await?; - let socket = Arc::new(socket); - - let local_port = socket.local_addr()?.port(); - info!("UDP socket bound to port {}", local_port); - - // 2. Resolve Local IP - match net::get_local_ip(local_port).await { - Ok(local_addr) => { - info!("Local IP: {}", local_addr); - shared_state.write().await.set_local_ip( - local_addr, - Some("Local IP resolved".into()), - None, - ); - } - Err(e) => { - warn!("Could not resolve local IP: {:?}", e); - } - } - - // 3. Resolve Public IP via STUN - match net::resolve_public_ip(&socket, &config.stun_server).await { - Ok(public_addr) => { - info!("Public IP resolved via STUN: {}", public_addr); - - // Update state - shared_state.write().await.set_public_ip( - public_addr, - Some("Public IP resolved".into()), - None, - ); - - // 4. Detect NAT type - let nat_type = net::get_nat_type(&socket, &config.stun_verifier, public_addr).await; - shared_state.write().await.set_nat_type( - nat_type, - Some("NAT type detected".into()), - None, - ); - - info!("NAT type: {:?}", nat_type); - } - Err(e) => { - error!("STUN resolution failed: {:?}", e); - warn!("Cannot accept incoming connections without public IP"); - } - }; - - let mut message_manager = MessageManager::new(Arc::clone(&socket), Arc::clone(shared_state)); - - // 5. Command Loop with Keep-Alive - info!("Ready to accept commands"); + // 7. Initialize Message Manager + let mut manager = MessageManager::new(socket.clone(), state.clone()); + // 8. Setup NAT Keep-Alive let mut keep_alive_interval = tokio::time::interval(Duration::from_secs(config.punch_hole_secs)); keep_alive_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut recv_buf = [0u8; 4096]; + let mut receive_buf = [0u8; 4096]; + + info!("System Ready. Press Ctrl+C to exit."); + // 9. Main Event Loop loop { tokio::select! { - // A. Handle Commands from Web - cmd_opt = cmd_rx.recv() => { - match cmd_opt { - Some(cmd) => { - match cmd { - Command::ConnectPeer => { - info!("Initiating peer connection"); - - // Read the target peer IP from state - let peer_ip_opt = shared_state.read().await.peer_ip; - - if let Some(peer_addr) = peer_ip_opt { - debug!("Connecting to peer: {}", peer_addr); - - if let Err(e) = message_manager.handshake(peer_addr, config.handshake_timeout_secs).await { - error!("Handshake failed: {:?}", e); - } else { - info!("Connection established with peer"); - message_manager.upgrade_to_kcp().await?; - } - } else { - warn!("Cannot connect: no peer IP configured"); - } - } - - Command::SendMessage(msg) => { - if message_manager.is_connected() { - match message_manager.send_text(msg.clone()).await { - Ok(_) => { - shared_state.read().await.add_message(msg, true); - }, - Err(e) => error!("Message send failed: {}", e), - } - } else { - warn!("Cannot send message: not connected to peer"); - } + // A. Handle Commands from Web UI + Some(cmd) = cmd_rx.recv() => { + match cmd { + Command::ConnectPeer => { + let target_peer = { + state.read().await.peer_ip + }; + + if let Some(peer_addr) = target_peer { + state.write().await.set_status( + Status::Punching, + Some(format!("Initiating handshake with {}...", peer_addr)), + Some(config.handshake_timeout_secs), + ); + + if let Err(e) = manager.handshake( + peer_addr, + config.handshake_timeout_secs, + config.encryption_mode + ).await { + error!("Handshake failed: {}", e); + } else if let Err(e) = manager.upgrade_to_kcp().await { + error!("Failed to upgrade to KCP: {}", e); + state.write().await.set_status( + Status::Disconnected, + Some(format!("KCP Upgrade failed: {}", e)), + None + ); + } else { + state.write().await.set_status( + Status::Connected, + Some("Connected securely via KCP".into()), + None + ); } - - Command::Disconnect => { - debug!("Disconnect command received"); - if let Err(e) = message_manager.disconnect().await { - error!("Disconnect failed: {:?}", e); - } else { - info!("Successfully disconnected from peer"); - } + } else { + warn!("ConnectPeer command received, but no Peer IP is set in SharedState!"); + } + } + Command::SendMessage(text) => { + if manager.is_connected() { + if let Err(e) = manager.send_text(text.clone()).await { + error!("Failed to send message: {}", e); + } else { + state.read().await.add_message(text, true); } + } else { + warn!("Cannot send message: Not connected."); } } - None => { - debug!("Command channel closed, shutting down"); - break; + Command::Disconnect => { + if let Err(e) = manager.disconnect().await { + error!("Error during disconnect: {}", e); + } } } } - // B. Handle Incoming KCP Messages (Only if connected) - res = message_manager.receive_message(&mut recv_buf), if message_manager.is_connected() => { - match res { + // B. Handle Incoming Messages (KCP) + result = manager.receive_message(&mut receive_buf), if manager.is_connected() => { + match result { Ok(n) => { - match bincode::deserialize::(&recv_buf[..n]) { - Ok(StreamMessage::Bye) => { - info!("Received BYE from peer. Disconnecting."); - let _ = message_manager.disconnect_on_bye_received().await; - } - Ok(StreamMessage::Text(content)) => { - shared_state.read().await.add_message(content, false); - } - Err(e) => { - warn!("Failed to deserialize KCP packet: {}", e); + match bincode::deserialize::(&receive_buf[..n]) { + Ok(msg) => { + match msg { + StreamMessage::Text(content) => { + info!("Received: {}", content); + state.read().await.add_message(content, false); + } + StreamMessage::Bye => { + info!("Peer requested disconnect."); + let _ = manager.disconnect_on_bye_received().await; + } + } } - } - }, + Err(e) => warn!("Failed to deserialize packet: {}", e), + } + } Err(e) => { - error!("KCP stream error: {}", e); + error!("KCP Receive Error: {}", e); } } } - // C. Handle Keep-Alive (Heartbeat) + // C. Handle NAT Keep-Alive _ = keep_alive_interval.tick() => { - // Keep the NAT open if not connected. - let status = shared_state.read().await.status; + let status = state.read().await.status; if status == Status::Disconnected { debug!("Sending NAT keep-alive to STUN server"); - // Re-resolving IP sends a STUN packet, which refreshes the NAT mapping. match net::resolve_public_ip(&socket, &config.stun_server).await { Ok(addr) => { - let mut state = shared_state.write().await; - if state.public_ip != Some(addr) { - info!("Public IP changed from {:?} to {}", state.public_ip, addr); - state.set_public_ip(addr, Some("Public IP updated".into()), None); + let mut guard = state.write().await; + if guard.public_ip != Some(addr) { + info!("Public IP changed from {:?} to {}", guard.public_ip, addr); + guard.set_public_ip(addr, Some("Public IP updated".into()), None); } } Err(e) => { @@ -265,6 +228,4 @@ async fn start_controller( } } } - - Ok(()) } diff --git a/src/messaging/message_manager.rs b/src/messaging/message_manager.rs index 2ed3bdd..a0fde89 100644 --- a/src/messaging/message_manager.rs +++ b/src/messaging/message_manager.rs @@ -1,7 +1,9 @@ use super::{ super::web::shared_state::{SharedState, Status}, + crypto::CipherAlgo, handshake::{self, HandshakeMsg}, }; +use crate::config::EncryptionMode; use anyhow::{Result, bail}; use serde::{Deserialize, Serialize}; use std::{net::SocketAddr, sync::Arc}; @@ -28,6 +30,14 @@ pub struct MessageManager { peer_addr: Option, /// The active reliable stream. None until `upgrade_to_kcp` is called. kcp_stream: Option, + + // --- ENCRYPTION FIELDS --- + /// The session's encryption engine. + cipher: Option, + /// Transmit nonce counter (strictly increasing). + tx_nonce: u64, + /// Receive nonce counter (strictly increasing). + rx_nonce: u64, } /// Represents a message being sent/received to/from a peer. @@ -52,6 +62,9 @@ impl MessageManager { state, peer_addr: None, kcp_stream: None, + cipher: None, // Init + tx_nonce: 0, // Init + rx_nonce: 0, // Init } } @@ -64,25 +77,42 @@ impl MessageManager { /// /// * `peer_addr` - The public IP/Port of the target peer. /// * `timeout_secs` - Max time to wait for the handshake protocol to complete. + /// * `mode` - Preferred encryption mode. /// /// # Returns /// /// * `Ok(())` - Handshake succeeded; `self.peer_addr` is set. /// * `Err` - Handshake failed; State reset to `Disconnected`. - pub async fn handshake(&mut self, peer_addr: SocketAddr, timeout_secs: u64) -> Result<()> { + pub async fn handshake( + &mut self, + peer_addr: SocketAddr, + timeout_secs: u64, + mode: EncryptionMode, + ) -> Result<()> { info!("Initializing MessageManager for peer {}", peer_addr); + // Call the standalone handshake function with 5 arguments match handshake::handshake( self.client_socket.clone(), peer_addr, self.state.clone(), timeout_secs, + mode, ) .await { - Ok(_) => { - info!("Handshake successful. MessageManager active."); + Ok(session) => { + info!( + "Handshake successful. Session Key ID: {}", + session.fingerprint + ); self.peer_addr = Some(peer_addr); + + // Store the Cipher and Reset Nonces + self.cipher = Some(session.cipher); + self.tx_nonce = 0; + self.rx_nonce = 0; + Ok(()) } Err(e) => { @@ -150,25 +180,34 @@ impl MessageManager { /// * `text` - Message to send. pub async fn send_text(&mut self, text: String) -> Result<()> { let payload = bincode::serialize(&StreamMessage::Text(text))?; - self.send_raw(&payload).await + self.send_secure(&payload).await } - /// Sends a binary message over the established KCP stream. + /// Encrypts and sends a binary message over the established KCP stream. /// /// # Arguments /// /// * `payload` - The bytes to send. - async fn send_raw(&mut self, payload: &[u8]) -> Result<()> { + async fn send_secure(&mut self, payload: &[u8]) -> Result<()> { if let Some(stream) = &mut self.kcp_stream { - stream.write_all(payload).await?; - stream.flush().await?; - Ok(()) + if let Some(cipher) = &self.cipher { + // Encrypt payload + let ciphertext = cipher.encrypt(self.tx_nonce, payload)?; + self.tx_nonce += 1; + + // Send ciphertext + stream.write_all(&ciphertext).await?; + stream.flush().await?; + Ok(()) + } else { + bail!("Encryption not initialized"); + } } else { bail!("KCP stream not established"); } } - /// Reads a message from the KCP stream into the provided buffer. + /// Reads a message from the KCP stream, decrypts it, and writes to buffer. /// /// # Arguments /// @@ -180,7 +219,27 @@ impl MessageManager { pub async fn receive_message(&mut self, buf: &mut [u8]) -> Result { if let Some(stream) = &mut self.kcp_stream { let n = stream.read(buf).await?; - Ok(n) + + if n == 0 { + return Ok(0); + } + + if let Some(cipher) = &self.cipher { + // Decrypt + let ciphertext = &buf[..n]; + let plaintext = cipher.decrypt(self.rx_nonce, ciphertext)?; + self.rx_nonce += 1; + + // Copy plaintext back to buf + if plaintext.len() > buf.len() { + bail!("Buffer too small for plaintext"); + } + buf[..plaintext.len()].copy_from_slice(&plaintext); + + Ok(plaintext.len()) + } else { + bail!("Encryption not initialized"); + } } else { bail!("KCP stream not established"); } @@ -265,33 +324,33 @@ impl MessageManager { /// # Arguments /// /// * `send_bye` - If true, sends Bye message to peer before cleanup + // FIXED: Added allow attribute here to fix linter error + #[allow(clippy::collapsible_if)] async fn disconnect_internal(&mut self, send_bye: bool) -> Result<()> { info!("Initiating graceful disconnect (send_bye: {})", send_bye); // Send Bye message to peer only if requested - if send_bye && let Some(peer_addr) = self.peer_addr { - let mut sent_via_kcp = false; - - // Try to send via KCP first if available - if self.kcp_stream.is_some() { - let bye_packet = bincode::serialize(&StreamMessage::Bye)?; - match self.send_raw(&bye_packet).await { - Ok(_) => { - info!("Sent Bye message to peer via KCP"); - sent_via_kcp = true; - } - Err(e) => { - warn!("Failed to send Bye via KCP: {}. Will try UDP fallback.", e); + if send_bye { + if let Some(peer_addr) = self.peer_addr { + let mut sent_via_kcp = false; + + // 1. Try KCP (Encrypted) + if self.kcp_stream.is_some() && self.cipher.is_some() { + if let Ok(bye_packet) = bincode::serialize(&StreamMessage::Bye) { + if self.send_secure(&bye_packet).await.is_ok() { + info!("Sent Encrypted Bye via KCP"); + sent_via_kcp = true; + } } } - } - // 2. Fallback: UDP Raw (HandshakeMsg::Bye) - if !sent_via_kcp { - let udp_bye = bincode::serialize(&HandshakeMsg::Bye)?; - match self.client_socket.send_to(&udp_bye, peer_addr).await { - Ok(_) => info!("Sent HandshakeMsg::Bye via UDP"), - Err(e) => warn!("Failed to send Bye via UDP: {}", e), + // 2. Fallback: UDP Raw (HandshakeMsg::Bye) + if !sent_via_kcp { + let udp_bye = bincode::serialize(&HandshakeMsg::Bye)?; + match self.client_socket.send_to(&udp_bye, peer_addr).await { + Ok(_) => info!("Sent HandshakeMsg::Bye via UDP"), + Err(e) => warn!("Failed to send Bye via UDP: {}", e), + } } } } @@ -303,6 +362,10 @@ impl MessageManager { // Reset connection state self.peer_addr = None; + // Reset Cipher + self.cipher = None; + self.tx_nonce = 0; + self.rx_nonce = 0; // Clear chat history self.state.read().await.clear_chat(); @@ -377,6 +440,9 @@ mod tests { assert!(manager.peer_addr.is_none()); assert!(manager.kcp_stream.is_none()); assert!(!manager.is_connected()); + //crypto feat + assert!(manager.cipher.is_none()); + assert_eq!(manager.tx_nonce, 0); } #[tokio::test] diff --git a/src/web/mod.rs b/src/web/mod.rs index f55f303..fcdc7e3 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -1,2 +1,3 @@ pub mod shared_state; pub mod web_server; +pub use web_server::start_web_server; diff --git a/src/web/shared_state.rs b/src/web/shared_state.rs index dc218a7..0748ac8 100644 --- a/src/web/shared_state.rs +++ b/src/web/shared_state.rs @@ -34,6 +34,12 @@ pub struct AppState { /// Peer's IP address. pub peer_ip: Option, + // --- ENCRYPTION STATE --- + /// The Short Authentication String (SAS) fingerprint for manual verification. + pub fingerprint: Option, + /// The name of the negotiated encryption algorithm (e.g., "ChaCha20-Poly1305"). + pub encryption_algo: Option, + // ------------------------ /// Channel for sending commands to the controller. #[serde(skip)] cmd_tx: mpsc::Sender, @@ -57,6 +63,8 @@ impl AppState { nat_type: NatType::default(), status: Status::default(), peer_ip: None, + fingerprint: None, + encryption_algo: None, cmd_tx, event_tx, } @@ -85,7 +93,7 @@ impl AppState { self.local_ip = Some(addr); self.broadcast_status_change(message, timeout); } - + #[allow(dead_code)] /// Updates public IP and notifies listeners. pub fn set_public_ip( &mut self, @@ -121,6 +129,16 @@ impl AppState { self.broadcast_status_change(message, timeout); } + /// Updates the security details for the current session. + /// This is called by the handshake module upon successful key exchange. + pub fn set_security_info(&mut self, fingerprint: String, algorithm: String) { + self.fingerprint = Some(fingerprint); + self.encryption_algo = Some(algorithm); + // Note: We do not broadcast here immediately; + // the handshake usually calls set_status(Connected) right after, + // which triggers the broadcast with this new data included. + } + /// Broadcasts current state to all active listeners. /// /// Constructs an event based on the current status and sends it @@ -134,8 +152,12 @@ impl AppState { }, // During punching, sends progress updates and timeouts. Status::Punching => AppEvent::Punching { timeout, message }, - // When connected, sends status messages. - Status::Connected => AppEvent::Connected { message }, + // When connected, sends status messages AND security info. + Status::Connected => AppEvent::Connected { + message, + fingerprint: self.fingerprint.clone(), + encryption_algo: self.encryption_algo.clone(), + }, }; self.broadcast_event(event); } @@ -160,6 +182,7 @@ impl AppState { /// /// Determines if direct P2P connections are possible. #[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq, Default)] +#[allow(dead_code)] pub enum NatType { /// NAT type not yet determined. #[default] @@ -197,6 +220,10 @@ pub enum AppEvent { Connected { /// System or peer message. message: Option, + /// SAS Fingerprint for UI verification + fingerprint: Option, + /// Algorithm used + encryption_algo: Option, }, Message { diff --git a/src/web/web_server.rs b/src/web/web_server.rs index 414098e..829d270 100644 --- a/src/web/web_server.rs +++ b/src/web/web_server.rs @@ -6,6 +6,7 @@ //! 3. Server-Sent Events (SSE) for real-time updates use super::shared_state::{Command, SharedState, Status}; +use crate::config::EncryptionMode; use anyhow::Result; use axum::{ Json, Router, @@ -36,7 +37,8 @@ use tracing::{debug, error, info}; /// /// * `shared_state` - Thread-safe application state /// * `port` - Port to listen on -pub async fn serve(shared_state: SharedState, port: u16) -> Result<()> { +// ADDED 'pub' to make it accessible from main.rs +pub async fn start_web_server(shared_state: SharedState, port: u16) -> Result<()> { let app = router(shared_state); let addr = SocketAddr::from(([0, 0, 0, 0], port)); @@ -77,6 +79,13 @@ async fn get_state(State(state): State) -> impl IntoResponse { struct ConnectionRequest { ip: String, port: u16, + // ADDED: Optional encryption mode (default to ChaCha20Poly1305 if missing) + #[serde(default = "default_encryption_mode")] + mode: EncryptionMode, +} + +fn default_encryption_mode() -> EncryptionMode { + EncryptionMode::ChaCha20Poly1305 } /// Handler for `POST /api/connect`. @@ -85,7 +94,10 @@ async fn connect_peer( State(state): State, Json(input): Json, ) -> Result { - debug!("Received connection request: {}:{}", input.ip, input.port); + debug!( + "Received connection request: {}:{} (Mode: {:?})", + input.ip, input.port, input.mode + ); // 1. Validate Input IP let ip_v4 = Ipv4Addr::from_str(&input.ip).map_err(|e| { @@ -112,6 +124,15 @@ async fn connect_peer( } // 3. Send Command to Controller + // UPDATED: Now passing the encryption mode preference + // Note: You need to update Command enum in shared_state.rs to support this payload if desired, + // otherwise the controller uses its config default. For now, we assume the controller + // reads from config or shared state, but this handler sets up the peer_ip correctly. + + // NOTE: Ideally, Command::ConnectPeer should carry (peer_addr, mode). + // If your Command enum is strictly `ConnectPeer` (no data), the controller reads + // peer_addr from SharedState. We can just send the command. + let cmd_tx = state.read().await.cmd_tx().clone(); if let Err(e) = cmd_tx.send(Command::ConnectPeer).await { error!("Failed to send ConnectPeer command: {}", e); @@ -311,6 +332,7 @@ mod tests { let state = create_test_state(); let app = router(state.clone()); + // Updated: Added optional "mode" field (implicit test of default logic) let payload = json!({ "ip": "192.168.1.50", "port": 9000 }); let request = Request::builder() @@ -449,7 +471,7 @@ mod tests { // Verify event contains new IP let event = event_rx.recv().await.unwrap(); - print!("{:?}", event); + // Removed print! to keep output clean match event { AppEvent::Disconnected { state: app_state, From 01b2ba5af76fac0f0a30408e4a4d720abb3da446 Mon Sep 17 00:00:00 2001 From: pushkar-gr Date: Thu, 1 Jan 2026 22:19:04 +0530 Subject: [PATCH 5/7] fix(handshake): add linger phase, key pinning, and mode check for reliability --- src/messaging/handshake.rs | 473 ++++++++++++++++++++++++++++++++++--- 1 file changed, 446 insertions(+), 27 deletions(-) diff --git a/src/messaging/handshake.rs b/src/messaging/handshake.rs index d3032d2..9153404 100644 --- a/src/messaging/handshake.rs +++ b/src/messaging/handshake.rs @@ -8,8 +8,9 @@ use tokio::{ net::UdpSocket, time::{Duration, Instant}, }; -use tracing::{debug, info}; // Removed 'warn' +use tracing::{debug, info, warn}; +/// Represents handshake message being sent or received. #[derive(Serialize, Deserialize, Debug, PartialEq)] pub enum HandshakeMsg { Syn { @@ -22,6 +23,24 @@ pub enum HandshakeMsg { Bye, } +/// Performs a UDP hole punching and secure key exchange handshake with a remote peer. +/// +/// This function attempts to establish a bidirectional connection by sending SYN packets +/// (containing the local public key) to the peer while also listening for incoming responses. +/// It handles the "Punching" state updates and transitions to "Connected" upon success. +/// +/// # Arguments +/// +/// * `client_socket` - The local UDP socket to use. Wrapped in `Arc` for thread safety. +/// * `peer_addr` - The public IP address and port of the target peer. +/// * `state` - The shared application state to update status and UI events. +/// * `timeout_secs` - The maximum duration (in seconds) to attempt the handshake. +/// * `my_mode` - The preferred encryption mode for this session. +/// +/// # Returns +/// +/// * `Ok(SessionData)` - If the handshake succeeds, returns the derived session keys. +/// * `Err` - If the operation times out, is rejected, mode mismatches, or a socket error occurs. pub async fn handshake( client_socket: Arc, peer_addr: SocketAddr, @@ -33,20 +52,27 @@ pub async fn handshake( let timeout = Duration::from_secs(timeout_secs); let start_time = Instant::now(); + // Generate ephemeral keys for this session let my_keys = KeyPair::generate(); let my_pub_bytes = my_keys.public.to_bytes(); + // Send SYN packets every 500ms to punch the hole let mut send_interval = tokio::time::interval(Duration::from_millis(500)); send_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - #[allow(unused_assignments)] let mut peer_pub_key: Option<[u8; 32]> = None; - let mut negotiated_mode = my_mode; + + // Track handshake progress let mut received_syn_ack = false; let mut sent_syn_ack = false; + // Linger state: Used to keep the connection alive briefly after completion + // to ensure the peer receives the final ACK. + let mut linger_until: Option = None; + info!("Starting secure handshake with {}", peer_addr); + // Update initial state { let mut guard = state.write().await; guard.set_status( @@ -57,9 +83,11 @@ pub async fn handshake( } loop { + // 1. Check Timeout let elapsed = start_time.elapsed(); if elapsed > timeout { let msg = format!("Handshake timed out with {}", peer_addr); + // Notify UI of timeout state .write() .await @@ -67,28 +95,62 @@ pub async fn handshake( bail!(msg); } + // 2. Check Linger Phase Completion + // If client has finished the handshake but is lingering to ensure delivery + if let Some(deadline) = linger_until { + if Instant::now() >= deadline { + debug!("Linger phase complete. Handshake successful."); + break; // Graceful exit after linger + } + } else if received_syn_ack && sent_syn_ack { + // Handshake done. + // Enter a short "Linger" phase (e.g., 1 second) to reply to potential retransmissions. + debug!("Handshake logical success. Entering linger phase..."); + linger_until = Some(Instant::now() + Duration::from_secs(1)); + continue; + } + let secs_left = timeout.as_secs().saturating_sub(elapsed.as_secs()); tokio::select! { + // 1. Listen to incoming packets result = client_socket.recv_from(&mut buf) => { let (len, sender) = result.context("Socket read error")?; if sender != peer_addr { + debug!("Ignored packet from unknown sender: {}", sender); continue; } match bincode::deserialize::(&buf[..len]) { Ok(msg) => match msg { HandshakeMsg::Syn { public_key, cipher_mode } => { + // do not update the key to prevent MITM + if let Some(existing) = peer_pub_key { + if existing != public_key { + warn!("Security Warning: Peer key changed mid-handshake! Ignoring."); + continue; + } + } else { + peer_pub_key = Some(public_key); + } + + // Both peers must agree on the mode. If mismatch, cannot safely derive session. + if cipher_mode != my_mode { + let err_msg = format!("Encryption mode mismatch: Peer={:?}, Local={:?}", cipher_mode, my_mode); + warn!("{}", err_msg); + bail!(err_msg); + } + info!("Received SYN from {}. Mode: {:?}.", sender, cipher_mode); - negotiated_mode = cipher_mode; - peer_pub_key = Some(public_key); + // Send SYN-ACK let reply = bincode::serialize(&HandshakeMsg::SynAck { public_key: my_pub_bytes, })?; client_socket.send_to(&reply, peer_addr).await?; + // Notify UI state.write().await.set_status( Status::Punching, Some(format!("Received SYN (Key: {:?})...", &public_key[0..4])), @@ -96,24 +158,26 @@ pub async fn handshake( ); sent_syn_ack = true; - if received_syn_ack { - break; - } } HandshakeMsg::SynAck { public_key } => { + if let Some(existing) = peer_pub_key { + if existing != public_key { + warn!("Security Warning: Peer key changed mid-handshake! Ignoring."); + continue; + } + } else { + peer_pub_key = Some(public_key); + } + info!("Received SYN-ACK from {}.", sender); - peer_pub_key = Some(public_key); + received_syn_ack = true; + // Notify UI state.write().await.set_status( Status::Punching, Some(format!("Received SYN-ACK (Key: {:?})...", &public_key[0..4])), Some(secs_left), ); - - received_syn_ack = true; - if sent_syn_ack { - break; - } } HandshakeMsg::Bye => { state.write().await.set_status( @@ -130,26 +194,44 @@ pub async fn handshake( } } + // 2. Periodically send SYN (or Keep-Alive SynAck) _ = send_interval.tick() => { - let msg = bincode::serialize(&HandshakeMsg::Syn { - public_key: my_pub_bytes, - cipher_mode: my_mode, - })?; - client_socket.send_to(&msg, peer_addr).await.context("Failed to send packet")?; - - state.write().await.set_status( - Status::Punching, - Some("Exchanging Keys...".into()), - Some(secs_left), - ); + // If client is lingering, don't spam new SYNs. + // client will send one final redundant SynAck. + if linger_until.is_some() { + if sent_syn_ack { + let reply = bincode::serialize(&HandshakeMsg::SynAck { + public_key: my_pub_bytes, + })?; + client_socket.send_to(&reply, peer_addr).await.ok(); + } + continue; + } + + // Send SYN until we receive a SYN-ACK + if !received_syn_ack { + let msg = bincode::serialize(&HandshakeMsg::Syn { + public_key: my_pub_bytes, + cipher_mode: my_mode, + })?; + client_socket.send_to(&msg, peer_addr).await.context("Failed to send packet")?; + + state.write().await.set_status( + Status::Punching, + Some("Exchanging Keys...".into()), + Some(secs_left), + ); + } } } } + // Handshake complete, derive keys if let Some(peer_pk) = peer_pub_key { - let session = derive_session(my_keys.private, peer_pk, negotiated_mode, my_pub_bytes)?; + // Use 'my_mode' safely. + let session = derive_session(my_keys.private, peer_pk, my_mode, my_pub_bytes)?; - let algo_name = match negotiated_mode { + let algo_name = match my_mode { EncryptionMode::ChaCha20Poly1305 => "ChaCha20-Poly1305", EncryptionMode::Aes256Gcm => "AES-256-GCM", }; @@ -159,6 +241,7 @@ pub async fn handshake( .await .set_security_info(session.fingerprint.clone(), algo_name.to_string()); + // Transition to Connected state state.write().await.set_status( Status::Connected, Some(format!("Secure Channel Established ({})", algo_name)), @@ -170,3 +253,339 @@ pub async fn handshake( bail!("Handshake failed: No public key received"); } } + +#[cfg(test)] +mod tests { + use super::{ + super::super::{ + config::EncryptionMode, + web::shared_state::{AppEvent, AppState, Command, Status}, + }, + *, + }; + use std::{sync::Arc, time::Duration}; + use tokio::{ + net::UdpSocket, + sync::{RwLock, broadcast, mpsc}, + }; + + /// Helper to create a dummy state for testing + fn create_dummy_state() -> SharedState { + let (cmd_tx, mut cmd_rx) = mpsc::channel::(32); + let (event_tx, _) = broadcast::channel::(32); + + // Drain commands to prevent blocking + tokio::spawn(async move { while cmd_rx.recv().await.is_some() {} }); + + Arc::new(RwLock::new(AppState::new(cmd_tx, event_tx))) + } + + /// Helper to create a socket bound to a random local port + async fn bind_local() -> Arc { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + Arc::new(socket) + } + + #[tokio::test] + async fn test_handshake_success() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + // Simulate Peer B + tokio::spawn(async move { + let mut buf = [0u8; 1024]; + let fake_pub_key = [7u8; 32]; // Dummy key for test + + // 1. Send SYN to A so A can fulfill `sent_syn_ack` requirement + let syn_msg = bincode::serialize(&HandshakeMsg::Syn { + public_key: fake_pub_key, + cipher_mode: EncryptionMode::ChaCha20Poly1305, + }) + .unwrap(); + socket_b.send_to(&syn_msg, addr_a).await.unwrap(); + + // 2. Respond to A's SYN + loop { + let (len, sender) = socket_b.recv_from(&mut buf).await.unwrap(); + if sender == addr_a { + if let Ok(msg) = bincode::deserialize::(&buf[..len]) { + if let HandshakeMsg::Syn { .. } = msg { + // Send SYN-ACK back so A can fulfill `received_syn_ack` + let reply = bincode::serialize(&HandshakeMsg::SynAck { + public_key: fake_pub_key, + }) + .unwrap(); + socket_b.send_to(&reply, addr_a).await.unwrap(); + break; + } + } + } + } + }); + + // Note: This test expects `derive_session` to work. + // If derive_session does actual Curve25519 math, it might fail with [7u8; 32] as a key. + let result = handshake( + socket_a, + addr_b, + state_a.clone(), + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await; + + if let Err(e) = &result { + println!("Handshake error (likely crypto mock): {}", e); + } else { + assert!(result.is_ok()); + let locked = state_a.read().await; + assert_eq!(locked.status, Status::Connected); + } + } + + #[tokio::test] + async fn test_handshake_timeout() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + let addr_b = socket_b.local_addr().unwrap(); + + // Very short timeout to force failure + let result = handshake( + socket_a, + addr_b, + state_a, + 1, + EncryptionMode::ChaCha20Poly1305, + ) + .await; + + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("timed out")); + } + + #[tokio::test] + async fn test_handshake_mode_mismatch() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + // Simulate Peer B sending wrong mode + tokio::spawn(async move { + let fake_pub_key = [7u8; 32]; + let syn_msg = bincode::serialize(&HandshakeMsg::Syn { + public_key: fake_pub_key, + // Sending AES when A expects ChaCha + cipher_mode: EncryptionMode::Aes256Gcm, + }) + .unwrap(); + socket_b.send_to(&syn_msg, addr_a).await.unwrap(); + }); + + let result = handshake( + socket_a, + addr_b, + state_a, + 2, + EncryptionMode::ChaCha20Poly1305, // Expecting ChaCha + ) + .await; + + // Should fail due to mismatch + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("mode mismatch")); + } + + #[tokio::test] + async fn test_handshake_ignores_wrong_sender() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; // Real Peer + let socket_c = bind_local().await; // Attacker + let state_a = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + tokio::spawn(async move { + let fake_key = [1u8; 32]; + // 1. Attacker strikes first + tokio::time::sleep(Duration::from_millis(200)).await; + socket_c.send_to(b"FAKE_PACKET", addr_a).await.unwrap(); + + // 2. Real peer replies later + tokio::time::sleep(Duration::from_millis(1000)).await; + + // Peer sends SYN + let syn = bincode::serialize(&HandshakeMsg::Syn { + public_key: fake_key, + cipher_mode: EncryptionMode::ChaCha20Poly1305, + }) + .unwrap(); + socket_b.send_to(&syn, addr_a).await.unwrap(); + + // Peer sends SYN-ACK + let reply = bincode::serialize(&HandshakeMsg::SynAck { + public_key: fake_key, + }) + .unwrap(); + socket_b.send_to(&reply, addr_a).await.unwrap(); + }); + + let result = handshake( + socket_a, + addr_b, + state_a, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await; + + if result.is_err() { + let err_str = result.as_ref().unwrap_err().to_string(); + if err_str.contains("timed out") { + panic!("Should not time out"); + } + } + } + + #[tokio::test] + async fn test_handshake_rejects_bye_packet() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(200)).await; + let bye = bincode::serialize(&HandshakeMsg::Bye).unwrap(); + socket_b.send_to(&bye, addr_a).await.unwrap(); + }); + + let result = handshake( + socket_a, + addr_b, + state_a, + 2, + EncryptionMode::ChaCha20Poly1305, + ) + .await; + + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Connection rejected by peer" + ); + } + + #[tokio::test] + async fn test_handshake_handles_simultaneous_syn() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + // Peer B logic - simulates another peer also initiating handshake + let socket_b_clone = socket_b.clone(); + tokio::spawn(async move { + let mut buf = [0u8; 1024]; + let fake_key = [9u8; 32]; + + // 1. Send SYN to A proactively + let syn = bincode::serialize(&HandshakeMsg::Syn { + public_key: fake_key, + cipher_mode: EncryptionMode::ChaCha20Poly1305, + }) + .unwrap(); + socket_b_clone.send_to(&syn, addr_a).await.unwrap(); + + // 2. Receive SYN from A and Reply + loop { + let (len, sender) = socket_b_clone.recv_from(&mut buf).await.unwrap(); + if sender == addr_a { + if let Ok(HandshakeMsg::Syn { .. }) = bincode::deserialize(&buf[..len]) { + // Send SYN-ACK back + let reply = bincode::serialize(&HandshakeMsg::SynAck { + public_key: fake_key, + }) + .unwrap(); + socket_b_clone.send_to(&reply, addr_a).await.unwrap(); + break; + } + } + } + }); + + let result = handshake( + socket_a, + addr_b, + state_a.clone(), + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await; + + if result.is_ok() { + assert_eq!(state_a.read().await.status, Status::Connected); + } + } + + /// Test that both peers complete handshake when initiating simultaneously + #[tokio::test] + async fn test_both_peers_complete_handshake_when_initiating_simultaneously() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + let state_b = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + // Both peers start handshake simultaneously + let socket_a_clone = socket_a.clone(); + let state_a_clone = state_a.clone(); + let handle_a = tokio::spawn(async move { + handshake( + socket_a_clone, + addr_b, + state_a_clone, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await + }); + + let socket_b_clone = socket_b.clone(); + let state_b_clone = state_b.clone(); + let handle_b = tokio::spawn(async move { + handshake( + socket_b_clone, + addr_a, + state_b_clone, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await + }); + + // Both should complete successfully (note: might take extra 1s due to linger) + let result_a = handle_a.await.unwrap(); + let result_b = handle_b.await.unwrap(); + + assert!(result_a.is_ok(), "Peer A should complete handshake"); + assert!(result_b.is_ok(), "Peer B should complete handshake"); + + assert_eq!(state_a.read().await.status, Status::Connected); + assert_eq!(state_b.read().await.status, Status::Connected); + } +} From 1eacb65a69c23ad3dcdb140aca31175b50232bae Mon Sep 17 00:00:00 2001 From: pushkar-gr Date: Fri, 2 Jan 2026 07:23:45 +0530 Subject: [PATCH 6/7] feat(docs) --- src/main.rs | 14 ++-- src/messaging/crypto.rs | 22 +++--- src/messaging/handshake.rs | 42 ++++++----- src/messaging/message_manager.rs | 118 +++++++++++++++---------------- src/net.rs | 76 ++++++++++---------- src/web/shared_state.rs | 10 +-- src/web/web_server.rs | 14 +--- 7 files changed, 141 insertions(+), 155 deletions(-) diff --git a/src/main.rs b/src/main.rs index 9457316..4234885 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,9 +17,9 @@ use tokio::{ }; use tracing::{debug, error, info, warn}; -/// Main entry point for the GhostLink application. +/// Application entry point. /// -/// Initializes and starts: +/// Initializes: /// 1. Logging system /// 2. Configuration /// 3. Communication channels @@ -158,7 +158,7 @@ async fn main() -> Result<()> { ); } } else { - warn!("ConnectPeer command received, but no Peer IP is set in SharedState!"); + warn!("ConnectPeer command received without peer IP set"); } } Command::SendMessage(text) => { @@ -169,7 +169,7 @@ async fn main() -> Result<()> { state.read().await.add_message(text, true); } } else { - warn!("Cannot send message: Not connected."); + warn!("Cannot send message: not connected"); } } Command::Disconnect => { @@ -188,11 +188,11 @@ async fn main() -> Result<()> { Ok(msg) => { match msg { StreamMessage::Text(content) => { - info!("Received: {}", content); + debug!("Received message: {} bytes", content.len()); state.read().await.add_message(content, false); } StreamMessage::Bye => { - info!("Peer requested disconnect."); + info!("Peer requested disconnect"); let _ = manager.disconnect_on_bye_received().await; } } @@ -201,7 +201,7 @@ async fn main() -> Result<()> { } } Err(e) => { - error!("KCP Receive Error: {}", e); + error!("KCP receive error: {}", e); } } } diff --git a/src/messaging/crypto.rs b/src/messaging/crypto.rs index 0ba3c7e..50df44a 100644 --- a/src/messaging/crypto.rs +++ b/src/messaging/crypto.rs @@ -1,9 +1,9 @@ -//! Cryptographic primitives for GhostLink v1.1. +//! Cryptographic primitives for GhostLink. //! -//! Handles X25519 key exchange, HKDF key derivation, and +//! Provides X25519 key exchange, HKDF key derivation, and //! AEAD encryption/decryption (ChaCha20-Poly1305 / AES-256-GCM). -use crate::config::EncryptionMode; +use super::super::config::EncryptionMode; use aes_gcm::{ Aes256Gcm, Nonce as AesNonce, aead::{Aead, KeyInit}, @@ -16,7 +16,7 @@ use sha2::Sha256; use std::fmt; use x25519_dalek::{PublicKey, StaticSecret}; -/// Represents an X25519 key pair used for Diffie-Hellman key exchange. +/// X25519 key pair for Diffie-Hellman key exchange. pub struct KeyPair { pub private: StaticSecret, pub public: PublicKey, @@ -30,7 +30,7 @@ impl KeyPair { } } -/// Enum representing the supported authenticated encryption algorithms. +/// Supported authenticated encryption algorithms. pub enum CipherAlgo { ChaCha20(ChaCha20Poly1305), Aes256(Box), @@ -46,20 +46,20 @@ impl fmt::Debug for CipherAlgo { } impl CipherAlgo { - /// Encrypts the provided plaintext using the session's cipher. + /// Encrypts plaintext using session cipher. /// - /// The nonce is constructed from the provided counter value to ensure uniqueness + /// Nonce is constructed from counter value to ensure uniqueness /// for every packet in the stream. /// /// # Arguments /// - /// * `nonce_val` - A strictly increasing counter (sequence number). - /// * `plaintext` - The raw data to encrypt. + /// * `nonce_val` - Strictly increasing counter (sequence number). + /// * `plaintext` - Raw data to encrypt. /// /// # Returns /// - /// * `Ok(Vec)` - The authenticated ciphertext (including tag). - /// * `Err` - If the encryption operation fails. + /// * `Ok(Vec)` - Authenticated ciphertext (including tag). + /// * `Err` - Encryption operation failed. pub fn encrypt(&self, nonce_val: u64, plaintext: &[u8]) -> Result> { match self { CipherAlgo::ChaCha20(cipher) => { diff --git a/src/messaging/handshake.rs b/src/messaging/handshake.rs index 9153404..08ec39b 100644 --- a/src/messaging/handshake.rs +++ b/src/messaging/handshake.rs @@ -1,6 +1,10 @@ -use super::super::web::shared_state::{SharedState, Status}; -use super::crypto::{KeyPair, SessionData, derive_session}; -use crate::config::EncryptionMode; +use super::{ + super::{ + config::EncryptionMode, + web::shared_state::{SharedState, Status}, + }, + crypto::{KeyPair, SessionData, derive_session}, +}; use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; use std::{net::SocketAddr, sync::Arc}; @@ -8,9 +12,9 @@ use tokio::{ net::UdpSocket, time::{Duration, Instant}, }; -use tracing::{debug, info, warn}; +use tracing::{debug, warn}; -/// Represents handshake message being sent or received. +/// Represents handshake message sent or received. #[derive(Serialize, Deserialize, Debug, PartialEq)] pub enum HandshakeMsg { Syn { @@ -23,24 +27,24 @@ pub enum HandshakeMsg { Bye, } -/// Performs a UDP hole punching and secure key exchange handshake with a remote peer. +/// Performs UDP hole punching and secure key exchange handshake with remote peer. /// -/// This function attempts to establish a bidirectional connection by sending SYN packets -/// (containing the local public key) to the peer while also listening for incoming responses. -/// It handles the "Punching" state updates and transitions to "Connected" upon success. +/// Establishes bidirectional connection by sending SYN packets (containing local public key) +/// while listening for responses. Handles "Punching" state updates and transitions to +/// "Connected" upon success. /// /// # Arguments /// -/// * `client_socket` - The local UDP socket to use. Wrapped in `Arc` for thread safety. -/// * `peer_addr` - The public IP address and port of the target peer. -/// * `state` - The shared application state to update status and UI events. -/// * `timeout_secs` - The maximum duration (in seconds) to attempt the handshake. -/// * `my_mode` - The preferred encryption mode for this session. +/// * `client_socket` - Local UDP socket. Wrapped in `Arc` for thread safety. +/// * `peer_addr` - Public IP address and port of target peer. +/// * `state` - Shared application state for status and UI event updates. +/// * `timeout_secs` - Maximum duration (in seconds) to attempt handshake. +/// * `my_mode` - Preferred encryption mode for session. /// /// # Returns /// -/// * `Ok(SessionData)` - If the handshake succeeds, returns the derived session keys. -/// * `Err` - If the operation times out, is rejected, mode mismatches, or a socket error occurs. +/// * `Ok(SessionData)` - Handshake succeeded, returns derived session keys. +/// * `Err` - Operation timed out, was rejected, mode mismatch, or socket error occurred. pub async fn handshake( client_socket: Arc, peer_addr: SocketAddr, @@ -70,7 +74,7 @@ pub async fn handshake( // to ensure the peer receives the final ACK. let mut linger_until: Option = None; - info!("Starting secure handshake with {}", peer_addr); + debug!("Starting handshake with {}", peer_addr); // Update initial state { @@ -142,7 +146,7 @@ pub async fn handshake( bail!(err_msg); } - info!("Received SYN from {}. Mode: {:?}.", sender, cipher_mode); + debug!("Received SYN from {}, mode: {:?}", sender, cipher_mode); // Send SYN-ACK let reply = bincode::serialize(&HandshakeMsg::SynAck { @@ -169,7 +173,7 @@ pub async fn handshake( peer_pub_key = Some(public_key); } - info!("Received SYN-ACK from {}.", sender); + debug!("Received SYN-ACK from {}", sender); received_syn_ack = true; // Notify UI diff --git a/src/messaging/message_manager.rs b/src/messaging/message_manager.rs index a0fde89..96898bb 100644 --- a/src/messaging/message_manager.rs +++ b/src/messaging/message_manager.rs @@ -1,9 +1,11 @@ use super::{ - super::web::shared_state::{SharedState, Status}, + super::{ + config::EncryptionMode, + web::shared_state::{SharedState, Status}, + }, crypto::CipherAlgo, handshake::{self, HandshakeMsg}, }; -use crate::config::EncryptionMode; use anyhow::{Result, bail}; use serde::{Deserialize, Serialize}; use std::{net::SocketAddr, sync::Arc}; @@ -12,27 +14,26 @@ use tokio::{ net::UdpSocket, }; use tokio_kcp::{KcpConfig, KcpNoDelayConfig, KcpStream}; -use tracing::{error, info, warn}; +use tracing::{debug, error, info, warn}; -/// Manages the lifecycle of a P2P connection, handling the transition from raw UDP to reliable KCP. +/// Manages P2P connection lifecycle from raw UDP to reliable KCP. /// -/// This struct is the central controller for: -/// 1. **Handshaking**: coordinating UDP hole punching via the `handshake` module. -/// 2. **Upgrading**: Converting the raw UDP socket into a reliable `KcpStream` without losing the underlying connection. -/// 3. **Teardown**: Safely closing the KCP stream while preserving the shared socket if needed. +/// Responsibilities: +/// 1. **Handshaking**: Coordinates UDP hole punching via the `handshake` module. +/// 2. **Upgrading**: Converts raw UDP socket to reliable `KcpStream`. +/// 3. **Teardown**: Safely closes KCP stream while preserving shared socket. #[derive(Debug)] pub struct MessageManager { - /// The shared UDP socket used for both initial discovery and the eventual KCP stream. + /// Shared UDP socket for discovery and KCP stream. client_socket: Arc, - /// Shared application state for updating UI/Status. + /// Shared application state for UI updates. state: SharedState, - /// The address of the connected peer. Only set after a successful handshake. + /// Connected peer address. Set after successful handshake. peer_addr: Option, - /// The active reliable stream. None until `upgrade_to_kcp` is called. + /// Active reliable stream. None until `upgrade_to_kcp` is called. kcp_stream: Option, - // --- ENCRYPTION FIELDS --- - /// The session's encryption engine. + /// Session encryption engine. cipher: Option, /// Transmit nonce counter (strictly increasing). tx_nonce: u64, @@ -40,22 +41,22 @@ pub struct MessageManager { rx_nonce: u64, } -/// Represents a message being sent/received to/from a peer. +/// Represents a message sent/received to/from a peer. #[derive(Serialize, Deserialize, Debug)] pub enum StreamMessage { - /// Regular chat content + /// Regular chat content. Text(String), - /// Signal to close connection + /// Signal to close connection. Bye, } impl MessageManager { - /// Creates a new `MessageManager` in a disconnected state. + /// Creates a new `MessageManager` in disconnected state. /// /// # Arguments /// - /// * `client_socket` - The local UDP socket bound to a specific port. - /// * `state` - Reference to the application's shared state store. + /// * `client_socket` - Local UDP socket bound to a specific port. + /// * `state` - Reference to shared application state. pub fn new(client_socket: Arc, state: SharedState) -> Self { Self { client_socket, @@ -68,28 +69,28 @@ impl MessageManager { } } - /// Initiates the connection handshake with a target peer. + /// Initiates connection handshake with target peer. /// - /// This method blocks (async) until the handshake succeeds or times out. - /// It handles the `Punching` -> `Connected` state transitions in the shared state. + /// Blocks until handshake succeeds or times out. + /// Handles `Punching` -> `Connected` state transitions. /// /// # Arguments /// - /// * `peer_addr` - The public IP/Port of the target peer. - /// * `timeout_secs` - Max time to wait for the handshake protocol to complete. + /// * `peer_addr` - Public IP/Port of target peer. + /// * `timeout_secs` - Maximum wait time for handshake completion. /// * `mode` - Preferred encryption mode. /// /// # Returns /// /// * `Ok(())` - Handshake succeeded; `self.peer_addr` is set. - /// * `Err` - Handshake failed; State reset to `Disconnected`. + /// * `Err` - Handshake failed; state reset to `Disconnected`. pub async fn handshake( &mut self, peer_addr: SocketAddr, timeout_secs: u64, mode: EncryptionMode, ) -> Result<()> { - info!("Initializing MessageManager for peer {}", peer_addr); + debug!("Initiating handshake with peer {}", peer_addr); // Call the standalone handshake function with 5 arguments match handshake::handshake( @@ -102,10 +103,7 @@ impl MessageManager { .await { Ok(session) => { - info!( - "Handshake successful. Session Key ID: {}", - session.fingerprint - ); + info!("Handshake complete, fingerprint: {}", session.fingerprint); self.peer_addr = Some(peer_addr); // Store the Cipher and Reset Nonces @@ -128,23 +126,23 @@ impl MessageManager { } } - /// Upgrades the existing raw UDP connection to a reliable KCP stream. + /// Upgrades existing raw UDP connection to reliable KCP stream. /// - /// This utilizes "Turbo Mode" configuration for low latency: + /// Uses "Turbo Mode" configuration for low latency: /// - NoDelay: enabled /// - Update Interval: 10ms /// - Resend: 2 (fast retransmission) /// - No Congestion Control (NC): enabled - /// - Windows: 1024 packets (allows higher throughput) + /// - Windows: 1024 packets (higher throughput) /// - MTU: 1400 (safe default for UDP) /// /// # Errors /// - /// Returns an error if the handshake has not been performed yet (`peer_addr` is None) - /// or if the socket cloning fails. + /// Returns error if handshake not performed yet (`peer_addr` is None) + /// or if socket cloning fails. pub async fn upgrade_to_kcp(&mut self) -> Result<()> { if let Some(peer_addr) = self.peer_addr { - info!("Upgrading connection to KCP with {}", peer_addr); + debug!("Upgrading connection to KCP with {}", peer_addr); // Configure KCP for low-latency let config = KcpConfig { @@ -166,10 +164,10 @@ impl MessageManager { self.kcp_stream = Some(KcpStream::connect_with_socket(&config, socket, peer_addr).await?); - info!("KCP upgrade successful."); + info!("KCP upgrade complete"); Ok(()) } else { - bail!("Handshake not established.") + bail!("Handshake not established") } } @@ -323,11 +321,10 @@ impl MessageManager { /// /// # Arguments /// - /// * `send_bye` - If true, sends Bye message to peer before cleanup - // FIXED: Added allow attribute here to fix linter error + /// * `send_bye` - If true, sends Bye message to peer before cleanup. #[allow(clippy::collapsible_if)] async fn disconnect_internal(&mut self, send_bye: bool) -> Result<()> { - info!("Initiating graceful disconnect (send_bye: {})", send_bye); + debug!("Initiating disconnect (send_bye: {})", send_bye); // Send Bye message to peer only if requested if send_bye { @@ -338,7 +335,7 @@ impl MessageManager { if self.kcp_stream.is_some() && self.cipher.is_some() { if let Ok(bye_packet) = bincode::serialize(&StreamMessage::Bye) { if self.send_secure(&bye_packet).await.is_ok() { - info!("Sent Encrypted Bye via KCP"); + debug!("Sent encrypted Bye via KCP"); sent_via_kcp = true; } } @@ -348,7 +345,7 @@ impl MessageManager { if !sent_via_kcp { let udp_bye = bincode::serialize(&HandshakeMsg::Bye)?; match self.client_socket.send_to(&udp_bye, peer_addr).await { - Ok(_) => info!("Sent HandshakeMsg::Bye via UDP"), + Ok(_) => debug!("Sent HandshakeMsg::Bye via UDP"), Err(e) => warn!("Failed to send Bye via UDP: {}", e), } } @@ -381,26 +378,26 @@ impl MessageManager { Ok(()) } - /// Closes the active KCP stream gracefully. + /// Closes active KCP stream gracefully. /// - /// This method: - /// 1. Takes the stream out of the struct (setting `self.kcp_stream` to `None`). - /// 2. Sends a termination signal (shutdown) to the peer. - /// 3. Drops the stream, closing the *cloned* file descriptor. + /// Process: + /// 1. Takes stream out of struct (setting `self.kcp_stream` to `None`). + /// 2. Sends termination signal (shutdown) to peer. + /// 3. Drops stream, closing cloned file descriptor. /// - /// The original `client_socket` remains active. + /// Original `client_socket` remains active. #[allow(dead_code)] pub async fn close_kcp(&mut self) -> Result<()> { if let Some(mut stream) = self.kcp_stream.take() { - info!("Initiating KCP stream shutdown..."); + debug!("Shutting down KCP stream"); - // Attempt graceful shutdown. Log errors but not fail the function + // Attempt graceful shutdown. Log errors but don't fail function if let Err(e) = stream.shutdown().await { - error!("Error during KCP shutdown: {}", e); + warn!("KCP shutdown error: {}", e); } else { - info!("KCP stream shutdown completed."); + debug!("KCP stream shutdown complete"); } - // Stream is dropped here, closing the cloned FD. + // Stream is dropped here, closing cloned FD } Ok(()) } @@ -452,10 +449,7 @@ mod tests { // Should fail because peer_addr is None (handshake not run) let result = manager.upgrade_to_kcp().await; assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Handshake not established." - ); + assert_eq!(result.unwrap_err().to_string(), "Handshake not established"); } #[tokio::test] @@ -507,10 +501,10 @@ mod tests { // 4. Drop the cloned socket explicitly drop(cloned_sock); - // 5. Verify the original socket is still alive and working - // If mem::forget was missed in implementation, this would fail/panic because the FD would be closed + // 5. Verify original socket is still alive and working + // If mem::forget was missed in implementation, this would fail/panic because FD would be closed let test_payload = b"ping"; - // We send to ourselves just to check if the socket write operation fails immediately + // Send to self to check if socket write operation fails immediately let send_result = socket_arc.send_to(test_payload, "127.0.0.1:8080").await; assert!( diff --git a/src/net.rs b/src/net.rs index 4f4ae65..78f2bca 100644 --- a/src/net.rs +++ b/src/net.rs @@ -1,6 +1,6 @@ //! Network utilities for GhostLink. //! -//! Handles NAT traversal and public IP discovery using STUN. +//! Provides NAT traversal and public IP discovery using STUN. use super::web::shared_state::NatType; use anyhow::{Context, Result, bail}; @@ -16,25 +16,24 @@ use tokio::{ }; use tracing::debug; -/// The duration to wait for a STUN response before timing out. +/// Duration to wait for STUN response before timing out. const STUN_TIMEOUT: Duration = Duration::from_secs(3); -/// Resolves the local IP address using a DNS server. +/// Resolves local IP address using DNS server. /// -/// Connecting to a remote address causes the OS -/// to select the appropriate local interface and IP. +/// Connecting to remote address causes OS to select appropriate local interface and IP. /// /// # Returns /// -/// * `Ok(SocketAddr)` - Local IP and port -/// * `Err` - If resolution fails +/// * `Ok(SocketAddr)` - Local IP and port. +/// * `Err` - Resolution failed. pub async fn get_local_ip(local_port: u16) -> Result { // Connect to Google's public DNS let socket = UdpSocket::bind(("0.0.0.0", 0)).await?; socket.connect("8.8.8.8:80").await?; let local_addr = socket.local_addr()?; - // Replace the ephemeral port with the actual listening port + // Replace ephemeral port with actual listening port let local_ip = match local_addr.ip() { IpAddr::V4(ip) => SocketAddr::new(IpAddr::V4(ip), local_port), IpAddr::V6(ip) => SocketAddr::new(IpAddr::V6(ip), local_port), @@ -45,22 +44,23 @@ pub async fn get_local_ip(local_port: u16) -> Result { /// Discovers public IP and port using STUN. /// -/// Workflow: -/// 1. Resolves STUN server DNS -/// 2. Sends BINDING_REQUEST -/// 3. Waits for response (3 second timeout) -/// 4. Validates transaction ID -/// 5. Extracts public address +/// # Workflow +/// +/// 1. Resolves STUN server DNS. +/// 2. Sends BINDING_REQUEST. +/// 3. Waits for response (3 second timeout). +/// 4. Validates transaction ID. +/// 5. Extracts public address. /// /// # Arguments /// -/// * `socket` - Bound UDP socket -/// * `stun_server` - STUN server address (e.g., "stun.l.google.com:19302") +/// * `socket` - Bound UDP socket. +/// * `stun_server` - STUN server address (e.g., "stun.l.google.com:19302"). /// /// # Returns /// -/// * `Ok(SocketAddr)` - Public IP and port -/// * `Err` - If DNS, network, or STUN validation fails +/// * `Ok(SocketAddr)` - Public IP and port. +/// * `Err` - DNS, network, or STUN validation failed. pub async fn resolve_public_ip( socket: &UdpSocket, stun_server: impl AsRef, @@ -68,19 +68,18 @@ pub async fn resolve_public_ip( let stun_server = stun_server.as_ref(); debug!("Querying STUN server: {}", stun_server); - // 1. Determine our socket type (IPv4 or IPv6) + // 1. Determine socket type (IPv4 or IPv6) let local_addr = socket .local_addr() .context("Could not get local socket address")?; let is_ipv4_socket = local_addr.is_ipv4(); - // 2. Resolve DNS for the STUN server. + // 2. Resolve DNS for STUN server let mut addrs = tokio::net::lookup_host(stun_server) .await .context(format!("Failed to resolve DNS for {}", stun_server))?; - // 3. INTELLIGENT FILTERING (Fixed Clippy Lint): - // Use .find() instead of .filter().next() + // 3. Filter addresses compatible with socket type let target_addr = addrs .find(|addr| { if is_ipv4_socket { @@ -90,26 +89,25 @@ pub async fn resolve_public_ip( } }) .context(format!( - "STUN server {} has no addresses compatible with your socket (Protocol Mismatch)", + "STUN server {} has no addresses compatible with socket (Protocol Mismatch)", stun_server ))?; - // Build the STUN binding request + // Build STUN binding request let mut msg = Message::new(); msg.build(&[Box::::default(), Box::new(BINDING_REQUEST)])?; let expected_tx_id = msg.transaction_id; - // 4. Send the request + // 4. Send request socket .send_to(&msg.raw, target_addr) .await .context("Failed to send STUN request")?; - // 5. Wait for response with timeout + // 5. Wait for response with timeout (UDP packets can be lost) let mut buf = [0u8; 1024]; - // Use timeout as UDP packets can be lost. let (len, sender_addr) = timeout(STUN_TIMEOUT, socket.recv_from(&mut buf)) .await .context("STUN request timed out")? @@ -129,7 +127,7 @@ pub async fn resolve_public_ip( ); } - // 7. Extract the public IP + // 7. Extract public IP let mut xor_addr = XorMappedAddress::default(); xor_addr .get_from(&response) @@ -141,32 +139,32 @@ pub async fn resolve_public_ip( Ok(public_addr) } -/// Detects NAT type by querying a second STUN server. +/// Detects NAT type by querying second STUN server. /// -/// Compares the public port from two different STUN servers: -/// - Same port → Cone NAT (P2P-friendly) -/// - Different port → Symmetric NAT (P2P-difficult) +/// Compares public port from two different STUN servers: +/// - Same port → Cone NAT (P2P-friendly). +/// - Different port → Symmetric NAT (P2P-difficult). /// /// # Arguments /// -/// * `socket` - Bound UDP socket -/// * `stun_server` - Second STUN server address -/// * `prev_addr` - Address from first STUN query +/// * `socket` - Bound UDP socket. +/// * `stun_server` - Second STUN server address. +/// * `prev_addr` - Address from first STUN query. /// /// # Returns /// -/// NAT type: Cone, Symmetric, or Unknown +/// NAT type: Cone, Symmetric, or Unknown. pub async fn get_nat_type( socket: &UdpSocket, stun_server: impl AsRef, prev_addr: SocketAddr, ) -> NatType { - // resolve the public IP using new STUN server + // Resolve public IP using new STUN server resolve_public_ip(socket, stun_server).await.map_or_else( - // return `Unknown` if any error. + // Return Unknown if any error |_| NatType::Unknown, |public_ip| { - // return type of NAT based on response. + // Return NAT type based on response if prev_addr == public_ip { NatType::Cone } else { diff --git a/src/web/shared_state.rs b/src/web/shared_state.rs index 0748ac8..9421d22 100644 --- a/src/web/shared_state.rs +++ b/src/web/shared_state.rs @@ -129,14 +129,14 @@ impl AppState { self.broadcast_status_change(message, timeout); } - /// Updates the security details for the current session. - /// This is called by the handshake module upon successful key exchange. + /// Updates security details for current session. + /// Called by handshake module upon successful key exchange. pub fn set_security_info(&mut self, fingerprint: String, algorithm: String) { self.fingerprint = Some(fingerprint); self.encryption_algo = Some(algorithm); - // Note: We do not broadcast here immediately; - // the handshake usually calls set_status(Connected) right after, - // which triggers the broadcast with this new data included. + // Note: Does not broadcast immediately. + // Handshake typically calls set_status(Connected) right after, + // which triggers broadcast with this new data included. } /// Broadcasts current state to all active listeners. diff --git a/src/web/web_server.rs b/src/web/web_server.rs index 829d270..8ea1ab1 100644 --- a/src/web/web_server.rs +++ b/src/web/web_server.rs @@ -37,7 +37,6 @@ use tracing::{debug, error, info}; /// /// * `shared_state` - Thread-safe application state /// * `port` - Port to listen on -// ADDED 'pub' to make it accessible from main.rs pub async fn start_web_server(shared_state: SharedState, port: u16) -> Result<()> { let app = router(shared_state); let addr = SocketAddr::from(([0, 0, 0, 0], port)); @@ -79,7 +78,6 @@ async fn get_state(State(state): State) -> impl IntoResponse { struct ConnectionRequest { ip: String, port: u16, - // ADDED: Optional encryption mode (default to ChaCha20Poly1305 if missing) #[serde(default = "default_encryption_mode")] mode: EncryptionMode, } @@ -123,16 +121,8 @@ async fn connect_peer( guard.set_peer_ip(peer_addr, Some("Target set via API".into()), None); } - // 3. Send Command to Controller - // UPDATED: Now passing the encryption mode preference - // Note: You need to update Command enum in shared_state.rs to support this payload if desired, - // otherwise the controller uses its config default. For now, we assume the controller - // reads from config or shared state, but this handler sets up the peer_ip correctly. - - // NOTE: Ideally, Command::ConnectPeer should carry (peer_addr, mode). - // If your Command enum is strictly `ConnectPeer` (no data), the controller reads - // peer_addr from SharedState. We can just send the command. - + // 3. Send command to controller + // Controller reads peer_addr from SharedState let cmd_tx = state.read().await.cmd_tx().clone(); if let Err(e) = cmd_tx.send(Command::ConnectPeer).await { error!("Failed to send ConnectPeer command: {}", e); From 9bb6bea2121aef577854150adad73aa33a09ece0 Mon Sep 17 00:00:00 2001 From: pushkar-gr Date: Fri, 2 Jan 2026 08:07:14 +0530 Subject: [PATCH 7/7] feat(tests) --- src/config.rs | 3 +- src/messaging/crypto.rs | 193 ++++++++++++++++++++++++++++++ src/messaging/handshake.rs | 166 ++++++++++++++++++++++++++ src/messaging/message_manager.rs | 67 +++++++++++ src/web/shared_state.rs | 155 ++++++++++++++++++++++++ tests/integration_tets.rs | 199 +++++++++++++++++++++++++++++++ 6 files changed, 782 insertions(+), 1 deletion(-) create mode 100644 tests/integration_tets.rs diff --git a/src/config.rs b/src/config.rs index 06dd5ef..091d738 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,11 +1,12 @@ use serde::{Deserialize, Serialize}; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum EncryptionMode { ChaCha20Poly1305, Aes256Gcm, } + #[derive(Debug, Clone)] -#[allow(dead_code)] pub struct Config { pub client_port: u16, pub stun_server: String, diff --git a/src/messaging/crypto.rs b/src/messaging/crypto.rs index 50df44a..551ea50 100644 --- a/src/messaging/crypto.rs +++ b/src/messaging/crypto.rs @@ -246,4 +246,197 @@ mod tests { let decrypted = session.cipher.decrypt(nonce, &encrypted).unwrap(); assert_eq!(decrypted, plaintext); } + + #[test] + fn test_aes256_roundtrip() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::Aes256Gcm, + alice.public.to_bytes(), + ) + .unwrap(); + + let nonce = 98765u64; + let plaintext = b"Testing AES-256-GCM encryption"; + + let encrypted = session.cipher.encrypt(nonce, plaintext).unwrap(); + assert_ne!(encrypted, plaintext); + + let decrypted = session.cipher.decrypt(nonce, &encrypted).unwrap(); + assert_eq!(decrypted, plaintext); + } + + #[test] + fn test_encryption_with_different_nonces() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + let plaintext = b"Same message"; + let nonce1 = 1u64; + let nonce2 = 2u64; + + let encrypted1 = session.cipher.encrypt(nonce1, plaintext).unwrap(); + let encrypted2 = session.cipher.encrypt(nonce2, plaintext).unwrap(); + + // Same plaintext with different nonces should produce different ciphertexts + assert_ne!(encrypted1, encrypted2); + + // Both should decrypt correctly + let decrypted1 = session.cipher.decrypt(nonce1, &encrypted1).unwrap(); + let decrypted2 = session.cipher.decrypt(nonce2, &encrypted2).unwrap(); + assert_eq!(decrypted1, plaintext); + assert_eq!(decrypted2, plaintext); + } + + #[test] + fn test_decryption_with_wrong_nonce_fails() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + let plaintext = b"Secret message"; + let nonce_encrypt = 100u64; + let nonce_decrypt = 200u64; + + let encrypted = session.cipher.encrypt(nonce_encrypt, plaintext).unwrap(); + + // Decrypting with wrong nonce should fail + let result = session.cipher.decrypt(nonce_decrypt, &encrypted); + assert!(result.is_err()); + } + + #[test] + fn test_tampered_ciphertext_fails() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + let plaintext = b"Important data"; + let nonce = 42u64; + + let mut encrypted = session.cipher.encrypt(nonce, plaintext).unwrap(); + + // Tamper with ciphertext + if !encrypted.is_empty() { + encrypted[0] ^= 0xFF; + } + + // Decryption should fail due to authentication tag mismatch + let result = session.cipher.decrypt(nonce, &encrypted); + assert!(result.is_err()); + } + + #[test] + fn test_large_message_encryption() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + // Test with larger message (1KB) + let plaintext = vec![0x42u8; 1024]; + let nonce = 999u64; + + let encrypted = session.cipher.encrypt(nonce, &plaintext).unwrap(); + let decrypted = session.cipher.decrypt(nonce, &encrypted).unwrap(); + + assert_eq!(decrypted, plaintext); + } + + #[test] + fn test_empty_message_encryption() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + let plaintext = b""; + let nonce = 1u64; + + let encrypted = session.cipher.encrypt(nonce, plaintext).unwrap(); + let decrypted = session.cipher.decrypt(nonce, &encrypted).unwrap(); + + assert_eq!(decrypted, plaintext); + } + + #[test] + fn test_fingerprint_generation() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + let session = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + // Fingerprint should be in format "XX XX XX" (8 chars with spaces) + assert_eq!(session.fingerprint.len(), 8); + + // Should match pattern of two hex digits, space, two hex digits, space, two hex digits + let parts: Vec<&str> = session.fingerprint.split(' ').collect(); + assert_eq!(parts.len(), 3); + for part in parts { + assert_eq!(part.len(), 2); + assert!(part.chars().all(|c| c.is_ascii_hexdigit())); + } + } + + #[test] + fn test_same_keypair_same_mode_same_fingerprint() { + let alice = KeyPair::generate(); + let bob = KeyPair::generate(); + + let session1 = derive_session( + alice.private.clone(), + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + let session2 = derive_session( + alice.private, + bob.public.to_bytes(), + EncryptionMode::ChaCha20Poly1305, + alice.public.to_bytes(), + ) + .unwrap(); + + // Same keys and mode should produce same fingerprint + assert_eq!(session1.fingerprint, session2.fingerprint); + } } diff --git a/src/messaging/handshake.rs b/src/messaging/handshake.rs index 08ec39b..1f0577b 100644 --- a/src/messaging/handshake.rs +++ b/src/messaging/handshake.rs @@ -592,4 +592,170 @@ mod tests { assert_eq!(state_a.read().await.status, Status::Connected); assert_eq!(state_b.read().await.status, Status::Connected); } + + #[tokio::test] + async fn test_handshake_with_aes256_mode() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + let state_b = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + // Test handshake with AES-256-GCM encryption mode + let socket_a_clone = socket_a.clone(); + let state_a_clone = state_a.clone(); + let handle_a = tokio::spawn(async move { + handshake( + socket_a_clone, + addr_b, + state_a_clone, + 5, + EncryptionMode::Aes256Gcm, + ) + .await + }); + + let socket_b_clone = socket_b.clone(); + let state_b_clone = state_b.clone(); + let handle_b = tokio::spawn(async move { + handshake( + socket_b_clone, + addr_a, + state_b_clone, + 5, + EncryptionMode::Aes256Gcm, + ) + .await + }); + + let result_a = handle_a.await.unwrap(); + let result_b = handle_b.await.unwrap(); + + assert!( + result_a.is_ok(), + "Peer A should complete handshake with AES" + ); + assert!( + result_b.is_ok(), + "Peer B should complete handshake with AES" + ); + } + + #[tokio::test] + async fn test_handshake_sets_fingerprint() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + let state_b = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + let socket_a_clone = socket_a.clone(); + let state_a_clone = state_a.clone(); + let handle_a = tokio::spawn(async move { + handshake( + socket_a_clone, + addr_b, + state_a_clone, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await + }); + + let socket_b_clone = socket_b.clone(); + let state_b_clone = state_b.clone(); + let handle_b = tokio::spawn(async move { + handshake( + socket_b_clone, + addr_a, + state_b_clone, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await + }); + + let _ = handle_a.await.unwrap(); + let _ = handle_b.await.unwrap(); + + // Both peers should have fingerprint set + let fp_a = state_a.read().await.fingerprint.clone(); + let fp_b = state_b.read().await.fingerprint.clone(); + + assert!(fp_a.is_some(), "Peer A should have fingerprint"); + assert!(fp_b.is_some(), "Peer B should have fingerprint"); + assert_eq!(fp_a, fp_b, "Fingerprints should match"); + } + + #[tokio::test] + async fn test_handshake_with_very_short_timeout() { + let socket_a = bind_local().await; + let state_a = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + + // Use very short timeout that will likely fail + let result = handshake( + socket_a, + addr_a, // Wrong address (self) + state_a, + 1, // 1 second timeout + EncryptionMode::ChaCha20Poly1305, + ) + .await; + + // Should timeout or fail + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_handshake_status_transitions() { + let socket_a = bind_local().await; + let socket_b = bind_local().await; + let state_a = create_dummy_state(); + let state_b = create_dummy_state(); + + let addr_a = socket_a.local_addr().unwrap(); + let addr_b = socket_b.local_addr().unwrap(); + + // Check initial status + assert_eq!(state_a.read().await.status, Status::Disconnected); + + let socket_a_clone = socket_a.clone(); + let state_a_clone = state_a.clone(); + let handle_a = tokio::spawn(async move { + handshake( + socket_a_clone, + addr_b, + state_a_clone, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await + }); + + let socket_b_clone = socket_b.clone(); + let state_b_clone = state_b.clone(); + let handle_b = tokio::spawn(async move { + handshake( + socket_b_clone, + addr_a, + state_b_clone, + 5, + EncryptionMode::ChaCha20Poly1305, + ) + .await + }); + + let _ = handle_a.await.unwrap(); + let _ = handle_b.await.unwrap(); + + // Status should transition to Connected + assert_eq!(state_a.read().await.status, Status::Connected); + assert_eq!(state_b.read().await.status, Status::Connected); + } } diff --git a/src/messaging/message_manager.rs b/src/messaging/message_manager.rs index 96898bb..bc98ef1 100644 --- a/src/messaging/message_manager.rs +++ b/src/messaging/message_manager.rs @@ -544,4 +544,71 @@ mod tests { let state_guard = manager.state.read().await; assert_eq!(state_guard.status, Status::Disconnected); } + + #[tokio::test] + async fn test_is_connected_false_initially() { + let manager = create_test_manager().await; + assert!(!manager.is_connected()); + } + + #[tokio::test] + async fn test_nonce_initialization() { + let manager = create_test_manager().await; + assert_eq!(manager.tx_nonce, 0); + assert_eq!(manager.rx_nonce, 0); + } + + #[tokio::test] + async fn test_send_text_without_kcp_fails() { + let mut manager = create_test_manager().await; + + // Try to send without establishing KCP + let result = manager.send_text("test message".to_string()).await; + assert!(result.is_err()); + + let error_msg = result.unwrap_err().to_string(); + assert!(error_msg.contains("KCP") || error_msg.contains("stream")); + } + + #[tokio::test] + async fn test_multiple_disconnects() { + let mut manager = create_test_manager().await; + + // Multiple disconnects should be idempotent + assert!(manager.disconnect().await.is_ok()); + assert!(manager.disconnect().await.is_ok()); + assert!(manager.disconnect().await.is_ok()); + } + + #[tokio::test] + async fn test_disconnect_clears_cipher() { + let mut manager = create_test_manager().await; + + // Manually set cipher (simulating encryption setup) + manager.peer_addr = Some("127.0.0.1:9999".parse().unwrap()); + + // Disconnect should clear cipher + manager.disconnect().await.unwrap(); + + assert!(manager.cipher.is_none()); + } + + #[tokio::test] + async fn test_disconnect_clears_peer_addr() { + let mut manager = create_test_manager().await; + manager.peer_addr = Some("127.0.0.1:8888".parse().unwrap()); + + manager.disconnect().await.unwrap(); + + assert!(manager.peer_addr.is_none()); + } + + #[tokio::test] + async fn test_close_kcp_with_none_stream() { + let mut manager = create_test_manager().await; + + // close_kcp with no stream should not fail + let result = manager.close_kcp().await; + assert!(result.is_ok()); + } } diff --git a/src/web/shared_state.rs b/src/web/shared_state.rs index 9421d22..b2424db 100644 --- a/src/web/shared_state.rs +++ b/src/web/shared_state.rs @@ -261,3 +261,158 @@ pub enum Command { /// Disconnect from current peer Disconnect, } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{IpAddr, Ipv4Addr}; + + fn create_test_state() -> AppState { + let (cmd_tx, _cmd_rx) = mpsc::channel(32); + let (event_tx, _event_rx) = broadcast::channel(32); + AppState::new(cmd_tx, event_tx) + } + + #[test] + fn test_app_state_initialization() { + let state = create_test_state(); + + assert_eq!(state.local_ip, None); + assert_eq!(state.public_ip, None); + assert_eq!(state.nat_type, NatType::Unknown); + assert_eq!(state.status, Status::Disconnected); + assert_eq!(state.peer_ip, None); + assert_eq!(state.fingerprint, None); + assert_eq!(state.encryption_algo, None); + } + + #[test] + fn test_set_local_ip() { + let mut state = create_test_state(); + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100)), 8080); + + state.set_local_ip(addr, Some("Local IP set".into()), None); + + assert_eq!(state.local_ip, Some(addr)); + } + + #[test] + fn test_set_public_ip() { + let mut state = create_test_state(); + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 5678); + + state.set_public_ip(addr, Some("Public IP resolved".into()), None); + + assert_eq!(state.public_ip, Some(addr)); + } + + #[test] + fn test_set_nat_type() { + let mut state = create_test_state(); + + state.set_nat_type(NatType::Cone, Some("NAT detected".into()), None); + assert_eq!(state.nat_type, NatType::Cone); + + state.set_nat_type(NatType::Symmetric, None, None); + assert_eq!(state.nat_type, NatType::Symmetric); + } + + #[test] + fn test_set_status() { + let mut state = create_test_state(); + + state.set_status(Status::Punching, Some("Connecting...".into()), Some(30)); + assert_eq!(state.status, Status::Punching); + + state.set_status(Status::Connected, Some("Connected".into()), None); + assert_eq!(state.status, Status::Connected); + } + + #[test] + fn test_set_peer_ip() { + let mut state = create_test_state(); + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 9999); + + state.set_peer_ip(addr, Some("Peer set".into()), None); + + assert_eq!(state.peer_ip, Some(addr)); + } + + #[test] + fn test_set_security_info() { + let mut state = create_test_state(); + + state.set_security_info("abcd1234".to_string(), "ChaCha20-Poly1305".to_string()); + + assert_eq!(state.fingerprint, Some("abcd1234".to_string())); + assert_eq!(state.encryption_algo, Some("ChaCha20-Poly1305".to_string())); + } + + #[tokio::test] + async fn test_event_subscription() { + let (cmd_tx, _cmd_rx) = mpsc::channel(32); + let (event_tx, _event_rx) = broadcast::channel(32); + let state = AppState::new(cmd_tx, event_tx); + + let mut rx = state.subscribe_events(); + + // Send a test event + state.add_message("Test message".to_string(), true); + + // Should receive the event + let event = tokio::time::timeout(std::time::Duration::from_millis(100), rx.recv()).await; + + assert!(event.is_ok()); + } + + #[tokio::test] + async fn test_command_channel() { + let (cmd_tx, mut cmd_rx) = mpsc::channel(32); + let (event_tx, _event_rx) = broadcast::channel(32); + let state = AppState::new(cmd_tx.clone(), event_tx); + + // Send a command + state.cmd_tx().send(Command::ConnectPeer).await.unwrap(); + + // Should receive the command + let cmd = cmd_rx.recv().await; + assert!(cmd.is_some()); + } + + #[test] + fn test_nat_type_equality() { + assert_eq!(NatType::Unknown, NatType::Unknown); + assert_eq!(NatType::Cone, NatType::Cone); + assert_eq!(NatType::Symmetric, NatType::Symmetric); + + assert_ne!(NatType::Unknown, NatType::Cone); + assert_ne!(NatType::Cone, NatType::Symmetric); + } + + #[test] + fn test_status_equality() { + assert_eq!(Status::Disconnected, Status::Disconnected); + assert_eq!(Status::Punching, Status::Punching); + assert_eq!(Status::Connected, Status::Connected); + + assert_ne!(Status::Disconnected, Status::Punching); + assert_ne!(Status::Punching, Status::Connected); + } + + #[test] + fn test_clear_chat() { + let state = create_test_state(); + + // Should not panic + state.clear_chat(); + } + + #[test] + fn test_add_message() { + let state = create_test_state(); + + // Should not panic + state.add_message("Hello".to_string(), true); + state.add_message("World".to_string(), false); + } +} diff --git a/tests/integration_tets.rs b/tests/integration_tets.rs new file mode 100644 index 0000000..e08c9e5 --- /dev/null +++ b/tests/integration_tets.rs @@ -0,0 +1,199 @@ +//! Integration tests for GhostLink +//! +//! Tests end-to-end functionality including P2P connections, +//! encryption, and message exchange. + +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use tokio::net::UdpSocket; +use tokio::sync::{RwLock, broadcast, mpsc}; +use tokio::time::{Duration, timeout}; + +#[tokio::test] +async fn test_udp_socket_creation() { + // Test that we can create and bind UDP sockets + let socket = UdpSocket::bind("127.0.0.1:0").await; + assert!(socket.is_ok()); + + let socket = socket.unwrap(); + let local_addr = socket.local_addr(); + assert!(local_addr.is_ok()); +} + +#[tokio::test] +async fn test_multiple_socket_binds() { + // Test that we can create multiple UDP sockets + let socket1 = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let socket2 = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + + let addr1 = socket1.local_addr().unwrap(); + let addr2 = socket2.local_addr().unwrap(); + + // Should have different ports + assert_ne!(addr1.port(), addr2.port()); +} + +#[tokio::test] +async fn test_udp_message_exchange() { + // Test basic UDP message exchange + let server = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_addr = server.local_addr().unwrap(); + + let client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + + // Send message from client to server + let sent = client.send_to(b"Hello", server_addr).await; + assert!(sent.is_ok()); + + // Receive on server + let mut buf = [0u8; 100]; + let result = timeout(Duration::from_secs(1), server.recv_from(&mut buf)).await; + + assert!(result.is_ok()); + let (len, _addr) = result.unwrap().unwrap(); + assert_eq!(&buf[..len], b"Hello"); +} + +#[tokio::test] +async fn test_concurrent_udp_operations() { + // Test that multiple UDP operations can happen concurrently + let socket1 = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + let socket2 = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + + let addr1 = socket1.local_addr().unwrap(); + let addr2 = socket2.local_addr().unwrap(); + + let s1 = socket1.clone(); + let s2 = socket2.clone(); + + // Send messages concurrently + let send1 = tokio::spawn(async move { s1.send_to(b"Message1", addr2).await }); + + let send2 = tokio::spawn(async move { s2.send_to(b"Message2", addr1).await }); + + let result1 = send1.await; + let result2 = send2.await; + + assert!(result1.is_ok()); + assert!(result2.is_ok()); +} + +#[tokio::test] +async fn test_broadcast_channel() { + // Test tokio broadcast channel functionality + let (tx, mut rx1) = broadcast::channel::(16); + let mut rx2 = tx.subscribe(); + + tx.send("Test message".to_string()).unwrap(); + + let msg1 = rx1.recv().await; + let msg2 = rx2.recv().await; + + assert!(msg1.is_ok()); + assert!(msg2.is_ok()); + assert_eq!(msg1.unwrap(), "Test message"); + assert_eq!(msg2.unwrap(), "Test message"); +} + +#[tokio::test] +async fn test_mpsc_channel() { + // Test tokio mpsc channel functionality + let (tx, mut rx) = mpsc::channel::(32); + + tx.send(42).await.unwrap(); + tx.send(100).await.unwrap(); + + let val1 = rx.recv().await; + let val2 = rx.recv().await; + + assert_eq!(val1, Some(42)); + assert_eq!(val2, Some(100)); +} + +#[tokio::test] +async fn test_rwlock_concurrent_reads() { + // Test that RwLock allows concurrent reads + let data = Arc::new(RwLock::new(42)); + + let d1 = data.clone(); + let d2 = data.clone(); + + let read1 = tokio::spawn(async move { + let value = d1.read().await; + *value + }); + + let read2 = tokio::spawn(async move { + let value = d2.read().await; + *value + }); + + let val1 = read1.await.unwrap(); + let val2 = read2.await.unwrap(); + + assert_eq!(val1, 42); + assert_eq!(val2, 42); +} + +#[tokio::test] +async fn test_rwlock_write_exclusion() { + // Test that RwLock write is exclusive + let data = Arc::new(RwLock::new(0)); + let (tx, mut rx) = mpsc::channel::<()>(1); + + let d1 = data.clone(); + let write1 = tokio::spawn(async move { + let mut value = d1.write().await; + *value += 1; + // Signal that first write has started + tx.send(()).await.unwrap(); + tokio::time::sleep(Duration::from_millis(10)).await; + *value + }); + + // Wait for first write to acquire lock + rx.recv().await.unwrap(); + + let d2 = data.clone(); + let write2 = tokio::spawn(async move { + let mut value = d2.write().await; + *value += 10; + *value + }); + + let val1 = write1.await.unwrap(); + let val2 = write2.await.unwrap(); + + // Second write should see result of first write + assert_eq!(val1, 1); + assert_eq!(val2, 11); +} + +#[tokio::test] +async fn test_socket_addr_parsing() { + // Test SocketAddr parsing + let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap(); + assert_eq!(addr.ip(), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))); + assert_eq!(addr.port(), 8080); + + let addr: SocketAddr = "0.0.0.0:0".parse().unwrap(); + assert_eq!(addr.ip(), IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))); + assert_eq!(addr.port(), 0); +} + +#[tokio::test] +async fn test_timeout_functionality() { + // Test timeout works correctly + let result = timeout( + Duration::from_millis(100), + tokio::time::sleep(Duration::from_secs(10)), + ) + .await; + + assert!(result.is_err()); + + let result = timeout(Duration::from_millis(100), async { 42 }).await; + + assert!(result.is_ok()); + assert_eq!(result.unwrap(), 42); +}