mitoxide_ssh/
bootstrap.rs

1//! Agent bootstrap and platform detection
2
3use crate::{Transport, TransportError};
4use tracing::{debug, info};
5
6/// Platform information detected from remote host
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct PlatformInfo {
9    /// Architecture (e.g., "x86_64", "aarch64")
10    pub arch: String,
11    /// Operating system (e.g., "Linux", "Darwin")
12    pub os: String,
13    /// OS version/distribution info
14    pub version: Option<String>,
15    /// Available bootstrap methods
16    pub bootstrap_methods: Vec<BootstrapMethod>,
17}
18
19/// Available bootstrap methods
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum BootstrapMethod {
22    /// Use memfd_create syscall (Linux only)
23    MemfdCreate,
24    /// Use temporary file in /tmp
25    TempFile,
26    /// Use temporary file in /dev/shm
27    DevShm,
28    /// Use Python for in-memory execution
29    Python,
30    /// Use shell script execution
31    Shell,
32}
33
34/// Bootstrap functionality for SSH transport
35pub struct Bootstrap {
36    /// Detected platform information
37    platform_info: Option<PlatformInfo>,
38    /// Custom bootstrap script template
39    custom_script: Option<String>,
40}
41
42impl Bootstrap {
43    /// Create a new bootstrap instance
44    pub fn new() -> Self {
45        Self {
46            platform_info: None,
47            custom_script: None,
48        }
49    }
50    
51    /// Set a custom bootstrap script template
52    pub fn with_custom_script(mut self, script: String) -> Self {
53        self.custom_script = Some(script);
54        self
55    }
56    
57    /// Detect platform information from the remote host
58    pub async fn detect_platform<T: Transport>(&mut self, transport: &mut T) -> Result<&PlatformInfo, TransportError> {
59        info!("Detecting remote platform");
60        
61        // Get basic platform info
62        let platform_cmd = "uname -m && uname -s && (lsb_release -d 2>/dev/null || cat /etc/os-release 2>/dev/null | head -1 || echo 'Unknown')";
63        let platform_output = self.execute_command(transport, platform_cmd).await?;
64        
65        let lines: Vec<&str> = platform_output.trim().split('\n').collect();
66        if lines.len() < 2 {
67            return Err(TransportError::Bootstrap("Failed to detect platform".to_string()));
68        }
69        
70        let arch = lines[0].trim().to_string();
71        let os = lines[1].trim().to_string();
72        let version = if lines.len() > 2 {
73            Some(lines[2].trim().to_string())
74        } else {
75            None
76        };
77        
78        debug!("Detected platform: {} {} {:?}", arch, os, version);
79        
80        // Detect available bootstrap methods
81        let bootstrap_methods = self.detect_bootstrap_methods(transport, &os).await?;
82        
83        let platform_info = PlatformInfo {
84            arch,
85            os,
86            version,
87            bootstrap_methods,
88        };
89        
90        self.platform_info = Some(platform_info);
91        Ok(self.platform_info.as_ref().unwrap())
92    }
93    
94    /// Detect available bootstrap methods
95    async fn detect_bootstrap_methods<T: Transport>(
96        &self, 
97        transport: &mut T, 
98        os: &str
99    ) -> Result<Vec<BootstrapMethod>, TransportError> {
100        let mut methods = Vec::new();
101        
102        // Check for memfd_create (Linux only)
103        if os == "Linux" {
104            let memfd_check = "python3 -c 'import ctypes; libc = ctypes.CDLL(\"libc.so.6\"); print(libc.syscall(319, b\"test\", 1) >= 0)' 2>/dev/null || echo 'False'";
105            if let Ok(output) = self.execute_command(transport, memfd_check).await {
106                if output.trim() == "True" {
107                    methods.push(BootstrapMethod::MemfdCreate);
108                    debug!("memfd_create available");
109                }
110            }
111        }
112        
113        // Check for Python
114        let python_check = "python3 --version 2>/dev/null || python --version 2>/dev/null";
115        if self.execute_command(transport, python_check).await.is_ok() {
116            methods.push(BootstrapMethod::Python);
117            debug!("Python available");
118        }
119        
120        // Check for /dev/shm
121        let devshm_check = "[ -d /dev/shm ] && [ -w /dev/shm ] && echo 'available'";
122        if let Ok(output) = self.execute_command(transport, devshm_check).await {
123            if output.trim() == "available" {
124                methods.push(BootstrapMethod::DevShm);
125                debug!("/dev/shm available");
126            }
127        }
128        
129        // Check for /tmp
130        let tmp_check = "[ -d /tmp ] && [ -w /tmp ] && echo 'available'";
131        if let Ok(output) = self.execute_command(transport, tmp_check).await {
132            if output.trim() == "available" {
133                methods.push(BootstrapMethod::TempFile);
134                debug!("/tmp available");
135            }
136        }
137        
138        // Shell is always available as fallback
139        methods.push(BootstrapMethod::Shell);
140        
141        Ok(methods)
142    }
143    
144    /// Generate bootstrap script for the detected platform
145    pub fn generate_bootstrap_script(&self, _agent_binary: &[u8]) -> Result<String, TransportError> {
146        let platform_info = self.platform_info.as_ref()
147            .ok_or_else(|| TransportError::Bootstrap("Platform not detected".to_string()))?;
148        
149        if let Some(custom_script) = &self.custom_script {
150            return Ok(custom_script.clone());
151        }
152        
153        // Choose the best available bootstrap method
154        let method = platform_info.bootstrap_methods.first()
155            .ok_or_else(|| TransportError::Bootstrap("No bootstrap methods available".to_string()))?;
156        
157        let script = match method {
158            BootstrapMethod::MemfdCreate => self.generate_memfd_script(),
159            BootstrapMethod::Python => self.generate_python_script(),
160            BootstrapMethod::DevShm => self.generate_devshm_script(),
161            BootstrapMethod::TempFile => self.generate_tempfile_script(),
162            BootstrapMethod::Shell => self.generate_shell_script(),
163        };
164        
165        debug!("Generated bootstrap script using method: {:?}", method);
166        Ok(script)
167    }
168    
169    /// Generate memfd_create bootstrap script
170    fn generate_memfd_script(&self) -> String {
171        r#"
172set -e
173python3 -c "
174import os, sys, ctypes
175try:
176    libc = ctypes.CDLL('libc.so.6')
177    fd = libc.syscall(319, b'mitoxide-agent', 1)  # memfd_create
178    if fd >= 0:
179        agent_data = sys.stdin.buffer.read()
180        os.write(fd, agent_data)
181        os.fexecve(fd, ['/proc/self/fd/%d' % fd], os.environ)
182    else:
183        raise Exception('memfd_create failed')
184except Exception as e:
185    print(f'memfd_create failed: {e}', file=sys.stderr)
186    sys.exit(1)
187"
188        "#.trim().to_string()
189    }
190    
191    /// Generate Python bootstrap script
192    fn generate_python_script(&self) -> String {
193        r#"
194set -e
195python3 -c "
196import os, sys, tempfile, stat
197try:
198    with tempfile.NamedTemporaryFile(delete=False, mode='wb') as f:
199        agent_data = sys.stdin.buffer.read()
200        f.write(agent_data)
201        f.flush()
202        os.chmod(f.name, stat.S_IRWXU)
203        os.execv(f.name, [f.name])
204except Exception as e:
205    print(f'Python bootstrap failed: {e}', file=sys.stderr)
206    sys.exit(1)
207"
208        "#.trim().to_string()
209    }
210    
211    /// Generate /dev/shm bootstrap script
212    fn generate_devshm_script(&self) -> String {
213        r#"
214set -e
215AGENT_PATH="/dev/shm/mitoxide-agent-$$-$(date +%s)"
216cat > "$AGENT_PATH"
217chmod +x "$AGENT_PATH"
218exec "$AGENT_PATH"
219        "#.trim().to_string()
220    }
221    
222    /// Generate /tmp bootstrap script
223    fn generate_tempfile_script(&self) -> String {
224        r#"
225set -e
226AGENT_PATH="/tmp/mitoxide-agent-$$-$(date +%s)"
227cat > "$AGENT_PATH"
228chmod +x "$AGENT_PATH"
229trap 'rm -f "$AGENT_PATH" 2>/dev/null || true' EXIT
230exec "$AGENT_PATH"
231        "#.trim().to_string()
232    }
233    
234    /// Generate shell bootstrap script (fallback)
235    fn generate_shell_script(&self) -> String {
236        r#"
237set -e
238# Try to find a writable directory
239for dir in /dev/shm /tmp /var/tmp; do
240    if [ -d "$dir" ] && [ -w "$dir" ]; then
241        AGENT_PATH="$dir/mitoxide-agent-$$-$(date +%s)"
242        cat > "$AGENT_PATH"
243        chmod +x "$AGENT_PATH"
244        trap 'rm -f "$AGENT_PATH" 2>/dev/null || true' EXIT
245        exec "$AGENT_PATH"
246        break
247    fi
248done
249echo "No writable directory found for agent bootstrap" >&2
250exit 1
251        "#.trim().to_string()
252    }
253    
254    /// Execute bootstrap on the remote host
255    pub async fn execute_bootstrap<T: Transport>(
256        &self, 
257        transport: &mut T, 
258        agent_binary: &[u8]
259    ) -> Result<(), TransportError> {
260        let script = self.generate_bootstrap_script(agent_binary)?;
261        
262        info!("Executing agent bootstrap");
263        debug!("Bootstrap script: {}", script);
264        
265        // This would be implemented by the specific transport
266        transport.bootstrap_agent(agent_binary).await
267    }
268    
269    /// Get platform information
270    pub fn platform_info(&self) -> Option<&PlatformInfo> {
271        self.platform_info.as_ref()
272    }
273    
274    /// Helper method to execute commands (would be implemented by transport)
275    async fn execute_command<T: Transport>(&self, _transport: &mut T, command: &str) -> Result<String, TransportError> {
276        // This is a placeholder - in reality, we'd need a way to execute commands
277        // through the transport without bootstrapping the agent
278        debug!("Would execute command: {}", command);
279        
280        // For now, return mock responses based on command
281        if command.contains("uname -m") {
282            Ok("x86_64\nLinux\nUbuntu 20.04.3 LTS".to_string())
283        } else if command.contains("python3 --version") {
284            Ok("Python 3.8.10".to_string())
285        } else if command.contains("memfd_create") {
286            Ok("True".to_string())
287        } else if command.contains("/dev/shm") || command.contains("/tmp") {
288            Ok("available".to_string())
289        } else {
290            Err(TransportError::CommandFailed { 
291                code: 1, 
292                message: "Command not found".to_string() 
293            })
294        }
295    }
296}
297
298impl Default for Bootstrap {
299    fn default() -> Self {
300        Self::new()
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::{Transport, TransportError, ConnectionInfo, TransportType};
308    use async_trait::async_trait;
309    
310    // Mock transport for testing
311    struct MockTransport {
312        should_fail: bool,
313    }
314    
315    impl MockTransport {
316        fn new(should_fail: bool) -> Self {
317            Self { should_fail }
318        }
319    }
320    
321    #[async_trait]
322    impl Transport for MockTransport {
323        async fn connect(&mut self) -> Result<crate::Connection, TransportError> {
324            if self.should_fail {
325                Err(TransportError::Connection("Mock connection failed".to_string()))
326            } else {
327                Ok(crate::Connection::new(None))
328            }
329        }
330        
331        async fn bootstrap_agent(&mut self, _agent_binary: &[u8]) -> Result<(), TransportError> {
332            if self.should_fail {
333                Err(TransportError::Bootstrap("Mock bootstrap failed".to_string()))
334            } else {
335                Ok(())
336            }
337        }
338        
339        fn connection_info(&self) -> ConnectionInfo {
340            ConnectionInfo {
341                host: "mock.example.com".to_string(),
342                port: 22,
343                username: "mockuser".to_string(),
344                transport_type: TransportType::Local,
345            }
346        }
347        
348        async fn test_connection(&mut self) -> Result<(), TransportError> {
349            if self.should_fail {
350                Err(TransportError::Connection("Mock test failed".to_string()))
351            } else {
352                Ok(())
353            }
354        }
355    }
356    
357    #[tokio::test]
358    async fn test_bootstrap_creation() {
359        let bootstrap = Bootstrap::new();
360        assert!(bootstrap.platform_info.is_none());
361        assert!(bootstrap.custom_script.is_none());
362    }
363    
364    #[tokio::test]
365    async fn test_custom_script() {
366        let custom_script = "echo 'custom bootstrap'".to_string();
367        let bootstrap = Bootstrap::new().with_custom_script(custom_script.clone());
368        assert_eq!(bootstrap.custom_script, Some(custom_script));
369    }
370    
371    #[tokio::test]
372    async fn test_platform_detection() {
373        let mut transport = MockTransport::new(false);
374        let mut bootstrap = Bootstrap::new();
375        
376        let platform_info = bootstrap.detect_platform(&mut transport).await.unwrap();
377        
378        assert_eq!(platform_info.arch, "x86_64");
379        assert_eq!(platform_info.os, "Linux");
380        assert!(!platform_info.bootstrap_methods.is_empty());
381    }
382    
383    #[test]
384    fn test_bootstrap_method_detection() {
385        let methods = vec![
386            BootstrapMethod::MemfdCreate,
387            BootstrapMethod::Python,
388            BootstrapMethod::DevShm,
389            BootstrapMethod::TempFile,
390            BootstrapMethod::Shell,
391        ];
392        
393        // Test that all methods are distinct
394        for (i, method1) in methods.iter().enumerate() {
395            for (j, method2) in methods.iter().enumerate() {
396                if i != j {
397                    assert_ne!(method1, method2);
398                }
399            }
400        }
401    }
402    
403    #[tokio::test]
404    async fn test_script_generation() {
405        let mut transport = MockTransport::new(false);
406        let mut bootstrap = Bootstrap::new();
407        
408        // Detect platform first
409        bootstrap.detect_platform(&mut transport).await.unwrap();
410        
411        // Generate bootstrap script
412        let agent_binary = b"fake agent binary";
413        let script = bootstrap.generate_bootstrap_script(agent_binary).unwrap();
414        
415        assert!(!script.is_empty());
416        assert!(script.contains("set -e")); // Should have error handling
417    }
418    
419    #[test]
420    fn test_memfd_script_generation() {
421        let bootstrap = Bootstrap::new();
422        let script = bootstrap.generate_memfd_script();
423        
424        assert!(script.contains("memfd_create"));
425        assert!(script.contains("python3"));
426        assert!(script.contains("syscall(319"));
427    }
428    
429    #[test]
430    fn test_python_script_generation() {
431        let bootstrap = Bootstrap::new();
432        let script = bootstrap.generate_python_script();
433        
434        assert!(script.contains("tempfile"));
435        assert!(script.contains("python3"));
436        assert!(script.contains("os.execv"));
437    }
438    
439    #[test]
440    fn test_tempfile_script_generation() {
441        let bootstrap = Bootstrap::new();
442        let script = bootstrap.generate_tempfile_script();
443        
444        assert!(script.contains("/tmp"));
445        assert!(script.contains("chmod +x"));
446        assert!(script.contains("exec"));
447    }
448    
449    #[test]
450    fn test_shell_script_generation() {
451        let bootstrap = Bootstrap::new();
452        let script = bootstrap.generate_shell_script();
453        
454        assert!(script.contains("/dev/shm"));
455        assert!(script.contains("/tmp"));
456        assert!(script.contains("chmod +x"));
457        assert!(script.contains("exec"));
458    }
459}
mitoxide_ssh/bootstrap.rs

mitoxide_ssh/
bootstrap.rs
