use super::{ client_state::{ClientState, CLIENT_STATES}, destroy_queue, scenegraph::Scenegraph, }; use crate::{ core::{registry::OwnedRegistry, task}, nodes::{audio, data, drawable, fields, hmd, input, items, root::Root, spatial, Node}, }; use color_eyre::eyre::{eyre, Result}; use lazy_static::lazy_static; use once_cell::sync::OnceCell; use parking_lot::Mutex; use rustc_hash::FxHashMap; use stardust_xr::{ messenger::{self, MessageSenderHandle}, schemas::flex::serialize, }; use std::{fmt::Debug, fs, iter::FromIterator, path::PathBuf, sync::Arc}; use tokio::{net::UnixStream, task::JoinHandle}; use tracing::info; lazy_static! { pub static ref CLIENTS: OwnedRegistry = OwnedRegistry::new(); pub static ref INTERNAL_CLIENT: Arc = CLIENTS.add(Client { pid: None, // env: None, exe: None, dispatch_join_handle: OnceCell::new(), flush_join_handle: OnceCell::new(), disconnect_status: OnceCell::new(), message_sender_handle: None, scenegraph: Default::default(), root: OnceCell::new(), base_resource_prefixes: Default::default(), state: Arc::new(ClientState::default()), }); } pub fn get_env(pid: i32) -> Result, std::io::Error> { let env = fs::read_to_string(format!("/proc/{pid}/environ"))?; Ok(FxHashMap::from_iter( env.split('\0') .filter_map(|var| var.split_once('=')) .map(|(k, v)| (k.to_string(), v.to_string())), )) } pub fn state(env: &FxHashMap) -> Option> { let token = env.get("STARDUST_STARTUP_TOKEN")?; CLIENT_STATES.lock().get(token).cloned() } pub struct Client { pub pid: Option, // env: Option>, exe: Option, dispatch_join_handle: OnceCell>>, flush_join_handle: OnceCell>>, disconnect_status: OnceCell>, pub message_sender_handle: Option, pub scenegraph: Arc, pub root: OnceCell>, pub base_resource_prefixes: Mutex>, pub state: Arc, } impl Client { pub fn from_connection(connection: UnixStream) -> Result> { let pid = connection.peer_cred().ok().and_then(|c| c.pid()); let env = pid.and_then(|pid| get_env(pid).ok()); let exe = pid.and_then(|pid| fs::read_link(format!("/proc/{}/exe", pid)).ok()); info!( pid, exe = exe .as_ref() .and_then(|exe| exe.to_str().map(|s| s.to_string())), "New client connected" ); let (mut messenger_tx, mut messenger_rx) = messenger::create(connection); let scenegraph = Arc::new(Scenegraph::default()); let state = env .as_ref() .and_then(state) .unwrap_or_else(|| Arc::new(ClientState::default())); let client = CLIENTS.add(Client { pid, // env, exe: exe.clone(), dispatch_join_handle: OnceCell::new(), flush_join_handle: OnceCell::new(), disconnect_status: OnceCell::new(), message_sender_handle: Some(messenger_tx.handle()), scenegraph: scenegraph.clone(), root: OnceCell::new(), base_resource_prefixes: Default::default(), state, }); let _ = client.scenegraph.client.set(Arc::downgrade(&client)); let _ = client.root.set(Root::create(&client)?); hmd::make_alias(&client)?; spatial::create_interface(&client)?; fields::create_interface(&client)?; drawable::create_interface(&client)?; audio::create_interface(&client)?; data::create_interface(&client)?; items::create_interface(&client)?; input::create_interface(&client)?; client .root .get() .unwrap() .node .send_remote_signal("restore_state", serialize(client.state.apply_to(&client))?)?; let pid_printable = pid .map(|pid| pid.to_string()) .unwrap_or_else(|| "??".to_string()); let exe_printable = exe .and_then(|exe| { exe.file_name() .and_then(|exe| exe.to_str()) .map(|exe| exe.to_string()) }) .unwrap_or_else(|| "??".to_string()); let _ = client.dispatch_join_handle.get_or_try_init(|| { task::new( || { format!( "client dispatch pid={} exe={}", &pid_printable, &exe_printable, ) }, { let client = client.clone(); async move { loop { match messenger_rx.dispatch(&*scenegraph).await { Err(e) => { client.disconnect(Err(e.into())); } _ => (), } } } }, ) }); let _ = client.flush_join_handle.get_or_try_init(|| { task::new( || format!("client flush pid={} exe={}", &pid_printable, &exe_printable,), { let client = client.clone(); async move { loop { match messenger_tx.flush().await { Err(e) => { client.disconnect(Err(e.into())); } _ => (), } } } }, ) }); Ok(client) } pub fn get_cmdline(&self) -> Option> { let pid = self.pid?; let exe_proc_path = format!("/proc/{pid}/exe"); let cmdline_proc_path = format!("/proc/{pid}/cmdline"); let exe = std::fs::read_link(exe_proc_path).ok()?; let cmdline = std::fs::read_to_string(cmdline_proc_path).ok()?; let mut cmdline_split: Vec<_> = cmdline.split('\0').map(ToString::to_string).collect(); cmdline_split.pop(); *cmdline_split.get_mut(0).unwrap() = exe.to_str()?.to_string(); Some(cmdline_split) } pub fn get_cwd(&self) -> Option { let pid = self.pid?; let cwd_proc_path = format!("/proc/{pid}/cwd"); std::fs::read_link(cwd_proc_path).ok() } pub async fn save_state(&self) -> Option { let internal = self.root.get()?.save_state().await.ok()?; Some(ClientState::from_deserialized(self, internal)) } #[inline] pub fn get_node(&self, name: &'static str, path: &str) -> Result> { self.scenegraph .get_node(path) .ok_or_else(|| eyre!("{} not found", name)) } pub fn disconnect(&self, reason: Result<()>) { let _ = self.disconnect_status.set(reason); if let Some(dispatch_join_handle) = self.dispatch_join_handle.get() { dispatch_join_handle.abort(); } if let Some(flush_join_handle) = self.flush_join_handle.get() { flush_join_handle.abort(); } if let Some(client) = CLIENTS.remove(self) { destroy_queue::add(client); } } } impl Debug for Client { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Client") .field("pid", &self.pid) .field("exe", &self.exe) .field("dispatch_join_handle", &self.dispatch_join_handle) .field("flush_join_handle", &self.flush_join_handle) .field("disconnect_status", &self.disconnect_status) .field("base_resource_prefixes", &self.base_resource_prefixes) .field("state", &self.state) .finish() } } impl Drop for Client { fn drop(&mut self) { info!( pid = self.pid, exe = self .exe .as_ref() .and_then(|exe| exe.to_str().map(|s| s.to_string())), disconnect_status = match self.disconnect_status.take() { Some(Ok(_)) => "Graceful disconnect".to_string(), Some(Err(e)) => format!("Error: {}", e.root_cause()), None => "Unknown".to_string(), }, "Client disconnected" ); } }