diff --git a/src/crypto.rs b/src/crypto.rs index b1dd985..a848622 100644 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -9,6 +9,7 @@ use curve25519_dalek::edwards::CompressedEdwardsY; use ed25519_dalek::{Signature, SigningKey, VerifyingKey}; use rand::rngs::OsRng; use sha2::{Digest, Sha512}; +use tokio::sync::Mutex; use x25519_dalek::{PublicKey as X25519Public, SharedSecret, StaticSecret as X25519Secret}; use crate::{server::Server, ws::EnclaveWebSocket}; @@ -138,7 +139,7 @@ impl SessionCipher { pub async fn crypto_handshake( server: &Arc, mut socket: WebSocket, -) -> anyhow::Result> { +) -> anyhow::Result { socket .send(axum::extract::ws::Message::Binary( server.x_keypair.0.to_bytes().to_vec().into(), @@ -160,7 +161,7 @@ pub async fn crypto_handshake( let shared_secret = server.x_keypair.1.diffie_hellman(&client_pubkey); - let cipher = SessionCipher::new(&shared_secret)?; + let cipher = Arc::new(Mutex::new(SessionCipher::new(&shared_secret)?)); - Ok(Arc::new(EnclaveWebSocket::new(socket, cipher))) + Ok(EnclaveWebSocket::new(socket, cipher)) } diff --git a/src/protocol/initialize.rs b/src/protocol/initialize.rs index 93b9f85..cb97f8d 100644 --- a/src/protocol/initialize.rs +++ b/src/protocol/initialize.rs @@ -13,7 +13,7 @@ use crate::server::UserConnections; impl UserConnections { pub async fn initialize( server: &Arc, - socket: &Arc, + socket: &EnclaveWebSocket, ) -> anyhow::Result<(VerifyingKey, ClientMeta)> { let Some(ServerMethod::Initialize { public_key: public_key_string, diff --git a/src/server.rs b/src/server.rs index 083d01f..201f93e 100644 --- a/src/server.rs +++ b/src/server.rs @@ -18,6 +18,7 @@ use tokio::{ }; use crate::{ + crypto::SessionCipher, data::{config::Config, messages::MessageStore, users::UserMetaStore}, protocol::{ClientMethod, read_loop}, types::ClientMeta, @@ -35,6 +36,7 @@ pub struct UserConnections { pub counter: AtomicU16, pub public_key: VerifyingKey, pub connections: Mutex>>, + pub cihper: Arc>, pub voice: Mutex>, } @@ -74,7 +76,7 @@ impl Server { let s = self.clone(); ws.on_upgrade(move |socket: WebSocket| async move { - let client = match crate::crypto::crypto_handshake(&s, socket).await { + let mut client = match crate::crypto::crypto_handshake(&s, socket).await { Ok(client) => client, Err(err) => { eprintln!("Failed to initialize crypto: {err}"); @@ -103,10 +105,15 @@ impl Server { counter: AtomicU16::new(0), connections: Mutex::new(HashMap::new()), voice: Mutex::new(None), + cihper: client.cipher.clone(), }) }) .clone(); + client.cipher = clients.cihper.clone(); + + let client = Arc::new(client); + let conid = clients .counter .fetch_add(1, std::sync::atomic::Ordering::Relaxed); diff --git a/src/ws.rs b/src/ws.rs index 07f2c8a..efae90e 100644 --- a/src/ws.rs +++ b/src/ws.rs @@ -1,4 +1,4 @@ -use std::borrow::Cow; +use std::{borrow::Cow, sync::Arc}; use axum::extract::ws::{Message, WebSocket}; use futures_util::{ @@ -15,17 +15,17 @@ use crate::{ pub struct EnclaveWebSocket { tx: Mutex>, rx: Mutex>, - cihper: Mutex, + pub cipher: Arc>, } impl EnclaveWebSocket { - pub fn new(ws: WebSocket, cipher: SessionCipher) -> Self { + pub fn new(ws: WebSocket, cipher: Arc>) -> Self { let (tx, rx) = ws.split(); Self { tx: Mutex::new(tx), rx: Mutex::new(rx), - cihper: Mutex::new(cipher), + cipher, } } @@ -45,7 +45,7 @@ impl EnclaveWebSocket { }, Some(Message::Binary(encrypted)) => { - let text = String::from_utf8(self.cihper.lock().await.decrypt(&encrypted)?)?; + let text = String::from_utf8(self.cipher.lock().await.decrypt(&encrypted)?)?; match serde_json::from_str(&text.to_string()) { Ok(msg) => Ok(Some(msg)), @@ -76,7 +76,7 @@ impl EnclaveWebSocket { pub async fn send(&self, message: &ClientMethod) -> anyhow::Result<()> { let text = serde_json::to_string(message)?; - let encrypted = self.cihper.lock().await.encrypt(text.as_bytes())?; + let encrypted = self.cipher.lock().await.encrypt(text.as_bytes())?; self.tx .lock()