From b7a02e30ffa77ce5e9ceceb603979bd2368d45c6 Mon Sep 17 00:00:00 2001 From: Josiah Glosson Date: Mon, 27 Jan 2025 20:38:48 -0600 Subject: [PATCH] Implement non-friend part of theseus tunneling --- apps/labrinth/src/routes/internal/statuses.rs | 10 +++- packages/app-lib/src/state/friends.rs | 52 +++++++++++++++---- packages/app-lib/src/state/mod.rs | 3 ++ packages/app-lib/src/state/tunnel.rs | 6 +++ .../rust-common/src/networking/message.rs | 6 ++- 5 files changed, 65 insertions(+), 12 deletions(-) create mode 100644 packages/app-lib/src/state/tunnel.rs diff --git a/apps/labrinth/src/routes/internal/statuses.rs b/apps/labrinth/src/routes/internal/statuses.rs index 4c38a8d2b..bf24649e7 100644 --- a/apps/labrinth/src/routes/internal/statuses.rs +++ b/apps/labrinth/src/routes/internal/statuses.rs @@ -206,7 +206,7 @@ pub async fn ws_init( active_socket.owned_tunnel_sockets.insert(socket); let _ = broadcast_friends( user.id, - ServerToClientMessage::FriendSocketOpened { + ServerToClientMessage::FriendSocketListening { user: user.id, socket, }, @@ -228,6 +228,12 @@ pub async fn ws_init( else { continue; }; + if !matches!( + other_tunnel.socket_type, + TunnelSocketType::Listening + ) { + continue; + } let Some(other_user) = db.sockets.get(&other_tunnel.owner) else { continue; @@ -276,7 +282,7 @@ pub async fn ws_init( TunnelSocketType::Listening => { let _ = broadcast_friends( user.id, - ServerToClientMessage::SocketClosed { socket }, + ServerToClientMessage::FriendSocketStoppedListening { socket }, &pool, &db, None, diff --git a/packages/app-lib/src/state/friends.rs b/packages/app-lib/src/state/friends.rs index 5f5055f09..332323fff 100644 --- a/packages/app-lib/src/state/friends.rs +++ b/packages/app-lib/src/state/friends.rs @@ -2,6 +2,7 @@ use crate::config::{MODRINTH_API_URL_V3, MODRINTH_SOCKET_URL}; use crate::data::ModrinthCredentials; use crate::event::emit::emit_friend; use crate::event::FriendPayload; +use crate::state::tunnel::TunnelSocket; use crate::state::{ProcessManager, Profile}; use crate::util::fetch::{fetch_advanced, fetch_json, FetchSemaphore}; use async_tungstenite::tokio::{connect_async, ConnectStream}; @@ -21,7 +22,10 @@ use rust_common::networking::message::{ use rust_common::users::{UserId, UserStatus}; use serde::{Deserialize, Serialize}; use std::sync::Arc; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpStream; use tokio::sync::RwLock; +use uuid::Uuid; type WriteSocket = Arc, Message>>>>; @@ -29,6 +33,7 @@ type WriteSocket = pub struct FriendsSocket { write: WriteSocket, user_statuses: Arc>, + tunnel_sockets: Arc>, } #[derive(Deserialize, Serialize)] @@ -50,6 +55,7 @@ impl FriendsSocket { Self { write: Arc::new(RwLock::new(None)), user_statuses: Arc::new(DashMap::new()), + tunnel_sockets: Arc::new(DashMap::new()), } } @@ -103,6 +109,7 @@ impl FriendsSocket { let write_handle = self.write.clone(); let statuses = self.user_statuses.clone(); + let sockets = self.tunnel_sockets.clone(); tokio::spawn(async move { let mut read_stream = read; @@ -163,10 +170,34 @@ impl FriendsSocket { } ServerToClientMessage::FriendRequestRejected { .. } => todo!(), - ServerToClientMessage::SocketConnected { .. } => todo!(), - ServerToClientMessage::SocketDisconnected { .. } => todo!(), - ServerToClientMessage::FriendSocketOpened { .. } => todo!(), - ServerToClientMessage::SocketData { .. } => todo!(), + ServerToClientMessage::FriendSocketListening { .. } => todo!(), + ServerToClientMessage::FriendSocketStoppedListening { .. } => todo!(), + + ServerToClientMessage::SocketConnected { to_socket, new_socket } => { + if let Some(connected_to) = sockets.get(&to_socket) { + if let TunnelSocket::Listening(connected_to) = connected_to.value() { + if let Ok(local_addr) = connected_to.local_addr() { + if let Ok(new_stream) = TcpStream::connect(local_addr).await { + sockets.insert(new_socket, TunnelSocket::Connected(new_stream)); + continue; + } + } + } + } + let _ = Self::send_message(&write_handle, ClientToServerMessage::SocketClose { socket: new_socket }).await; + }, + ServerToClientMessage::SocketClosed { socket } => { + if let Some((_, TunnelSocket::Connected(mut stream))) = sockets.remove(&socket) { + let _ = stream.shutdown().await; + } + }, + ServerToClientMessage::SocketData { socket, data } => { + if let Some(mut socket) = sockets.get_mut(&socket) { + if let TunnelSocket::Connected(ref mut stream) = socket.value_mut() { + let _ = stream.write_all(&data).await; + } + } + }, } } } @@ -254,8 +285,11 @@ impl FriendsSocket { &self, profile_name: Option, ) -> crate::Result<()> { - self.send_message(ClientToServerMessage::StatusUpdate { profile_name }) - .await + Self::send_message( + &self.write, + ClientToServerMessage::StatusUpdate { profile_name }, + ) + .await } #[tracing::instrument(skip_all)] @@ -324,9 +358,9 @@ impl FriendsSocket { Ok(()) } - #[tracing::instrument(skip(self))] + #[tracing::instrument(skip(write))] async fn send_message( - &self, + write: &WriteSocket, message: ClientToServerMessage, ) -> crate::Result<()> { let serialized = match message.serialize()? { @@ -334,7 +368,7 @@ impl FriendsSocket { Either::Right(bytes) => Message::binary(bytes), }; - let mut write_lock = self.write.write().await; + let mut write_lock = write.write().await; if let Some(ref mut write_half) = *write_lock { write_half.send(serialized).await?; } diff --git a/packages/app-lib/src/state/mod.rs b/packages/app-lib/src/state/mod.rs index 56bf7e85f..75f73f689 100644 --- a/packages/app-lib/src/state/mod.rs +++ b/packages/app-lib/src/state/mod.rs @@ -34,6 +34,9 @@ pub use self::cache::*; mod friends; pub use self::friends::*; +mod tunnel; +pub use self::tunnel::*; + pub mod db; pub mod fs_watcher; mod mr_auth; diff --git a/packages/app-lib/src/state/tunnel.rs b/packages/app-lib/src/state/tunnel.rs new file mode 100644 index 000000000..c76d04bc8 --- /dev/null +++ b/packages/app-lib/src/state/tunnel.rs @@ -0,0 +1,6 @@ +use tokio::net::{TcpListener, TcpStream}; + +pub enum TunnelSocket { + Listening(TcpListener), + Connected(TcpStream), +} diff --git a/packages/rust-common/src/networking/message.rs b/packages/rust-common/src/networking/message.rs index 7da6c5696..951948bb1 100644 --- a/packages/rust-common/src/networking/message.rs +++ b/packages/rust-common/src/networking/message.rs @@ -46,10 +46,14 @@ pub enum ServerToClientMessage { from: UserId, }, - FriendSocketOpened { + FriendSocketListening { user: UserId, socket: Uuid, }, + FriendSocketStoppedListening { + socket: Uuid, + }, + SocketConnected { to_socket: Uuid, new_socket: Uuid,