enclave-chat/enclave-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Better websocket

Browse Source
This commit is contained in:
selimaj-dev
2026-08-27 16:18:19 +02:00
parent 02a444d6ed
commit a4d675bda9
10 changed files with 141 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/protocol/initialize.rs
+30 -44
View File
@@ -3,10 +3,9 @@ use std::{
time::{SystemTime, UNIX_EPOCH},
};
use axum::extract::ws::WebSocket;
use ed25519_dalek::{Signer, VerifyingKey};
use crate::server::Server;
use crate::{server::Server, ws::EnclaveWebSocket};
use super::*;
use crate::server::UserConnections;
@@ -14,23 +13,21 @@ use crate::server::UserConnections;
impl UserConnections {
pub async fn initialize(
server: &Arc<Server>,
mut socket: WebSocket,
) -> anyhow::Result<(WebSocket, VerifyingKey, ClientMeta)> {
socket: Arc<EnclaveWebSocket>,
) -> anyhow::Result<(Arc<EnclaveWebSocket>, VerifyingKey, ClientMeta)> {
let Some(ServerMethod::Initialize {
public_key: public_key_string,
signature,
timestamp,
hostname,
}) = read_socket(&mut socket).await?
}) = socket.read().await?
else {
send_socket(
&mut socket,
&ClientMethod::Error {
socket
.send(&ClientMethod::Error {
error: Cow::Borrowed("Initialization required"),
},
)
.await?;
})
.await?;
return Err(anyhow::anyhow!(
"Failed to initialize: Client sent the wrong method"
@@ -40,23 +37,20 @@ impl UserConnections {
let server_timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64;
if server_timestamp.saturating_sub(timestamp) > 2000 {
send_socket(
&mut socket,
&ClientMethod::Error {
socket
.send(&ClientMethod::Error {
error: Cow::Borrowed(
"Timestamp doesn't match, make sure it's in secs and is (<= 2secs)",
),
},
)
.await?;
})
.await?;
return Err(anyhow::anyhow!("Client tampstamp wasn't correct"));
}
if hostname != server.config.public_hostname || !server.config.hostnames.contains(&hostname)
{
send_socket(
&mut socket,
socket.send(
&ClientMethod::Error {
error: Cow::Owned(format!("Invalid Hostname, to avoid man-in-the-middle attacks, please use the correct hostname: {}", server.config.public_hostname)),
},
@@ -67,13 +61,11 @@ impl UserConnections {
}
let Ok(public_key) = crate::crypto::from_string(&public_key_string) else {
send_socket(
&mut socket,
&ClientMethod::Error {
socket
.send(&ClientMethod::Error {
error: Cow::Borrowed("Invalid public key"),
},
)
.await?;
})
.await?;
return Err(anyhow::anyhow!("Invalid public key"));
};
@@ -85,21 +77,18 @@ impl UserConnections {
)
.is_err()
{
send_socket(
&mut socket,
&ClientMethod::Error {
socket
.send(&ClientMethod::Error {
error: Cow::Borrowed("Invalid signature"),
},
)
.await?;
})
.await?;
return Err(anyhow::anyhow!("Invalid signature"));
}
{
send_socket(
&mut socket,
&ClientMethod::Initialized {
socket
.send(&ClientMethod::Initialized {
public_key: crate::crypto::to_string(&server.key.verifying_key()),
signature: crate::crypto::to_string_sig(&server.key.sign(
format!("{server_timestamp}@{hostname}@{public_key_string}").as_bytes(),
@@ -107,19 +96,16 @@ impl UserConnections {
timestamp: server_timestamp,
hostname,
},
)
.await?;
})
.await?;
}
let Some(ServerMethod::Meta(meta)) = read_socket(&mut socket).await? else {
send_socket(
&mut socket,
&ClientMethod::Error {
let Some(ServerMethod::Meta(meta)) = socket.read().await? else {
socket
.send(&ClientMethod::Error {
error: Cow::Borrowed("Expected meta"),
},
)
.await?;
})
.await?;
return Err(anyhow::anyhow!(
"Expected meta, client called another method"
src/protocol/message.rs
+7 -11
View File
@@ -1,17 +1,15 @@
use crate::data::messages::{MessageData, StoredMessage};
use crate::protocol::{ClientMethod, send_socket};
use crate::protocol::ClientMethod;
use crate::server::Server;
use axum::extract::ws::WebSocket;
use ed25519_dalek::{Verifier, VerifyingKey};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::Mutex;
pub async fn send_message(
server: &Arc<Server>,
verifying_key: VerifyingKey,
_socket: &Arc<Mutex<WebSocket>>,
_socket: &Arc<crate::ws::EnclaveWebSocket>,
message: MessageData,
channel_id: String,
) -> anyhow::Result<()> {
@@ -58,7 +56,7 @@ pub async fn send_message(
pub async fn get_messages(
server: &Arc<Server>,
_verifying_key: VerifyingKey,
socket: &Arc<Mutex<WebSocket>>,
socket: &Arc<crate::ws::EnclaveWebSocket>,
channel_id: String,
chunk: u32,
) -> anyhow::Result<()> {
@@ -68,13 +66,11 @@ pub async fn get_messages(
.message_store
.get_recent_messages(&channel_id, CHUNK_SIZE, chunk)?;
send_socket(
&mut *socket.lock().await,
&ClientMethod::Messages {
socket
.send(&ClientMethod::Messages {
messages: HashMap::from([(channel_id, messages)]),
},
)
.await?;
})
.await?;
Ok(())
}
src/protocol/mod.rs
+10 -62
View File
@@ -1,9 +1,7 @@
use std::{borrow::Cow, collections::HashMap, sync::Arc};
use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
use ed25519_dalek::VerifyingKey;
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use crate::{data::messages::StoredMessage, server::Server, types::ClientMeta};
@@ -110,22 +108,16 @@ pub enum ServerMethod {
pub async fn read_loop(
server: &Arc<Server>,
verifying_key: VerifyingKey,
socket: &Arc<Mutex<WebSocket>>,
socket: &Arc<crate::ws::EnclaveWebSocket>,
) -> anyhow::Result<()> {
let mut socket_lock = socket.lock().await;
while let Some(message) = read_socket(&mut *socket_lock).await? {
drop(socket_lock);
while let Some(message) = socket.read().await? {
match message {
ServerMethod::Initialize { .. } => {
send_socket(
&mut *socket.lock().await,
&ClientMethod::Error {
socket
.send(&ClientMethod::Error {
error: Cow::Borrowed("Already initialized"),
},
)
.await?;
})
.await?;
}
#[allow(unused_variables)]
@@ -181,14 +173,12 @@ pub async fn read_loop(
.await
.insert(pin, (verifying_key, channel_id.clone()));
send_socket(
&mut *socket.lock().await,
&ClientMethod::JoinVoice {
socket
.send(&ClientMethod::JoinVoice {
channel_id: channel_id.clone(),
pin,
},
)
.await?;
})
.await?;
}
server
@@ -199,49 +189,7 @@ pub async fn read_loop(
.await?;
}
}
socket_lock = socket.lock().await;
}
Ok(())
}
pub async fn read_socket(socket: &mut WebSocket) -> anyhow::Result<Option<ServerMethod>> {
match socket.recv().await.transpose()? {
Some(Message::Text(text)) => match serde_json::from_str(&text.to_string()) {
Ok(msg) => Ok(Some(msg)),
Err(e) => {
send_socket(
socket,
&ClientMethod::Error {
error: Cow::Owned(format!("Unable to parse message: {e}")),
},
)
.await?;
Ok(None)
}
},
Some(Message::Ping(v)) => {
socket.send(Message::Pong(v)).await?;
Ok(None)
}
Some(_) => Ok(None),
None => Ok(None),
}
}
pub async fn send_socket(socket: &mut WebSocket, message: &ClientMethod) -> anyhow::Result<()> {
socket
.send(Message::Text(Utf8Bytes::from(serde_json::to_string(
message,
)?)))
.await?;
Ok(())
}
src/protocol/user.rs
+3 -8
View File
@@ -1,23 +1,18 @@
use std::sync::Arc;
use axum::extract::ws::WebSocket;
use ed25519_dalek::VerifyingKey;
use tokio::sync::Mutex;
use crate::{
protocol::{ClientMethod, send_socket},
server::Server,
};
use crate::{protocol::ClientMethod, server::Server};
pub async fn get_users(
server: &Arc<Server>,
_verifying_key: VerifyingKey,
socket: &Arc<Mutex<WebSocket>>,
socket: &Arc<crate::ws::EnclaveWebSocket>,
pubkeys: Vec<String>,
) -> anyhow::Result<()> {
let users = server.user_store.get_users(&pubkeys).await?;
send_socket(&mut *socket.lock().await, &ClientMethod::Users { users }).await?;
socket.send(&ClientMethod::Users { users }).await?;
Ok(())
}