feat(server): tracker ws em tempo real com registro e broadcast de peers

server/src/api/mod.rs

@@ -1,7 +1,9 @@
 pub mod announce;
 pub mod files;
 pub mod health;
+pub mod peers;
 pub mod search;
+pub mod ws;
 
 use axum::{routing::{get, post}, Router};
 use tower_http::cors::CorsLayer;
@@ -13,6 +15,8 @@ pub fn router(state: SharedState) -> Router {
         .route("/search", get(search::search))
         .route("/files/{file_id}", get(files::get_file))
         .route("/announce", post(announce::announce))
+        .route("/peers/{file_id}", get(peers::peers))
+        .route("/tracker/ws", get(ws::ws_handler))
         .with_state(state);
 
     Router::new()

server/src/api/peers.rs

@@ -0,0 +1,13 @@
+use axum::extract::{Path, State};
+use axum::Json;
+use crate::error::AppResult;
+use crate::state::SharedState;
+
+pub async fn peers(
+    State(state): State<SharedState>,
+    Path(file_id): Path<String>,
+) -> AppResult<Json<serde_json::Value>> {
+    let t = state.tracker.lock().unwrap();
+    let peers = t.peers_of(&file_id);
+    Ok(Json(serde_json::json!({ "file_id": file_id, "peers": peers })))
+}

server/src/api/ws.rs

@@ -0,0 +1,94 @@
+use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
+use axum::extract::State;
+use axum::response::Response;
+use futures_util::{SinkExt, StreamExt};
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::UnboundedReceiverStream;
+use tracing::{debug, warn};
+
+use crate::state::SharedState;
+
+#[derive(serde::Deserialize, Debug)]
+struct RegisterMsg {
+    peer_id: String,
+    file_id: String,
+}
+
+pub async fn ws_handler(
+    ws: WebSocketUpgrade,
+    State(state): State<SharedState>,
+) -> Response {
+    ws.on_upgrade(move |socket| handle_socket(socket, state))
+}
+
+async fn handle_socket(socket: WebSocket, state: SharedState) {
+    let (mut sender, mut receiver) = socket.split();
+    let (tx, rx) = mpsc::unbounded_channel::<Message>();
+    let mut recv_stream = UnboundedReceiverStream::new(rx);
+
+    let registration = wait_for_registration(&mut receiver).await;
+    let (file_id, peer_id) = match registration {
+        Some((file_id, peer_id)) => {
+            let mut t = state.tracker.lock().unwrap();
+            t.register(file_id.clone(), peer_id.clone(), tx.clone());
+            (file_id, peer_id)
+        }
+        None => return,
+    };
+
+    // notifica os demais pares do mesmo arquivo (exceto o próprio)
+    {
+        let t = state.tracker.lock().unwrap();
+        t.broadcast_except(
+            &file_id,
+            &peer_id,
+            &format!(
+                r#"{{"type":"peer_online","file_id":"{}","peer_id":"{}"}}"#,
+                file_id, peer_id
+            ),
+        );
+    }
+
+    loop {
+        tokio::select! {
+            msg = receiver.next() => {
+                match msg {
+                    Some(Ok(Message::Close(_))) | None => break,
+                    Some(Ok(Message::Ping(p))) => {
+                        if sender.send(Message::Pong(p)).await.is_err() { break; }
+                    }
+                    Some(Ok(_)) => {} // ignoramos frames não relevantes
+                    Some(Err(e)) => {
+                        warn!("ws error: {e}");
+                        break;
+                    }
+                }
+            }
+            msg = recv_stream.next() => {
+                if let Some(m) = msg {
+                    if sender.send(m).await.is_err() { break; }
+                }
+            }
+        }
+    }
+
+    let mut t = state.tracker.lock().unwrap();
+    t.unregister(&file_id, &peer_id);
+    debug!("peer {peer_id} removed from {file_id}");
+}
+
+/// Awaits the first text frame carrying `peer_id` + `file_id`.
+async fn wait_for_registration(receiver: &mut futures_util::stream::SplitStream<WebSocket>) -> Option<(String, String)> {
+    loop {
+        match receiver.next().await {
+            Some(Ok(Message::Text(text))) => {
+                if let Ok(reg) = serde_json::from_str::<RegisterMsg>(&text) {
+                    return Some((reg.file_id, reg.peer_id));
+                }
+            }
+            Some(Ok(Message::Close(_))) | None => return None,
+            Some(Err(_)) => return None,
+            _ => continue,
+        }
+    }
+}

server/src/lib.rs

@@ -6,6 +6,7 @@ pub mod search;
 pub mod state;
 pub mod store;
 pub mod validate;
+pub mod ws;
 
 use axum::Router;
 use crate::config::Config;

server/src/state.rs

@@ -1,12 +1,14 @@
 use std::sync::{Arc, Mutex};
 use rusqlite::Connection;
 use crate::config::Config;
+use crate::ws::tracker::PeerTracker;
 
 pub type DbConn = Mutex<Connection>;
 
 pub struct AppState {
     pub config: Config,
     pub db: DbConn,
+    pub tracker: Mutex<PeerTracker>,
 }
 
 impl AppState {
@@ -14,6 +16,7 @@ impl AppState {
         Arc::new(Self {
             config,
             db: Mutex::new(db),
+            tracker: Mutex::new(PeerTracker::default()),
         })
     }
 }

server/src/ws/mod.rs

@@ -0,0 +1 @@
+pub mod tracker;

server/src/ws/tracker.rs

@@ -0,0 +1,50 @@
+use std::collections::HashMap;
+use tokio::sync::mpsc;
+use axum::extract::ws::Message;
+
+pub type WsSender = mpsc::UnboundedSender<Message>;
+
+#[derive(Default)]
+pub struct PeerTracker {
+    pub by_file: HashMap<String, HashMap<String, WsSender>>,
+}
+
+impl PeerTracker {
+    pub fn register(&mut self, file_id: String, peer_id: String, tx: WsSender) {
+        self.by_file.entry(file_id).or_default().insert(peer_id, tx);
+    }
+
+    pub fn unregister(&mut self, file_id: &str, peer_id: &str) {
+        if let Some(peers) = self.by_file.get_mut(file_id) {
+            peers.remove(peer_id);
+            if peers.is_empty() {
+                self.by_file.remove(file_id);
+            }
+        }
+    }
+
+    pub fn peers_of(&self, file_id: &str) -> Vec<String> {
+        self.by_file
+            .get(file_id)
+            .map(|peers| peers.keys().cloned().collect())
+            .unwrap_or_default()
+    }
+
+    pub fn online_total(&self) -> usize {
+        self.by_file
+            .values()
+            .flat_map(|peers| peers.keys())
+            .collect::<std::collections::HashSet<_>>()
+            .len()
+    }
+
+    pub fn broadcast_except(&self, file_id: &str, except_peer: &str, payload: &str) {
+        if let Some(peers) = self.by_file.get(file_id) {
+            for (pid, tx) in peers {
+                if pid != except_peer {
+                    let _ = tx.send(Message::Text(payload.to_string().into()));
+                }
+            }
+        }
+    }
+}