use crate::engine::Engine; use pulse_sdk::{map_postcard_err, prelude::*}; use std::{ collections::HashMap, sync::{Arc, Weak}, }; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::{ UnixListener, unix::{OwnedReadHalf, OwnedWriteHalf}, }, sync::Mutex, }; #[derive(Debug)] pub struct TerminalServer { clients: Mutex>, logs: Mutex>, engine: Weak, } impl TerminalServer { pub fn new(engine: Weak) -> Arc { Arc::new(Self { clients: Mutex::new(HashMap::new()), logs: Mutex::new(Vec::new()), engine, }) } pub async fn run(self: &Arc) -> tokio::io::Result<()> { let path = pulse_sdk::server_path(); if path.exists() { tokio::fs::remove_file(&path).await?; } let listener = UnixListener::bind(&path)?; println!("Terminal server listening on {:?}", path); loop { let (stream, _) = listener.accept().await?; let (reader, writer) = stream.into_split(); let id = rand::random(); self.clients.lock().await.insert(id, writer); let s = self.clone(); tokio::spawn(async move { if let Err(err) = s.handle_client(&id, reader).await { eprintln!("Terminal connection error: {err}"); } }); } } async fn initialize_client(self: &Arc, id: &usize) -> tokio::io::Result<()> { let engine = self.get_engine(); self.send_to( id, pulse_sdk::terminal::TerminalServerMessage::StrategyUpdated(Strategy { strategy: engine.strategy.strategy.manifest.lock().await.clone(), risk: engine.strategy.risk.manifest.lock().await.clone(), mode: Mode::Auto, state: ItemState::Running, cooldown: engine.config.lock().await.cooldown, }), ) .await?; self.send_to( id, pulse_sdk::terminal::TerminalServerMessage::SetLogs(self.logs.lock().await.clone()), ) .await?; Ok(()) } async fn handle_client( self: &Arc, id: &usize, mut reader: OwnedReadHalf, ) -> tokio::io::Result<()> { self.initialize_client(id).await?; loop { let mut len_buf = [0u8; size_of::()]; let size = reader.read_exact(&mut len_buf).await?; let len = usize::from_le_bytes(len_buf); if size == 0 || len == 0 { break; } let mut buffer = vec![0u8; len]; reader.read_exact(&mut buffer).await?; match map_postcard_err(postcard::from_bytes(&buffer))? { TerminalClientMessage::ExecuteCommand(command) => { let command = command.as_str(); let (command, args) = if let Some((command, args)) = command.split_once(" ") { (command, args.split(" ").collect()) } else { (command, Vec::new()) }; self.get_engine().execute_command(command, args).await?; } } } Ok(()) } pub async fn broadcast( self: &Arc, message: pulse_sdk::terminal::TerminalServerMessage, ) -> tokio::io::Result<()> { let msg = map_postcard_err(postcard::to_allocvec(&message))?; let mut clients = self.clients.lock().await; let mut remove_clients = Vec::new(); for (id, client) in clients.iter_mut() { if let Err(e) = Self::send_to_client(client, &msg).await { remove_clients.push(*id); println!("{e:?}"); } } for id in remove_clients { clients.remove(&id); } Ok(()) } pub async fn send_to( self: &Arc, id: &usize, message: pulse_sdk::terminal::TerminalServerMessage, ) -> tokio::io::Result<()> { Self::send_to_client( self.clients.lock().await.get_mut(id).ok_or_else(|| { tokio::io::Error::new( std::io::ErrorKind::Other, format!("Client({id}) does not exist"), ) })?, &map_postcard_err(postcard::to_allocvec(&message))?, ) .await } pub async fn send_to_client(client: &mut OwnedWriteHalf, msg: &[u8]) -> tokio::io::Result<()> { client.write_all(&msg.len().to_le_bytes()).await?; client.write_all(msg).await?; client.flush().await?; Ok(()) } pub async fn log( self: &Arc, kind: LogKind, name: &str, message: &str, ) -> tokio::io::Result<()> { let log = EventLog { kind, name: name.to_string(), message: message.to_string(), }; self.logs.lock().await.push(log.clone()); self.broadcast(pulse_sdk::terminal::TerminalServerMessage::AddLog(log)) .await } pub async fn log_raw(self: &Arc, log: EventLog) -> tokio::io::Result<()> { self.logs.lock().await.push(log.clone()); self.broadcast(pulse_sdk::terminal::TerminalServerMessage::AddLog(log)) .await } pub async fn info(self: &Arc, name: &str, message: &str) -> tokio::io::Result<()> { self.log(LogKind::Info, name, message).await } pub async fn warn(self: &Arc, name: &str, message: &str) -> tokio::io::Result<()> { self.log(LogKind::Warn, name, message).await } pub async fn error(self: &Arc, name: &str, message: &str) -> tokio::io::Result<()> { self.log(LogKind::Err, name, message).await } pub async fn debug(self: &Arc, name: &str, message: &str) -> tokio::io::Result<()> { self.log(LogKind::Debug, name, message).await } pub fn get_engine(&self) -> Arc { self.engine .upgrade() .expect("Failed to upgrade engine(Weak) to Arc") } }