Clients list

src/protocol/initialize.rs

@@ -11,7 +11,10 @@ use crate::server::Server;
 use super::*;
 
 impl super::Client {
-    pub async fn initialize(server: &Arc<Server>, mut socket: WebSocket) -> anyhow::Result<Self> {
+    pub async fn initialize(
+        server: &Arc<Server>,
+        mut socket: WebSocket,
+    ) -> anyhow::Result<(Self, ClientMeta)> {
         let Some(ServerMethod::Initialize {
             public_key: public_key_string,
             signature,
@@ -78,10 +81,12 @@ impl super::Client {
             ));
         };
 
-        Ok(Self {
-            socket,
-            meta,
+        Ok((
+            Self {
+                socket: Arc::new(Mutex::new(socket)),
                 public_key,
-        })
+            },
+            meta,
+        ))
     }
 }

src/protocol/mod.rs

@@ -1,14 +1,15 @@
-use std::borrow::Cow;
+use std::{borrow::Cow, sync::Arc};
 
 use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
 use ed25519_dalek::VerifyingKey;
 use serde::{Deserialize, Serialize};
+use tokio::sync::Mutex;
 
 pub mod initialize;
 
+#[derive(Clone)]
 pub struct Client {
-    pub socket: WebSocket,
-    pub meta: ClientMeta,
+    pub socket: Arc<Mutex<WebSocket>>,
     pub public_key: VerifyingKey,
 }
 
@@ -118,10 +119,10 @@ impl Client {
     }
 
     pub async fn read(&mut self) -> anyhow::Result<Option<ServerMethod>> {
-        Self::read_socket(&mut self.socket).await
+        Self::read_socket(&mut *self.socket.lock().await).await
     }
 
     pub async fn send(&mut self, message: ClientMethod) -> anyhow::Result<()> {
-        Self::send_socket(&mut self.socket, message).await
+        Self::send_socket(&mut *self.socket.lock().await, message).await
     }
 }

src/server.rs

@@ -1,16 +1,30 @@
-use std::sync::Arc;
+use std::{
+    collections::HashMap,
+    sync::{Arc, atomic::AtomicU16},
+};
 
 use axum::{
     extract::{WebSocketUpgrade, ws::WebSocket},
     response::Response,
 };
-use ed25519_dalek::SigningKey;
+use ed25519_dalek::{SigningKey, VerifyingKey};
+use tokio::sync::Mutex;
 
-use crate::{config::Config, protocol::Client};
+use crate::{
+    config::Config,
+    protocol::{Client, ClientMeta},
+};
+
+pub struct OnlineClientMeta {
+    pub meta: ClientMeta,
+    pub counter: AtomicU16,
+    pub connections: HashMap<u16, Client>,
+}
 
 pub struct Server {
     pub key: SigningKey,
     pub config: Config,
+    pub clients: Mutex<HashMap<VerifyingKey, OnlineClientMeta>>,
 }
 
 impl Server {
@@ -18,6 +32,7 @@ impl Server {
         Ok(Arc::new(Self {
             key: crate::signature::get().await?,
             config: Config::get().await?,
+            clients: Mutex::new(HashMap::new()),
         }))
     }
 }
@@ -28,7 +43,25 @@ impl Server {
 
         ws.on_upgrade(move |socket: WebSocket| async move {
             match Client::initialize(&s, socket).await {
-                Ok(mut client) => {
+                Ok((mut client, meta)) => {
+                    let mut clients_meta = s.clients.lock().await;
+
+                    let client_meta =
+                        clients_meta
+                            .entry(client.public_key)
+                            .or_insert_with(|| OnlineClientMeta {
+                                meta,
+                                counter: AtomicU16::new(0),
+                                connections: HashMap::new(),
+                            });
+
+                    client_meta.connections.insert(
+                        client_meta
+                            .counter
+                            .fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                        client.clone(),
+                    );
+
                     if let Err(e) = client.read_loop().await {
                         eprintln!("Failed to handle client: {e}");
                     } else {