mitoxide_ssh/
connection.rs

1//! SSH connection management
2
3use crate::TransportError;
4use tokio::process::Child;
5use tracing::{debug, warn};
6
7/// SSH connection wrapper
8#[derive(Debug)]
9pub struct Connection {
10    /// SSH process handle
11    ssh_process: Option<Child>,
12    /// Connection state
13    connected: bool,
14}
15
16impl Connection {
17    /// Create a new connection from an SSH process
18    pub fn new(ssh_process: Option<Child>) -> Self {
19        let connected = ssh_process.is_some();
20        Self {
21            ssh_process,
22            connected,
23        }
24    }
25    
26    /// Check if the connection is active
27    pub fn is_connected(&self) -> bool {
28        self.connected
29    }
30    
31    /// Get mutable reference to the SSH process
32    pub fn process_mut(&mut self) -> Option<&mut Child> {
33        self.ssh_process.as_mut()
34    }
35    
36    /// Take ownership of the SSH process
37    pub fn take_process(&mut self) -> Option<Child> {
38        self.connected = false;
39        self.ssh_process.take()
40    }
41    
42    /// Close the connection
43    pub async fn close(&mut self) -> Result<(), TransportError> {
44        if let Some(mut child) = self.ssh_process.take() {
45            debug!("Closing SSH connection");
46            
47            // Try to terminate gracefully
48            if let Err(e) = child.kill().await {
49                warn!("Failed to kill SSH process: {}", e);
50            }
51            
52            // Wait for the process to exit
53            match child.wait().await {
54                Ok(status) => {
55                    debug!("SSH process exited with status: {}", status);
56                }
57                Err(e) => {
58                    warn!("Error waiting for SSH process: {}", e);
59                }
60            }
61        }
62        
63        self.connected = false;
64        Ok(())
65    }
66    
67    /// Get stdin handle for writing to the remote process
68    pub fn stdin(&mut self) -> Option<&mut tokio::process::ChildStdin> {
69        self.ssh_process.as_mut()?.stdin.as_mut()
70    }
71    
72    /// Get stdout handle for reading from the remote process
73    pub fn stdout(&mut self) -> Option<&mut tokio::process::ChildStdout> {
74        self.ssh_process.as_mut()?.stdout.as_mut()
75    }
76    
77    /// Get stderr handle for reading errors from the remote process
78    pub fn stderr(&mut self) -> Option<&mut tokio::process::ChildStderr> {
79        self.ssh_process.as_mut()?.stderr.as_mut()
80    }
81}
82
83impl Drop for Connection {
84    fn drop(&mut self) {
85        if let Some(mut child) = self.ssh_process.take() {
86            // Try to kill the process if it's still running
87            let _ = child.start_kill();
88        }
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    
96    #[test]
97    fn test_connection_creation() {
98        let conn = Connection::new(None);
99        assert!(!conn.is_connected());
100    }
101    
102    #[tokio::test]
103    async fn test_connection_close() {
104        let mut conn = Connection::new(None);
105        let result = conn.close().await;
106        assert!(result.is_ok());
107        assert!(!conn.is_connected());
108    }
109}