Improved cipher storage

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<Server>,
     mut socket: WebSocket,
-) -> anyhow::Result<Arc<EnclaveWebSocket>> {
+) -> anyhow::Result<EnclaveWebSocket> {
     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))
 }

src/protocol/initialize.rs

@@ -13,7 +13,7 @@ use crate::server::UserConnections;
 impl UserConnections {
     pub async fn initialize(
         server: &Arc<Server>,
-        socket: &Arc<EnclaveWebSocket>,
+        socket: &EnclaveWebSocket,
     ) -> anyhow::Result<(VerifyingKey, ClientMeta)> {
         let Some(ServerMethod::Initialize {
             public_key: public_key_string,

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<HashMap<u16, Arc<crate::ws::EnclaveWebSocket>>>,
+    pub cihper: Arc<Mutex<SessionCipher>>,
     pub voice: Mutex<Option<VoiceConnection>>,
 }
 
@@ -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);

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<SplitSink<WebSocket, Message>>,
     rx: Mutex<SplitStream<WebSocket>>,
-    cihper: Mutex<SessionCipher>,
+    pub cipher: Arc<Mutex<SessionCipher>>,
 }
 
 impl EnclaveWebSocket {
-    pub fn new(ws: WebSocket, cipher: SessionCipher) -> Self {
+    pub fn new(ws: WebSocket, cipher: Arc<Mutex<SessionCipher>>) -> 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()