109 lines
3.1 KiB
Rust
109 lines
3.1 KiB
Rust
use std::{
|
|
collections::HashMap,
|
|
sync::{Arc, atomic::AtomicU16},
|
|
};
|
|
|
|
use axum::{
|
|
extract::{WebSocketUpgrade, ws::WebSocket},
|
|
response::Response,
|
|
};
|
|
use ed25519_dalek::{SigningKey, VerifyingKey};
|
|
use tokio::sync::Mutex;
|
|
|
|
use crate::{
|
|
config::Config,
|
|
protocol::{ClientMethod, read_loop, send_socket},
|
|
types::ClientMeta,
|
|
};
|
|
|
|
pub struct UserConnections {
|
|
pub meta: ClientMeta,
|
|
pub counter: AtomicU16,
|
|
pub public_key: VerifyingKey,
|
|
pub connections: HashMap<u16, Arc<Mutex<WebSocket>>>,
|
|
}
|
|
|
|
pub struct Server {
|
|
pub key: SigningKey,
|
|
pub config: Config,
|
|
pub clients: Mutex<HashMap<VerifyingKey, UserConnections>>,
|
|
}
|
|
|
|
impl Server {
|
|
pub async fn new() -> anyhow::Result<Arc<Self>> {
|
|
Ok(Arc::new(Self {
|
|
key: crate::signature::get().await?,
|
|
config: Config::get().await?,
|
|
clients: Mutex::new(HashMap::new()),
|
|
}))
|
|
}
|
|
}
|
|
|
|
impl Server {
|
|
pub async fn ws_handler(self: &Arc<Self>, ws: WebSocketUpgrade) -> Response {
|
|
let s = self.clone();
|
|
|
|
ws.on_upgrade(move |socket: WebSocket| async move {
|
|
match UserConnections::initialize(&s, socket).await {
|
|
Ok((client, public_key, meta)) => {
|
|
let client = Arc::new(Mutex::new(client));
|
|
|
|
let mut clients_meta = s.clients.lock().await;
|
|
|
|
let clients =
|
|
clients_meta
|
|
.entry(public_key)
|
|
.or_insert_with(|| UserConnections {
|
|
meta,
|
|
public_key: public_key,
|
|
counter: AtomicU16::new(0),
|
|
connections: HashMap::new(),
|
|
});
|
|
|
|
let conid = clients
|
|
.counter
|
|
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
|
|
|
clients.connections.insert(conid, client.clone());
|
|
|
|
if let Err(e) = read_loop(&client).await {
|
|
eprintln!("Failed to handle client: {e}");
|
|
} else {
|
|
println!("Client connection closed")
|
|
}
|
|
|
|
clients.connections.remove(&conid);
|
|
|
|
if clients.connections.len() == 0 {
|
|
clients_meta.remove(&public_key);
|
|
}
|
|
}
|
|
|
|
Err(e) => {
|
|
eprintln!("Failed to initialize client: {e}")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
impl UserConnections {
|
|
pub async fn send(&self, message: &ClientMethod) -> anyhow::Result<()> {
|
|
for (_, conn) in &self.connections {
|
|
send_socket(&mut *conn.lock().await, message).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn send_to(&self, id: u16, message: &ClientMethod) -> anyhow::Result<bool> {
|
|
if let Some(conn) = self.connections.get(&id) {
|
|
send_socket(&mut *conn.lock().await, message).await?;
|
|
|
|
Ok(true)
|
|
} else {
|
|
Ok(false)
|
|
}
|
|
}
|
|
}
|