//! TCP Client component for outbound connections. //! //! Provides a TreeElement for managing TCP client connections with //! configuration, status queries, reconnection support, and protocol stacking. use std::io::{Read, Write}; use std::net::TcpStream; use std::sync::{Arc, Mutex}; use std::time::Duration; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use crate::protocols::{ProtocolConfig, ProtocolStack}; use crate::tcp::config::{ConnectionStatus, TcpClientConfig}; use unshell::tree::component::Component; use unshell::tree::message::TreeMessage; use unshell::tree::symbols; use unshell::tree::{Branch, TreeElement}; /// TCP Client component with protocol stacking support. /// /// This component can: /// - Connect to remote TCP servers /// - Apply protocol stacks (base64, http, etc.) /// - Send/receive messages via RPC /// - Auto-reconnect on failure #[derive(Debug, Serialize, Deserialize)] pub struct TcpClient { /// Unique name for this client pub name: String, /// Connection configuration pub config: TcpClientConfig, /// Protocol stack configuration #[serde(default)] pub protocols: Vec, /// Current connection status #[serde(skip)] status: ConnectionStatus, /// Active TCP stream #[serde(skip)] stream: Option>>, /// Protocol stack (runtime) #[serde(skip)] protocol_stack: ProtocolStack, /// Internal tree structure #[serde(skip)] branch: Branch, } impl TcpClient { /// Create a new TCP client with default settings pub fn new(name: impl Into) -> Self { Self::with_config(name, TcpClientConfig::default()) } /// Create a new TCP client with custom configuration pub fn with_config(name: impl Into, config: TcpClientConfig) -> Self { let name = name.into(); let mut branch = Branch::new("TCPClient"); let state_branch = Branch::new("state"); branch.add_child("state", Box::new(state_branch)); Self { name: name.clone(), config, protocols: Vec::new(), status: ConnectionStatus::disconnected(), stream: None, protocol_stack: ProtocolStack::new(), branch, } } /// Set protocol stack configuration pub fn set_protocols(&mut self, protocols: Vec) -> Result<(), String> { self.protocols = protocols.clone(); self.protocol_stack = ProtocolStack::from_configs(&protocols).map_err(|e| e.to_string())?; Ok(()) } /// Connect to the configured address pub fn connect(&mut self) -> Result<(), String> { let addr = format!("{}:{}", self.config.address, self.config.port); let stream = TcpStream::connect_timeout( &addr .parse() .map_err(|e| format!("Invalid address: {}", e))?, Duration::from_millis(self.config.timeout_ms), ) .map_err(|e| format!("Connection failed: {}", e))?; stream .set_nonblocking(false) .map_err(|e| format!("Failed to set blocking: {}", e))?; let local = stream .local_addr() .map(|a| a.to_string()) .unwrap_or_default(); let remote = stream .peer_addr() .map(|a| a.to_string()) .unwrap_or_default(); self.status = ConnectionStatus::connected(remote, local); self.stream = Some(Arc::new(Mutex::new(stream))); Ok(()) } /// Disconnect from server pub fn disconnect(&mut self) -> Result<(), String> { self.stream = None; self.status = ConnectionStatus::disconnected(); Ok(()) } /// Check if connected pub fn is_connected(&self) -> bool { self.status.connected } /// Send raw data over the connection pub fn send_raw(&mut self, data: &[u8]) -> Result { let stream = self.stream.as_ref().ok_or("Not connected")?; let mut stream = stream.lock().map_err(|e| format!("Lock failed: {}", e))?; stream .write(data) .map_err(|e| format!("Write failed: {}", e)) } /// Receive raw data from the connection pub fn recv_raw(&mut self, buffer_size: usize) -> Result, String> { let stream = self.stream.as_ref().ok_or("Not connected")?; let mut stream = stream.lock().map_err(|e| format!("Lock failed: {}", e))?; let mut buffer = vec![0u8; buffer_size]; let n = stream .read(&mut buffer) .map_err(|e| format!("Read failed: {}", e))?; buffer.truncate(n); Ok(buffer) } /// Send a TreeMessage through the protocol stack pub fn send_message_raw(&mut self, message: &TreeMessage) -> Result<(), String> { let encoded = self .protocol_stack .encode_message(message) .map_err(|e| format!("Encoding failed: {}", e))?; self.send_raw(&encoded)?; Ok(()) } /// Receive and decode a TreeMessage pub fn recv_message(&mut self, buffer_size: usize) -> Result { let data = self.recv_raw(buffer_size)?; self.protocol_stack .decode_message(&data) .map_err(|e| format!("Decoding failed: {}", e)) } /// Send and wait for response (RPC pattern) pub fn rpc_call(&mut self, message: &TreeMessage) -> Result { let id = message.id.clone(); self.send_message_raw(message)?; // Simple blocking receive - in production would have timeout let response = self.recv_message(4096)?; // Verify it's a response to our message if let Some(response_to) = &response.response_to { if response_to != id.as_deref().unwrap_or("") { return Err("Response ID mismatch".to_string()); } } Ok(response) } /// Get current configuration pub fn config(&self) -> &TcpClientConfig { &self.config } /// Get mutable configuration pub fn config_mut(&mut self) -> &mut TcpClientConfig { &mut self.config } /// Get status as JSON pub fn get_status(&self) -> Value { json!({ "connected": self.status.connected, "remote_address": self.status.remote_address, "local_address": self.status.local_address, "bytes_sent": self.status.bytes_sent, "bytes_received": self.status.bytes_received, "config": self.config, "protocols": self.protocol_stack.to_configs(), }) } /// Handle RPC call from message fn handle_rpc(&mut self, payload: &Value) -> Value { let method = match payload.get("method").and_then(|m| m.as_str()) { Some(m) => m, None => return json!({"success": false, "error": "missing method"}), }; let params = payload.get("params").cloned().unwrap_or(Value::Null); match method { "connect" => { // Allow override of address/port if let Some(addr) = params.get("address").and_then(|a| a.as_str()) { self.config.address = addr.to_string(); } if let Some(port) = params.get("port").and_then(|p| p.as_u64()) { self.config.port = port as u16; } match self.connect() { Ok(_) => json!({"success": true, "status": self.status}), Err(e) => json!({"success": false, "error": e}), } } "disconnect" => match self.disconnect() { Ok(_) => json!({"success": true}), Err(e) => json!({"success": false, "error": e}), }, "send" => { let data = params .get("data") .and_then(|d| d.as_str()) .map(|s| s.as_bytes().to_vec()); match data { Some(data) => match self.send_raw(&data) { Ok(n) => json!({"success": true, "bytes_sent": n}), Err(e) => json!({"success": false, "error": e}), }, None => json!({"success": false, "error": "missing data"}), } } "recv" => { let size = params .get("size") .and_then(|s| s.as_u64()) .map(|s| s as usize) .unwrap_or(4096); match self.recv_raw(size) { Ok(data) => json!({ "success": true, "data": String::from_utf8_lossy(&data), "bytes": data.len() }), Err(e) => json!({"success": false, "error": e}), } } "status" => self.get_status(), "set_protocols" => { if let Some(protocols) = params.get("protocols") { match serde_json::from_value(protocols.clone()) { Ok(p) => match self.set_protocols(p) { Ok(_) => json!({"success": true}), Err(e) => json!({"success": false, "error": e}), }, Err(e) => json!({"success": false, "error": e.to_string()}), } } else { json!({"success": false, "error": "missing protocols"}) } } _ => json!({"success": false, "error": format!("unknown method: {}", method)}), } } } impl Component for TcpClient { fn name(&self) -> &str { &self.name } fn status(&self) -> Value { self.get_status() } fn init(&mut self, config: Value) -> Result<(), String> { // Support both legacy config and new format if let Some(client_config) = config.get("config") { self.config = serde_json::from_value(client_config.clone()) .map_err(|e| format!("Invalid config: {}", e))?; } else { self.config = serde_json::from_value(config.clone()) .map_err(|e| format!("Invalid config: {}", e))?; } if let Some(protocols) = config.get("protocols") { let p: Vec = serde_json::from_value(protocols.clone()) .map_err(|e| format!("Invalid protocols: {}", e))?; self.set_protocols(p)?; } self.connect()?; Ok(()) } fn shutdown(&mut self) -> Result<(), String> { self.disconnect() } } impl TreeElement for TcpClient { fn get_type(&self) -> Value { json!({ "type": "TCPClient", "name": self.name, }) } fn send_message(&mut self, target: Value, message: Value) -> Value { match target { Value::Null => { // Check for RPC call format if message.get("method").is_some() { return self.handle_rpc(&message); } // Legacy string commands if let Some(cmd) = message.as_str() { match cmd { "Connect" => match self.connect() { Ok(_) => json!({"success": true}), Err(e) => json!({"success": false, "error": e}), }, "Disconnect" => match self.disconnect() { Ok(_) => json!({"success": true}), Err(e) => json!({"success": false, "error": e}), }, "Status" => self.get_status(), symbols::CMD_GET_CHILDREN => { let children = self .branch .children() .keys() .map(|k| json!(k)) .collect::>(); json!(children) } _ => json!(symbols::ERR_UNSUPPORTED_METHOD), } } else if let Value::Object(obj) = message { // Handle configuration changes if let Some(config) = obj.get("config") { match serde_json::from_value(config.clone()) { Ok(cfg) => { self.config = cfg; json!({"success": true}) } Err(e) => json!({"success": false, "error": e.to_string()}), } } else if obj.get("method").is_some() { let payload = Value::Object(obj.clone()); self.handle_rpc(&payload) } else { json!(symbols::ERR_INVALID_COMMAND) } } else { json!(symbols::ERR_INVALID_COMMAND) } } Value::String(subtarget) => match subtarget.as_str() { "config" => json!(self.config), "state" => json!({ "connected": self.status.connected, "remote": self.status.remote_address, }), "protocols" => json!(self.protocol_stack.to_configs()), _ => json!(symbols::ERR_CHILD_NOT_FOUND), }, _ => json!(symbols::ERR_INVALID_TARGET), } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_client_creation() { let client = TcpClient::new("test-client"); assert_eq!(client.name(), "test-client"); assert!(!client.is_connected()); } #[test] fn test_config_serialization() { let client = TcpClient::with_config("test", TcpClientConfig::new("127.0.0.1", 8080)); let json = serde_json::to_string(&client).unwrap(); assert!(json.contains("test")); assert!(json.contains("127.0.0.1")); } #[test] fn test_rpc_status() { let mut client = TcpClient::new("test"); let result = client.send_message(json!(null), json!({"method": "status"})); let obj = result.as_object().unwrap(); assert!(obj.contains_key("connected")); } }