From 5ecae9f41365d83d8498be4daa868fc36858257c Mon Sep 17 00:00:00 2001 From: Josiah Glosson Date: Mon, 27 Jan 2025 19:49:44 -0600 Subject: [PATCH] Implement Labrinth side of tunneling --- apps/labrinth/src/queue/socket.rs | 24 +- apps/labrinth/src/routes/internal/statuses.rs | 239 ++++++++++++++++-- .../rust-common/src/networking/message.rs | 16 +- 3 files changed, 246 insertions(+), 33 deletions(-) diff --git a/apps/labrinth/src/queue/socket.rs b/apps/labrinth/src/queue/socket.rs index 32dc62070..654003d7f 100644 --- a/apps/labrinth/src/queue/socket.rs +++ b/apps/labrinth/src/queue/socket.rs @@ -1,16 +1,20 @@ //! "Database" for Hydra + use crate::models::users::{UserId, UserStatus}; use actix_ws::Session; -use dashmap::DashMap; +use dashmap::{DashMap, DashSet}; +use uuid::Uuid; pub struct ActiveSockets { pub sockets: DashMap, + pub tunnel_sockets: DashMap, } impl Default for ActiveSockets { fn default() -> Self { Self { sockets: DashMap::new(), + tunnel_sockets: DashMap::new(), } } } @@ -18,6 +22,7 @@ impl Default for ActiveSockets { pub struct ActiveSocket { pub status: UserStatus, pub socket: Session, + pub owned_tunnel_sockets: DashSet, } impl ActiveSocket { @@ -25,6 +30,23 @@ impl ActiveSocket { Self { status, socket: session, + owned_tunnel_sockets: DashSet::new(), } } } + +pub struct TunnelSocket { + pub owner: UserId, + pub socket_type: TunnelSocketType, +} + +impl TunnelSocket { + pub fn new(owner: UserId, socket_type: TunnelSocketType) -> Self { + Self { owner, socket_type } + } +} + +pub enum TunnelSocketType { + Listening, + Connected { connected_to: Uuid }, +} diff --git a/apps/labrinth/src/routes/internal/statuses.rs b/apps/labrinth/src/routes/internal/statuses.rs index 8e5804d4a..4c38a8d2b 100644 --- a/apps/labrinth/src/routes/internal/statuses.rs +++ b/apps/labrinth/src/routes/internal/statuses.rs @@ -5,12 +5,15 @@ use crate::database::redis::RedisPool; use crate::models::pats::Scopes; use crate::models::users::User; use crate::queue::session::AuthQueue; -use crate::queue::socket::{ActiveSocket, ActiveSockets}; +use crate::queue::socket::{ + ActiveSocket, ActiveSockets, TunnelSocket, TunnelSocketType, +}; use crate::routes::ApiError; use actix_web::web::{Data, Payload}; use actix_web::{get, web, HttpRequest, HttpResponse}; use actix_ws::Message; use chrono::Utc; +use dashmap::mapref::entry::Entry; use either::Either; use futures_util::{StreamExt, TryStreamExt}; use rust_common::ids::UserId; @@ -20,6 +23,8 @@ use rust_common::networking::message::{ use rust_common::users::UserStatus; use serde::Deserialize; use sqlx::PgPool; +use uuid::Uuid; +use Entry::Vacant; pub fn config(cfg: &mut web::ServiceConfig) { cfg.service(ws_init); @@ -60,9 +65,8 @@ pub async fn ws_init( let user = User::from_full(db_user); - if let Some((_, ActiveSocket { socket, .. })) = db.sockets.remove(&user.id) - { - let _ = socket.close(None).await; + if let Some((_, socket)) = db.sockets.remove(&user.id) { + let _ = socket.socket.close(None).await; } let (res, mut session, msg_stream) = match actix_ws::handle(&req, body) { @@ -93,7 +97,7 @@ pub async fn ws_init( .into(), ) }) - .map(|x| x.value().status.clone()) + .map(|x| x.status.clone()) .collect::>() } else { Vec::new() @@ -107,8 +111,10 @@ pub async fn ws_init( )?) .await; + let db = db.clone(); db.sockets - .insert(user.id, ActiveSocket::new(status.clone(), session)); + .entry(user.id) + .insert(ActiveSocket::new(status.clone(), session)); broadcast_friends( user.id, @@ -122,6 +128,9 @@ pub async fn ws_init( let mut stream = msg_stream.into_stream(); actix_web::rt::spawn(async move { + let Some(active_socket) = db.sockets.get(&user.id) else { + return; + }; // receive messages from websocket while let Some(msg) = stream.next().await { let message = match msg { @@ -140,8 +149,7 @@ pub async fn ws_init( Ok(Message::Ping(msg)) => { if let Some(socket) = db.sockets.get(&user.id) { - let ActiveSocket { socket, .. } = socket.value(); - let _ = socket.clone().pong(&msg).await; + let _ = socket.socket.clone().pong(&msg).await; } continue; } @@ -187,9 +195,143 @@ pub async fn ws_init( } } - ClientToServerMessage::SocketOpen { .. } => todo!(), - ClientToServerMessage::SocketClose { .. } => todo!(), - ClientToServerMessage::SocketSend { .. } => todo!(), + ClientToServerMessage::SocketListen { socket } => { + let Vacant(entry) = db.tunnel_sockets.entry(socket) else { + continue; + }; + entry.insert(TunnelSocket::new( + user.id, + TunnelSocketType::Listening, + )); + active_socket.owned_tunnel_sockets.insert(socket); + let _ = broadcast_friends( + user.id, + ServerToClientMessage::FriendSocketOpened { + user: user.id, + socket, + }, + &pool, + &db, + None, + ) + .await; + } + ClientToServerMessage::SocketConnect { + from_socket, + to_socket, + } => { + let Vacant(entry) = db.tunnel_sockets.entry(from_socket) + else { + continue; + }; + let Some(other_tunnel) = db.tunnel_sockets.get(&to_socket) + else { + continue; + }; + let Some(other_user) = db.sockets.get(&other_tunnel.owner) + else { + continue; + }; + let new_socket_id = Uuid::new_v4(); + entry.insert(TunnelSocket::new( + user.id, + TunnelSocketType::Connected { + connected_to: new_socket_id, + }, + )); + active_socket.owned_tunnel_sockets.insert(from_socket); + db.tunnel_sockets.insert( + new_socket_id, + TunnelSocket::new( + *other_user.key(), + TunnelSocketType::Connected { + connected_to: from_socket, + }, + ), + ); + other_user.owned_tunnel_sockets.insert(new_socket_id); + let _ = send_message( + &other_user, + &ServerToClientMessage::SocketConnected { + to_socket, + new_socket: new_socket_id, + }, + ) + .await; + } + ClientToServerMessage::SocketClose { socket } => { + if active_socket + .owned_tunnel_sockets + .remove(&socket) + .is_none() + { + continue; + } + let Some((_, tunnel_socket)) = + db.tunnel_sockets.remove(&socket) + else { + continue; + }; + match tunnel_socket.socket_type { + TunnelSocketType::Listening => { + let _ = broadcast_friends( + user.id, + ServerToClientMessage::SocketClosed { socket }, + &pool, + &db, + None, + ) + .await; + } + TunnelSocketType::Connected { connected_to } => { + let Some((_, other)) = + db.tunnel_sockets.remove(&connected_to) + else { + continue; + }; + let Some(other_user) = db.sockets.get(&other.owner) + else { + continue; + }; + let _ = send_message( + &other_user, + &ServerToClientMessage::SocketClosed { socket }, + ) + .await; + } + } + } + ClientToServerMessage::SocketSend { socket, data } => { + let Some(tunnel_socket) = db.tunnel_sockets.get(&socket) + else { + continue; + }; + if tunnel_socket.owner != user.id { + continue; + } + let TunnelSocketType::Connected { connected_to } = + tunnel_socket.socket_type + else { + continue; + }; + let Some(other_tunnel) = + db.tunnel_sockets.get(&connected_to) + else { + continue; + }; + let Some(other_user) = db.sockets.get(&other_tunnel.owner) + else { + continue; + }; + let _ = send_message( + &other_user, + &ServerToClientMessage::SocketData { + socket: connected_to, + data, + }, + ) + .await; + } } } @@ -222,16 +364,7 @@ pub async fn broadcast_friends( if friend.accepted { if let Some(socket) = sockets.sockets.get(&friend_id.into()) { - let ActiveSocket { socket, .. } = socket.value(); - - // FIXME Probably shouldn't swallow sending errors - let _ = match message.serialize() { - Ok(Either::Left(text)) => socket.clone().text(text).await, - Ok(Either::Right(bytes)) => { - socket.clone().binary(bytes).await - } - Err(_) => Ok(()), // TODO: Maybe should log these? Though it is the backend - }; + let _ = send_message(socket.value(), &message).await; } } } @@ -239,23 +372,77 @@ pub async fn broadcast_friends( Ok(()) } +pub async fn send_message( + socket: &ActiveSocket, + message: &ServerToClientMessage, +) -> Result<(), crate::database::models::DatabaseError> { + let mut socket = socket.socket.clone(); + + // FIXME Probably shouldn't swallow sending errors + let _ = match message.serialize() { + Ok(Either::Left(text)) => socket.text(text).await, + Ok(Either::Right(bytes)) => socket.binary(bytes).await, + Err(_) => Ok(()), // TODO: Maybe should log these? Though it is the backend + }; + + Ok(()) +} + pub async fn close_socket( id: UserId, pool: &PgPool, - sockets: &ActiveSockets, + db: &ActiveSockets, ) -> Result<(), crate::database::models::DatabaseError> { - if let Some((_, ActiveSocket { socket, .. })) = sockets.sockets.remove(&id) - { - let _ = socket.close(None).await; + if let Some((_, socket)) = db.sockets.remove(&id) { + let _ = socket.socket.close(None).await; broadcast_friends( id, ServerToClientMessage::UserOffline { id }, pool, - sockets, + db, None, ) .await?; + + for owned_socket in socket.owned_tunnel_sockets { + let Some((_, tunnel_socket)) = + db.tunnel_sockets.remove(&owned_socket) + else { + continue; + }; + match tunnel_socket.socket_type { + TunnelSocketType::Listening => { + let _ = broadcast_friends( + id, + ServerToClientMessage::SocketClosed { + socket: owned_socket, + }, + pool, + db, + None, + ) + .await; + } + TunnelSocketType::Connected { connected_to } => { + let Some((_, other)) = + db.tunnel_sockets.remove(&connected_to) + else { + continue; + }; + let Some(other_user) = db.sockets.get(&other.owner) else { + continue; + }; + let _ = send_message( + &other_user, + &ServerToClientMessage::SocketClosed { + socket: connected_to, + }, + ) + .await; + } + } + } } Ok(()) diff --git a/packages/rust-common/src/networking/message.rs b/packages/rust-common/src/networking/message.rs index 557f16dae..7da6c5696 100644 --- a/packages/rust-common/src/networking/message.rs +++ b/packages/rust-common/src/networking/message.rs @@ -10,9 +10,13 @@ pub enum ClientToServerMessage { profile_name: Option, }, - SocketOpen { + SocketListen { socket: Uuid, }, + SocketConnect { + from_socket: Uuid, + to_socket: Uuid, + }, SocketClose { socket: Uuid, }, @@ -42,15 +46,15 @@ pub enum ServerToClientMessage { from: UserId, }, + FriendSocketOpened { + user: UserId, + socket: Uuid, + }, SocketConnected { to_socket: Uuid, new_socket: Uuid, }, - SocketDisconnected { - socket: Uuid, - }, - FriendSocketOpened { - user: UserId, + SocketClosed { socket: Uuid, }, SocketData {