Skip to main content

ssh_mcp/
server.rs

1//! MCP Server implementation
2//!
3//! This module provides the main MCP server that integrates SSH connection
4//! management with the `exec` and `sudo-exec` tools.
5
6use std::path::{Path, PathBuf};
7use std::sync::Arc;
8use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
9use std::time::Duration;
10use std::time::{SystemTime, UNIX_EPOCH};
11
12use rmcp::{
13    ErrorData as McpError,
14    handler::server::ServerHandler,
15    model::*,
16    service::{RequestContext, RoleServer},
17};
18use tokio::sync::Mutex;
19use tracing::{debug, error, info, warn};
20
21use crate::background::detach::{DetachMode, DetachProbeOutput, DetachProbeRequest};
22use crate::background::job::NewRunningJob;
23use crate::background::wrapper::{
24    build_background_wrapper_script_full, build_background_wrapper_script_portable,
25};
26use crate::background::{JobRegistry, JobState, LocalLogSpooler, SharedJobState};
27use crate::config::Config;
28use crate::error::{Result, SshMcpError};
29#[cfg(unix)]
30use crate::platform::O_NOFOLLOW_FLAG;
31#[cfg(test)]
32use crate::server::validation::read_file::{
33    READ_FILE_BYTES_PER_TOKEN, READ_FILE_DEFAULT_PREVIEW_LINES, READ_FILE_HARD_MAX_BYTES,
34    READ_FILE_MAX_LINE_WINDOW,
35};
36#[cfg(test)]
37use crate::server::validation::read_file::{
38    estimate_tokens_from_bytes, resolve_read_file_line_limit, resolve_read_file_max_bytes,
39};
40#[cfg(test)]
41use crate::server::validation::validate_background_log_path;
42use crate::ssh::{
43    CommandOutput, SshConfig, SshConnectionManager, sanitize_command, wrap_sudo_command,
44};
45use crate::ticket::TicketSigner;
46use crate::tools::{
47    CheckProcessParams, ReadFileMode, ReadFileParams, ReplaceInFileParams, WriteFileParams,
48};
49use crate::transfer::{TransferEngine, TransferParams, TransferRunContext, TransferSshOptions};
50
51mod args;
52mod exec;
53mod handlers;
54mod testing;
55mod tools;
56mod validation;
57
58const BACKGROUND_START_TIMEOUT: Duration = Duration::from_secs(20);
59const READ_FILE_ERROR_MARKER: &str = "__SSH_MCP_READ_FILE_ERR__";
60
61const JOB_COMPLETED_RETENTION: Duration = Duration::from_secs(60 * 60);
62
63static JOB_COUNTER: AtomicU64 = AtomicU64::new(0);
64
65fn make_job_id() -> String {
66    let counter = JOB_COUNTER.fetch_add(1, Ordering::Relaxed);
67    let epoch_ms = SystemTime::now()
68        .duration_since(UNIX_EPOCH)
69        .map(|d| d.as_millis())
70        .unwrap_or(0);
71    format!("{}-{}", epoch_ms, counter)
72}
73
74fn build_background_wrapper_script(
75    mode: DetachMode,
76    job_id: &str,
77    user_command: &str,
78    log_path: &str,
79) -> String {
80    match mode {
81        DetachMode::Full | DetachMode::Unknown => {
82            build_background_wrapper_script_full(job_id, user_command, log_path)
83        }
84        DetachMode::Portable => {
85            build_background_wrapper_script_portable(job_id, user_command, log_path)
86        }
87        DetachMode::DirectOnly => {
88            build_background_wrapper_script_portable(job_id, user_command, log_path)
89        }
90    }
91}
92
93/// SSH MCP Server
94///
95/// The main server implementation that provides MCP tools for remote SSH
96/// command execution.
97#[derive(Clone)]
98pub struct SshMcpServer {
99    /// Server configuration
100    config: Config,
101
102    /// SSH connection manager
103    connection: Arc<SshConnectionManager>,
104
105    /// Command execution timeout
106    timeout: Duration,
107
108    /// Maximum command length
109    max_chars: Option<usize>,
110
111    detach_mode: Arc<AtomicU8>,
112    detach_mode_lock: Arc<Mutex<()>>,
113
114    spooler: Arc<LocalLogSpooler>,
115    job_registry: Arc<JobRegistry>,
116
117    transfer: TransferEngine,
118    ticket_signer: Arc<TicketSigner>,
119}
120
121impl SshMcpServer {
122    /// Create a new SSH MCP Server
123    ///
124    /// This sets up the SSH connection manager based on the provided configuration.
125    /// Connection is not established until a tool is actually used.
126    pub async fn new(config: Config) -> Result<Self> {
127        let local_root = std::env::current_dir()?;
128
129        let spooler = Arc::new(LocalLogSpooler::new_default());
130        spooler.ensure_dir().await.map_err(|e| {
131            SshMcpError::Config(format!(
132                "failed to initialize local log spool dir {}: {e}",
133                spooler.base_dir().display()
134            ))
135        })?;
136        let job_registry = Arc::new(JobRegistry::new(JOB_COMPLETED_RETENTION));
137
138        // Build SSH configuration
139        let mut ssh_config = SshConfig::new(&config.host, &config.user).with_port(config.port);
140
141        // Add authentication
142        if let Some(ref password) = config.password {
143            ssh_config = ssh_config.with_password(password);
144        }
145
146        if let Some(ref key_path) = config.key {
147            // Read the key file
148            let key_content = tokio::fs::read_to_string(key_path)
149                .await
150                .map_err(SshMcpError::Io)?;
151            ssh_config = ssh_config.with_private_key(&key_content);
152        }
153
154        // Add elevation passwords if provided
155        if let Some(ref su_password) = config.su_password {
156            ssh_config = ssh_config.with_su_password(su_password);
157        }
158
159        if let Some(ref sudo_password) = config.sudo_password {
160            ssh_config = ssh_config.with_sudo_password(sudo_password);
161        }
162
163        // Add keepalive settings for human-like connection persistence
164        ssh_config = ssh_config
165            .with_keepalive_interval(config.keepalive_interval)
166            .with_keepalive_max(config.keepalive_max);
167
168        // Add reconnect and health probe settings
169        ssh_config = ssh_config
170            .with_reconnect_retries(config.reconnect_retries)
171            .with_reconnect_backoff_ms(config.reconnect_backoff_ms)
172            .with_health_probe_timeout_ms(config.health_probe_timeout_ms);
173
174        // Add host key verification settings
175        ssh_config = ssh_config
176            .with_host_key_checking(config.strict_host_key_checking)
177            .with_known_hosts(config.known_hosts.clone());
178
179        // Add output token limit for OOM protection
180        ssh_config = ssh_config.with_max_output_tokens(config.max_output_tokens);
181
182        // Create connection manager
183        let connection = Arc::new(SshConnectionManager::new(ssh_config).await);
184
185        let timeout = Duration::from_millis(config.timeout_ms);
186        let max_chars = config.max_chars;
187
188        Ok(Self {
189            config,
190            connection,
191            timeout,
192            max_chars,
193            detach_mode: Arc::new(AtomicU8::new(DetachMode::Unknown.as_u8())),
194            detach_mode_lock: Arc::new(Mutex::new(())),
195            spooler,
196            job_registry,
197            transfer: TransferEngine::new(local_root),
198            ticket_signer: Arc::new(TicketSigner::new()),
199        })
200    }
201
202    fn connection_id(&self) -> String {
203        format!(
204            "{}@{}:{}",
205            self.config.user, self.config.host, self.config.port
206        )
207    }
208
209    fn default_local_log_path(
210        &self,
211        job_id: &str,
212    ) -> std::result::Result<(PathBuf, String), String> {
213        let path = self
214            .spooler
215            .log_path_for(job_id)
216            .map_err(|e| format!("failed to generate local log path for job_id='{job_id}': {e}"))?;
217        let path_str = path.to_string_lossy().to_string();
218        Ok((path, path_str))
219    }
220
221    async fn ensure_local_log_file(&self, log_path: &Path) -> std::result::Result<(), SshMcpError> {
222        self.spooler.ensure_dir().await.map_err(|e| {
223            SshMcpError::Config(format!(
224                "failed to ensure local log spool dir {}: {e}",
225                self.spooler.base_dir().display()
226            ))
227        })?;
228
229        if log_path.parent() != Some(self.spooler.base_dir()) {
230            return Err(SshMcpError::InvalidParams(format!(
231                "log_path must be directly under {}",
232                self.spooler.base_dir().display()
233            )));
234        }
235
236        match tokio::fs::symlink_metadata(log_path).await {
237            Ok(meta) => {
238                let ft = meta.file_type();
239                if ft.is_symlink() {
240                    return Err(SshMcpError::invalid_params(
241                        "log_path is a symlink (refusing to follow it)",
242                    ));
243                }
244                if !ft.is_file() {
245                    return Err(SshMcpError::invalid_params(
246                        "log_path exists but is not a regular file",
247                    ));
248                }
249            }
250            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
251            Err(e) => return Err(SshMcpError::Io(e)),
252        }
253
254        let mut opts = tokio::fs::OpenOptions::new();
255        opts.write(true).create(true).truncate(true);
256
257        #[cfg(unix)]
258        {
259            opts.custom_flags(O_NOFOLLOW_FLAG);
260        }
261
262        let file = match opts.open(log_path).await {
263            Ok(f) => f,
264            Err(e) => {
265                if let Ok(meta) = tokio::fs::symlink_metadata(log_path).await
266                    && meta.file_type().is_symlink()
267                {
268                    return Err(SshMcpError::invalid_params(
269                        "log_path is a symlink (refusing to follow it)",
270                    ));
271                }
272                return Err(SshMcpError::Io(e));
273            }
274        };
275
276        file.sync_all().await.map_err(SshMcpError::Io)
277    }
278
279    async fn register_running_job(
280        &self,
281        job_id: &str,
282        pid: u32,
283        log_path: PathBuf,
284        command: &str,
285    ) -> SharedJobState {
286        let job = Arc::new(Mutex::new(JobState::new_running(NewRunningJob {
287            job_id: job_id.to_string(),
288            pid,
289            log_path,
290            command: command.to_string(),
291            connection_id: self.connection_id(),
292        })));
293
294        self.job_registry
295            .insert(job_id.to_string(), Arc::clone(&job))
296            .await;
297
298        let persisted = {
299            let guard = job.lock().await;
300            guard.clone()
301        };
302        if let Err(e) = self.spooler.persist_job_state(&persisted).await {
303            warn!(job_id = ?job_id, error = ?e, "failed to persist running job state");
304        }
305
306        job
307    }
308
309    /// Get a reference to the SSH connection manager
310    pub fn connection(&self) -> &Arc<SshConnectionManager> {
311        &self.connection
312    }
313
314    /// Close the server and cleanup resources
315    pub async fn shutdown(&self) {
316        info!("Shutting down SSH MCP Server...");
317        self.connection.close().await;
318    }
319
320    async fn determine_detach_mode(&self) -> Result<DetachMode> {
321        let server = self.clone();
322        crate::background::detach::determine_detach_mode(
323            self.detach_mode.as_ref(),
324            self.detach_mode_lock.as_ref(),
325            make_job_id,
326            move |req, timeout| {
327                let server = server.clone();
328                async move { server.exec_detach_probe(req, timeout).await }
329            },
330        )
331        .await
332    }
333
334    async fn exec_detach_probe(
335        &self,
336        req: DetachProbeRequest,
337        timeout: Duration,
338    ) -> Result<DetachProbeOutput> {
339        let output = self.connection.exec_command(&req.wrapper, timeout).await?;
340        Ok(DetachProbeOutput {
341            stdout: output.stdout,
342            stderr: output.stderr,
343            exit_code: output.exit_code,
344        })
345    }
346
347    /// Execute a command (used by exec tool)
348    async fn execute_command_with_timeout(
349        &self,
350        command: &str,
351        timeout: Duration,
352    ) -> std::result::Result<CallToolResult, McpError> {
353        debug!(
354            "exec tool called: cmd_len={}, background=false, sudo=false, timeout_ms={}",
355            command.len(),
356            timeout.as_millis()
357        );
358
359        // Sanitize the command
360        let sanitized = match self.sanitize_or_tool_error(command) {
361            Ok(cmd) => cmd,
362            Err(result) => return Ok(result),
363        };
364
365        // Foreground execution is detachable-by-design:
366        // - Start the command on a dedicated SSH channel
367        // - Stream remote stdout/stderr into a local spool file
368        // - If timeout elapses, return JSON with job_id/pid/log_path while the stream continues
369
370        let requires_elevation = self.connection.get_su_password().is_some();
371        if requires_elevation {
372            if let Err(e) = self.connection.ensure_connected().await {
373                error!(error = ?e, "Failed to ensure SSH connection");
374                return Ok(CallToolResult::error(vec![Content::text(e.to_string())]));
375            }
376
377            if let Err(e) = self.connection.ensure_elevated().await {
378                debug!(error = ?e, "Elevation failed, will run as normal user");
379            }
380        }
381
382        let detach_mode = match self.determine_detach_mode().await {
383            Ok(mode) => mode,
384            Err(e) => {
385                debug!(error = ?e, "detach-mode probe failed; falling back to direct foreground exec");
386                DetachMode::DirectOnly
387            }
388        };
389        if detach_mode == DetachMode::DirectOnly {
390            match self.connection.exec_command(&sanitized, timeout).await {
391                Ok(output) => return Ok(Self::calltool_from_command_output(output)),
392                Err(e) => {
393                    error!(error = ?e, "Command execution failed");
394                    let mut msg = format!("Error: {}", e);
395                    if matches!(e, SshMcpError::Timeout(_)) {
396                        msg.push_str("\nHint: background detach is not supported on this target; rerun with background=true or a larger timeout_ms.");
397                    }
398                    return Ok(CallToolResult::error(vec![Content::text(msg)]));
399                }
400            }
401        }
402
403        // Ensure connection is established for detached foreground execution path.
404        if !requires_elevation && let Err(e) = self.connection.ensure_connected().await {
405            error!(error = ?e, "Failed to ensure SSH connection");
406            return Ok(CallToolResult::error(vec![Content::text(e.to_string())]));
407        }
408
409        self.execute_detachable_foreground_impl(detach_mode, &sanitized, &sanitized, timeout)
410            .await
411    }
412
413    async fn execute_command(
414        &self,
415        command: &str,
416    ) -> std::result::Result<CallToolResult, McpError> {
417        self.execute_command_with_timeout(command, self.timeout)
418            .await
419    }
420
421    async fn execute_background_command(
422        &self,
423        command: &str,
424        log_path: Option<&str>,
425    ) -> std::result::Result<CallToolResult, McpError> {
426        self.execute_background_impl(command, log_path, exec::BackgroundPrivilege::Normal)
427            .await
428    }
429
430    /// Execute a command with sudo (used by sudo-exec tool)
431    async fn execute_sudo_command_with_timeout(
432        &self,
433        command: &str,
434        timeout: Duration,
435    ) -> std::result::Result<CallToolResult, McpError> {
436        debug!(
437            "sudo-exec tool called: cmd_len={}, background=false, sudo=true, timeout_ms={}",
438            command.len(),
439            timeout.as_millis()
440        );
441
442        // Sanitize the command
443        let sanitized = match self.sanitize_or_tool_error(command) {
444            Ok(cmd) => cmd,
445            Err(result) => return Ok(result),
446        };
447
448        // Wrap the command with sudo
449        let sudo_password = self.connection.get_sudo_password();
450        let wrapped_command = wrap_sudo_command(&sanitized, sudo_password);
451        debug!(
452            "Wrapped sudo command (password hidden): sudo -n sh -c '...' or printf '...' | sudo ..."
453        );
454
455        if let Err(e) = self.connection.ensure_connected().await {
456            error!(error = ?e, "Failed to ensure SSH connection");
457            return Ok(CallToolResult::error(vec![Content::text(e.to_string())]));
458        }
459
460        let detach_mode = match self.determine_detach_mode().await {
461            Ok(mode) => mode,
462            Err(e) => {
463                debug!(error = ?e, "detach-mode probe failed; falling back to direct sudo foreground exec");
464                DetachMode::DirectOnly
465            }
466        };
467        if detach_mode == DetachMode::DirectOnly {
468            match self
469                .connection
470                .exec_command(&wrapped_command, timeout)
471                .await
472            {
473                Ok(output) => Ok(Self::calltool_from_command_output(output)),
474                Err(e) => {
475                    error!(error = ?e, "Sudo command execution failed");
476                    let mut msg = format!("Error: {}", e);
477                    if matches!(e, SshMcpError::Timeout(_)) {
478                        msg.push_str("\nHint: background detach is not supported on this target; rerun with background=true or a larger timeout_ms.");
479                    }
480                    Ok(CallToolResult::error(vec![Content::text(msg)]))
481                }
482            }
483        } else {
484            self.execute_detachable_foreground_impl(
485                detach_mode,
486                &wrapped_command,
487                &format!("sudo {sanitized}"),
488                timeout,
489            )
490            .await
491        }
492    }
493
494    async fn execute_sudo_command(
495        &self,
496        command: &str,
497    ) -> std::result::Result<CallToolResult, McpError> {
498        self.execute_sudo_command_with_timeout(command, self.timeout)
499            .await
500    }
501
502    async fn execute_background_sudo_command(
503        &self,
504        command: &str,
505        log_path: Option<&str>,
506    ) -> std::result::Result<CallToolResult, McpError> {
507        let sudo_password = self.connection.get_sudo_password();
508        self.execute_background_impl(
509            command,
510            log_path,
511            exec::BackgroundPrivilege::Sudo {
512                password: sudo_password,
513            },
514        )
515        .await
516    }
517
518    fn sanitize_or_tool_error(&self, command: &str) -> std::result::Result<String, CallToolResult> {
519        sanitize_command(command, self.max_chars).map_err(|e| {
520            error!(error = ?e, "Command sanitization failed");
521            CallToolResult::error(vec![Content::text(format!("Error: {}", e))])
522        })
523    }
524
525    fn calltool_from_command_output(output: CommandOutput) -> CallToolResult {
526        // Combine stdout and stderr for the response
527        let mut result_text = output.stdout;
528        if !output.stderr.is_empty() {
529            if !result_text.is_empty() {
530                result_text.push_str("\n--- stderr ---\n");
531            }
532            result_text.push_str(&output.stderr);
533        }
534
535        // Check for error exit code.
536        // exit_code=None means the SSH channel was torn down without delivering
537        // an exit status or exit signal — treat as error, not success.
538        if output.exit_code.map(|code| code != 0).unwrap_or(true) {
539            CallToolResult::error(vec![Content::text(result_text)])
540        } else {
541            CallToolResult::success(vec![Content::text(result_text)])
542        }
543    }
544
545    /// Build exec tool definition (compact)
546    fn exec_tool() -> Tool {
547        tools::exec_tool()
548    }
549
550    /// Build sudo-exec tool definition (compact)
551    fn sudo_exec_tool() -> Tool {
552        tools::sudo_exec_tool()
553    }
554
555    /// Build transfer tool definition (compact)
556    fn transfer_tool() -> Tool {
557        tools::transfer_tool()
558    }
559
560    /// Build check-process tool definition
561    fn check_process_tool() -> Tool {
562        tools::check_process_tool()
563    }
564
565    /// Build read-file tool definition
566    fn read_file_tool() -> Tool {
567        tools::read_file_tool()
568    }
569
570    /// Build write-file tool definition
571    fn write_file_tool() -> Tool {
572        tools::write_file_tool()
573    }
574
575    /// Build replace-in-file tool definition
576    fn replace_in_file_tool() -> Tool {
577        tools::replace_in_file_tool()
578    }
579
580    /// Get extended documentation for a tool by name
581    ///
582    /// Returns the full documentation text that was removed from compact tool definitions
583    /// to save tokens in the MCP protocol.
584    pub fn get_tool_documentation(tool_name: &str) -> Option<&'static str> {
585        tools::get_tool_documentation(tool_name)
586    }
587
588    /// Resolve timeout duration from optional milliseconds, falling back to server default.
589    fn resolve_timeout(&self, timeout_ms: Option<u64>) -> Duration {
590        timeout_ms
591            .map(Duration::from_millis)
592            .unwrap_or(self.timeout)
593    }
594
595    /// Parse tool parameters from JSON with standardized error handling.
596    fn parse_tool_params<T: serde::de::DeserializeOwned>(
597        &self,
598        args: serde_json::Map<String, serde_json::Value>,
599        tool_name: &str,
600    ) -> std::result::Result<T, McpError> {
601        serde_json::from_value(serde_json::Value::Object(args))
602            .map_err(|e| McpError::invalid_params(format!("invalid {tool_name} params: {e}"), None))
603    }
604
605    /// Execute transfer tool with connection management and JSON serialization.
606    async fn execute_transfer(
607        &self,
608        params: TransferParams,
609        verbose: bool,
610    ) -> std::result::Result<CallToolResult, McpError> {
611        let timeout = self.resolve_timeout(params.timeout_ms);
612        let key_path = self.config.key.clone();
613
614        // Ensure connection is established (so errors are deterministic).
615        if let Err(e) = self.connection.ensure_connected().await {
616            let resp = crate::transfer::TransferResponse::error(
617                params,
618                self.transfer.local_root(),
619                &e.to_string(),
620            );
621            let body = resp
622                .to_json(verbose)
623                .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
624            return Ok(CallToolResult::success(vec![Content::text(body)]));
625        }
626
627        let resp = self
628            .transfer
629            .run(
630                &self.connection,
631                params,
632                TransferRunContext {
633                    timeout,
634                    ssh: TransferSshOptions {
635                        host: self.config.host.clone(),
636                        port: self.config.port,
637                        user: self.config.user.clone(),
638                        key_path,
639                        host_key_checking: self.config.strict_host_key_checking,
640                        known_hosts: self.config.known_hosts.clone(),
641                    },
642                },
643            )
644            .await;
645        let body = resp
646            .to_json(verbose)
647            .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
648        Ok(CallToolResult::success(vec![Content::text(body)]))
649    }
650}
651
652impl ServerHandler for SshMcpServer {
653    /// Return server information
654    fn get_info(&self) -> ServerInfo {
655        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
656            .with_protocol_version(ProtocolVersion::LATEST)
657            .with_server_info(Implementation::from_build_env())
658            .with_instructions(format!(
659                "SSH MCP Server v{} - Execute commands on {}@{}:{}",
660                env!("CARGO_PKG_VERSION"),
661                self.config.user,
662                self.config.host,
663                self.config.port,
664            ))
665    }
666
667    /// List available tools
668    async fn list_tools(
669        &self,
670        _request: Option<PaginatedRequestParams>,
671        _context: RequestContext<RoleServer>,
672    ) -> std::result::Result<ListToolsResult, McpError> {
673        debug!("list_tools called");
674
675        let mut tools = vec![Self::exec_tool()];
676
677        // Docs/expected order: exec, (optional) sudo-exec, check-process, transfer, read-file, write-file, replace-in-file.
678        if !self.config.disable_sudo {
679            tools.push(Self::sudo_exec_tool());
680        }
681        tools.push(Self::check_process_tool());
682        tools.push(Self::transfer_tool());
683        tools.push(Self::read_file_tool());
684        tools.push(Self::write_file_tool());
685        tools.push(Self::replace_in_file_tool());
686
687        Ok(ListToolsResult {
688            tools,
689            next_cursor: None,
690            meta: Default::default(),
691        })
692    }
693
694    /// Call a tool
695    async fn call_tool(
696        &self,
697        request: CallToolRequestParams,
698        _context: RequestContext<RoleServer>,
699    ) -> std::result::Result<CallToolResult, McpError> {
700        let tool_name: &str = request.name.as_ref();
701        debug!("call_tool called: {:?}", tool_name);
702
703        let args = request.arguments.unwrap_or_default();
704
705        // Route to the appropriate tool
706        match tool_name {
707            "exec" => {
708                let parsed = self.parse_common_tool_args(&args)?;
709                let timeout = self.resolve_timeout(parsed.timeout_ms);
710
711                if parsed.background {
712                    self.execute_background_command(&parsed.command, parsed.log_path.as_deref())
713                        .await
714                } else {
715                    self.execute_command_with_timeout(&parsed.command, timeout)
716                        .await
717                }
718            }
719            "sudo_exec" | "sudo-exec" => {
720                if self.config.disable_sudo {
721                    return Err(McpError::invalid_params("sudo-exec tool is disabled", None));
722                }
723
724                let parsed = self.parse_common_tool_args(&args)?;
725                let timeout = self.resolve_timeout(parsed.timeout_ms);
726
727                if parsed.background {
728                    self.execute_background_sudo_command(
729                        &parsed.command,
730                        parsed.log_path.as_deref(),
731                    )
732                    .await
733                } else {
734                    self.execute_sudo_command_with_timeout(&parsed.command, timeout)
735                        .await
736                }
737            }
738            "transfer" => {
739                let params: TransferParams = self.parse_tool_params(args, "transfer")?;
740                let verbose = params.verbose;
741                self.execute_transfer(params, verbose).await
742            }
743            "check-process" | "check_process" => {
744                let params: CheckProcessParams = self.parse_tool_params(args, "check-process")?;
745                self.execute_check_process(params).await
746            }
747            "read-file" | "read_file" => {
748                let params: ReadFileParams = self.parse_tool_params(args, "read-file")?;
749                self.execute_read_file(params).await
750            }
751            "write-file" => {
752                let params: WriteFileParams = self.parse_tool_params(args, "write-file")?;
753                self.execute_write_file(
754                    params,
755                    crate::server::handlers::file_edit_common::FileEditFaultInjection::None,
756                )
757                .await
758            }
759            "replace-in-file" => {
760                let params: ReplaceInFileParams =
761                    self.parse_tool_params(args, "replace-in-file")?;
762                self.execute_replace_in_file(
763                    params,
764                    crate::server::handlers::file_edit_common::FileEditFaultInjection::None,
765                )
766                .await
767            }
768            _ => Err(McpError::invalid_params(
769                format!("Unknown tool: {}", tool_name),
770                None,
771            )),
772        }
773    }
774}
775
776#[cfg(test)]
777mod tests {
778    use super::*;
779    use crate::background::response::{
780        BACKGROUND_JSON_SNIPPET_LIMIT_CHARS, background_json_err, background_json_timeout,
781    };
782    use crate::background::wrapper::remote_job_log_path;
783    use crate::server::validation::common::validate_read_file_path;
784    use crate::server::validation::read_file::sanitize_read_file_stderr_snippet;
785    use crate::server::validation::read_file::{
786        normalize_optional_sha256_hex, normalize_sha256_hex,
787    };
788
789    fn extract_text_from_result(result: &CallToolResult) -> String {
790        result
791            .content
792            .iter()
793            .filter_map(|c| c.raw.as_text().map(|text| text.text.clone()))
794            .collect::<Vec<_>>()
795            .join("\n")
796    }
797
798    // Note: Real tests would require a mock SSH server or testcontainers
799    // These are placeholder tests
800
801    #[test]
802    fn test_server_info() {
803        // Verify the package version is defined
804        assert!(!env!("CARGO_PKG_VERSION").is_empty());
805    }
806
807    #[test]
808    fn test_exec_tool_definition() {
809        let tool = SshMcpServer::exec_tool();
810        assert_eq!(tool.name.as_ref(), "exec");
811        assert!(tool.description.is_some());
812    }
813
814    #[test]
815    fn test_sudo_exec_tool_definition() {
816        let tool = SshMcpServer::sudo_exec_tool();
817        assert_eq!(tool.name.as_ref(), "sudo-exec");
818        assert!(tool.description.is_some());
819    }
820
821    #[test]
822    fn test_read_file_tool_definition() {
823        let tool = SshMcpServer::read_file_tool();
824        assert_eq!(tool.name.as_ref(), "read-file");
825        assert!(tool.description.is_some());
826    }
827
828    #[test]
829    fn test_write_file_tool_definition() {
830        let tool = SshMcpServer::write_file_tool();
831        assert_eq!(tool.name.as_ref(), "write-file");
832        assert!(tool.description.is_some());
833    }
834
835    #[test]
836    fn test_replace_in_file_tool_definition() {
837        let tool = SshMcpServer::replace_in_file_tool();
838        assert_eq!(tool.name.as_ref(), "replace-in-file");
839        assert!(tool.description.is_some());
840    }
841
842    #[test]
843    fn test_build_background_wrapper_full_escapes_single_quotes_in_user_command() {
844        let remote_log = remote_job_log_path("job-1");
845        let script =
846            build_background_wrapper_script_full("job-1", "echo 'hello world'", &remote_log);
847        assert!(script.contains("exec sh -lc 'set +m; echo '\"'\"'hello world'\"'\"''"));
848    }
849
850    #[test]
851    fn test_build_background_wrapper_portable_is_busybox_friendly() {
852        let remote_log = remote_job_log_path("job-1");
853        let script = build_background_wrapper_script_portable("job-1", "echo test", &remote_log);
854        assert!(!script.contains("dirname --"));
855        assert!(!script.contains("mkdir -p --"));
856        assert!(!script.contains("sh -lc"));
857        assert!(script.contains("exec sh -c"));
858        assert!(!script.contains("nohup"));
859    }
860
861    #[test]
862    fn test_background_wrappers_emit_markers_and_exec() {
863        let remote_log = remote_job_log_path("job-1");
864
865        let full = build_background_wrapper_script_full("job-1", "echo test", &remote_log);
866        assert!(full.contains("__SSH_MCP_JOB_ID=job-1"));
867        assert!(full.contains("__SSH_MCP_PID=$$"));
868        assert!(full.contains("__SSH_MCP_LOG=$LOG"));
869        assert!(full.contains("exec sh -lc"));
870
871        let portable = build_background_wrapper_script_portable("job-1", "echo test", &remote_log);
872        assert!(portable.contains("__SSH_MCP_JOB_ID=job-1"));
873        assert!(portable.contains("__SSH_MCP_PID=$$"));
874        assert!(portable.contains("__SSH_MCP_LOG=$LOG"));
875        assert!(portable.contains("exec sh -c"));
876    }
877
878    #[test]
879    fn test_background_wrappers_do_not_redirect_remote_output() {
880        let remote_log = remote_job_log_path("job-1");
881        let full = build_background_wrapper_script_full("job-1", "echo test", &remote_log);
882        assert!(!full.contains(">$LOG"));
883        assert!(!full.contains("2>&1"));
884        assert!(!full.contains("$EXIT"));
885        assert!(!full.contains("nohup"));
886
887        let portable = build_background_wrapper_script_portable("job-1", "echo test", &remote_log);
888        assert!(!portable.contains(">$LOG"));
889        assert!(!portable.contains("2>&1"));
890        assert!(!portable.contains("$EXIT"));
891        assert!(!portable.contains("nohup"));
892    }
893
894    #[test]
895    fn test_validate_background_log_path_rejects_leading_dash() {
896        let err =
897            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "-not-a-path").unwrap_err();
898        assert!(err.contains("start with '-'") || err.contains("start with"));
899    }
900
901    #[test]
902    fn test_validate_background_log_path_rejects_newlines() {
903        assert!(
904            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\nrm -rf /").is_err()
905        );
906        assert!(
907            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\rrm -rf /").is_err()
908        );
909    }
910
911    #[test]
912    fn test_validate_read_file_path_requires_absolute() {
913        let err = validate_read_file_path("relative/path").unwrap_err();
914        assert!(err.contains("absolute"));
915    }
916
917    #[test]
918    fn test_validate_read_file_path_rejects_trailing_slash() {
919        let err = validate_read_file_path("/etc/").unwrap_err();
920        assert!(err.contains("must not end with '/'"));
921    }
922
923    #[test]
924    fn test_normalize_sha256_hex_accepts_uppercase_input() {
925        let input = "AABBCCDDEEFF00112233445566778899AABBCCDDEEFF00112233445566778899";
926        let normalized = normalize_sha256_hex(input, "expected_sha256").unwrap();
927        assert_eq!(
928            normalized,
929            "aabbccddeeff00112233445566778899aabbccddeeff00112233445566778899"
930        );
931    }
932
933    #[test]
934    fn test_normalize_sha256_hex_rejects_invalid_length() {
935        let err = normalize_sha256_hex("abcd", "expected_sha256").unwrap_err();
936        assert!(err.contains("64-character"));
937    }
938
939    #[test]
940    fn test_normalize_optional_sha256_hex_treats_blank_as_absent() {
941        let normalized = normalize_optional_sha256_hex(Some("   "), "expected_sha256").unwrap();
942        assert!(normalized.is_none());
943    }
944
945    #[test]
946    fn test_resolve_read_file_max_bytes_uses_token_limit() {
947        assert_eq!(
948            resolve_read_file_max_bytes(Some(12_000)),
949            12_000 * READ_FILE_BYTES_PER_TOKEN
950        );
951    }
952
953    #[test]
954    fn test_resolve_read_file_max_bytes_none_uses_hard_cap() {
955        assert_eq!(resolve_read_file_max_bytes(None), READ_FILE_HARD_MAX_BYTES);
956    }
957
958    #[test]
959    fn test_resolve_read_file_max_bytes_applies_hard_cap() {
960        let very_large_tokens = READ_FILE_HARD_MAX_BYTES;
961        assert_eq!(
962            resolve_read_file_max_bytes(Some(very_large_tokens)),
963            READ_FILE_HARD_MAX_BYTES
964        );
965    }
966
967    #[test]
968    fn test_estimate_tokens_from_bytes_rounds_up() {
969        assert_eq!(estimate_tokens_from_bytes(0), 0);
970        assert_eq!(estimate_tokens_from_bytes(1), 1);
971        assert_eq!(estimate_tokens_from_bytes(4), 1);
972        assert_eq!(estimate_tokens_from_bytes(5), 2);
973    }
974
975    #[test]
976    fn test_resolve_read_file_line_limit_defaults_to_preview_window() {
977        let preview = resolve_read_file_line_limit(ReadFileMode::Preview, None)
978            .expect("preview lines should resolve");
979        assert_eq!(preview, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
980
981        let head = resolve_read_file_line_limit(ReadFileMode::Head, None)
982            .expect("head lines should resolve");
983        assert_eq!(head, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
984
985        let tail = resolve_read_file_line_limit(ReadFileMode::Tail, None)
986            .expect("tail lines should resolve");
987        assert_eq!(tail, Some(READ_FILE_DEFAULT_PREVIEW_LINES));
988    }
989
990    #[test]
991    fn test_resolve_read_file_line_limit_for_full_ignores_lines() {
992        let full = resolve_read_file_line_limit(ReadFileMode::Full, Some(123))
993            .expect("full mode should ignore lines");
994        assert_eq!(full, None);
995    }
996
997    #[test]
998    fn test_resolve_read_file_line_limit_rejects_zero() {
999        let err = resolve_read_file_line_limit(ReadFileMode::Head, Some(0)).unwrap_err();
1000        assert!(err.contains("positive"));
1001    }
1002
1003    #[test]
1004    fn test_resolve_read_file_line_limit_rejects_too_large() {
1005        let err =
1006            resolve_read_file_line_limit(ReadFileMode::Tail, Some(READ_FILE_MAX_LINE_WINDOW + 1))
1007                .unwrap_err();
1008        assert!(err.contains("<="));
1009    }
1010
1011    #[test]
1012    fn test_sanitize_read_file_stderr_snippet_normalizes_whitespace_and_controls() {
1013        let stderr = "line1\nline2\t\u{0007}bad\rline3";
1014        let snippet = sanitize_read_file_stderr_snippet(stderr)
1015            .expect("snippet should be present for non-empty stderr");
1016        assert_eq!(snippet, "line1 line2 bad line3");
1017    }
1018
1019    #[test]
1020    fn test_background_json_err_sets_truncation_flag_and_hint() {
1021        let long_error = "e".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
1022        let long_stderr = "s".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
1023
1024        let result =
1025            background_json_err("job-1", "/tmp/ssh-mcp/job-1.log", &long_error, &long_stderr);
1026        let text = extract_text_from_result(&result);
1027
1028        let value: serde_json::Value =
1029            serde_json::from_str(text.trim()).expect("background_json_err should return JSON");
1030
1031        assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
1032        assert_eq!(
1033            value.get("background").and_then(|v| v.as_bool()),
1034            Some(true)
1035        );
1036        assert_eq!(value.get("truncated").and_then(|v| v.as_bool()), Some(true));
1037
1038        let fields = value
1039            .get("truncated_fields")
1040            .expect("expected truncated_fields");
1041        assert_eq!(fields.get("error").and_then(|v| v.as_bool()), Some(true));
1042        assert_eq!(fields.get("stderr").and_then(|v| v.as_bool()), Some(true));
1043
1044        let hint = value
1045            .get("hint")
1046            .and_then(|v| v.as_str())
1047            .expect("expected hint when truncated");
1048        assert!(
1049            hint.contains("check-process") && hint.contains("job_id=job-1"),
1050            "hint should point to check-process job_id; got: '{hint}'"
1051        );
1052
1053        let error_snippet = value
1054            .get("error")
1055            .and_then(|v| v.as_str())
1056            .expect("expected error field");
1057        assert_eq!(
1058            error_snippet.chars().count(),
1059            BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
1060        );
1061        let stderr_snippet = value
1062            .get("stderr")
1063            .and_then(|v| v.as_str())
1064            .expect("expected stderr field");
1065        assert_eq!(
1066            stderr_snippet.chars().count(),
1067            BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
1068        );
1069    }
1070
1071    #[test]
1072    fn test_background_json_timeout_hint_contains_pid_and_check_process_tool() {
1073        let result = background_json_timeout(
1074            "job-42",
1075            4242,
1076            "/tmp/ssh-mcp/local.log",
1077            &crate::background::response::BackgroundTimeoutSnapshot {
1078                state: "running",
1079                still_running: true,
1080                exit_code: None,
1081                state_reason: None,
1082                elapsed_time: "00:01",
1083                log_exists: true,
1084                log_tail: "tail line",
1085                tail_lines_used: 50,
1086            },
1087        );
1088        let text = extract_text_from_result(&result);
1089
1090        let value: serde_json::Value =
1091            serde_json::from_str(text.trim()).expect("background_json_timeout should return JSON");
1092
1093        assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
1094        assert_eq!(value.get("timeout").and_then(|v| v.as_bool()), Some(true));
1095        assert_eq!(
1096            value.get("background").and_then(|v| v.as_bool()),
1097            Some(true)
1098        );
1099        assert_eq!(
1100            value.get("still_running").and_then(|v| v.as_bool()),
1101            Some(true)
1102        );
1103        assert_eq!(value.get("state").and_then(|v| v.as_str()), Some("running"));
1104        assert_eq!(
1105            value.get("tail_lines_used").and_then(|v| v.as_u64()),
1106            Some(50)
1107        );
1108        assert_eq!(
1109            value.get("elapsed_time").and_then(|v| v.as_str()),
1110            Some("00:01")
1111        );
1112        assert_eq!(
1113            value.get("log_tail").and_then(|v| v.as_str()),
1114            Some("tail line")
1115        );
1116
1117        let hint = value
1118            .get("hint")
1119            .and_then(|v| v.as_str())
1120            .expect("expected hint field");
1121
1122        // Hint should contain the actual job_id value
1123        assert!(
1124            hint.contains("job_id=job-42"),
1125            "hint should contain the actual job_id value; got: '{hint}'"
1126        );
1127        // Hint should mention the check-process tool
1128        assert!(
1129            hint.contains("check-process"),
1130            "hint should mention check-process tool; got: '{hint}'"
1131        );
1132        // Hint should warn against restarting
1133        assert!(
1134            hint.contains("DO NOT restart"),
1135            "hint should warn against restarting; got: '{hint}'"
1136        );
1137        // Hint should use TIMEOUT_RECOVERY prefix
1138        assert!(
1139            hint.contains("TIMEOUT_RECOVERY"),
1140            "hint should start with TIMEOUT_RECOVERY; got: '{hint}'"
1141        );
1142        // Hint should NOT contain old placeholders
1143        assert!(
1144            !hint.contains("<pid>"),
1145            "hint should not contain <pid> placeholder; got: '{hint}'"
1146        );
1147        assert!(
1148            !hint.contains("<log_path>"),
1149            "hint should not contain <log_path> placeholder; got: '{hint}'"
1150        );
1151    }
1152
1153    #[test]
1154    fn test_tool_documentation_available() {
1155        // Verify that extended documentation is available for all tools
1156        assert!(SshMcpServer::get_tool_documentation("exec").is_some());
1157        assert!(SshMcpServer::get_tool_documentation("sudo-exec").is_some());
1158        assert!(SshMcpServer::get_tool_documentation("transfer").is_some());
1159        assert!(SshMcpServer::get_tool_documentation("read-file").is_some());
1160        assert!(SshMcpServer::get_tool_documentation("write-file").is_some());
1161        assert!(SshMcpServer::get_tool_documentation("replace-in-file").is_some());
1162        assert!(SshMcpServer::get_tool_documentation("unknown").is_none());
1163    }
1164
1165    #[test]
1166    fn test_exec_documentation_content() {
1167        let docs = SshMcpServer::get_tool_documentation("exec").unwrap();
1168        assert!(docs.contains("EXEC TOOL"));
1169        assert!(docs.contains("PARAMETERS:"));
1170        assert!(docs.contains("BACKGROUND MODE:"));
1171        assert!(docs.contains("command"));
1172        assert!(docs.contains("background"));
1173        assert!(docs.contains("still_running"));
1174    }
1175
1176    #[test]
1177    fn test_sudo_exec_documentation_content() {
1178        let docs = SshMcpServer::get_tool_documentation("sudo-exec").unwrap();
1179        assert!(docs.contains("SUDO-EXEC TOOL"));
1180        assert!(docs.contains("sudo"));
1181    }
1182
1183    #[test]
1184    fn test_transfer_documentation_content() {
1185        let docs = SshMcpServer::get_tool_documentation("transfer").unwrap();
1186        assert!(docs.contains("TRANSFER TOOL"));
1187        assert!(docs.contains("put"));
1188        assert!(docs.contains("get"));
1189        assert!(docs.contains("TRANSPORTS:"));
1190    }
1191
1192    #[test]
1193    fn test_read_file_documentation_content() {
1194        let docs = SshMcpServer::get_tool_documentation("read-file").unwrap();
1195        assert!(docs.contains("READ-FILE TOOL"));
1196        assert!(docs.contains("remote_path"));
1197        assert!(docs.contains("mode"));
1198        assert!(docs.contains("UTF-8"));
1199    }
1200
1201    #[test]
1202    fn test_write_file_documentation_content() {
1203        let docs = SshMcpServer::get_tool_documentation("write-file").unwrap();
1204        assert!(docs.contains("WRITE-FILE TOOL"));
1205        assert!(docs.contains("expected_sha256"));
1206        assert!(docs.contains("atomic"));
1207    }
1208
1209    #[test]
1210    fn test_replace_in_file_documentation_content() {
1211        let docs = SshMcpServer::get_tool_documentation("replace-in-file").unwrap();
1212        assert!(docs.contains("REPLACE-IN-FILE TOOL"));
1213        assert!(docs.contains("old_text"));
1214        assert!(docs.contains("replace_all"));
1215    }
1216
1217    #[test]
1218    fn test_compact_tool_descriptions() {
1219        // Verify that tool descriptions are compact (not verbose)
1220        let exec = SshMcpServer::exec_tool();
1221        let sudo_exec = SshMcpServer::sudo_exec_tool();
1222        let transfer = SshMcpServer::transfer_tool();
1223        let read_file = SshMcpServer::read_file_tool();
1224        let write_file = SshMcpServer::write_file_tool();
1225        let replace_in_file = SshMcpServer::replace_in_file_tool();
1226
1227        // Descriptions should be present but concise (under 100 chars)
1228        if let Some(desc) = exec.description {
1229            assert!(
1230                desc.len() < 100,
1231                "exec description too long: {} chars",
1232                desc.len()
1233            );
1234        }
1235        if let Some(desc) = sudo_exec.description {
1236            assert!(
1237                desc.len() < 100,
1238                "sudo-exec description too long: {} chars",
1239                desc.len()
1240            );
1241        }
1242        if let Some(desc) = transfer.description {
1243            assert!(
1244                desc.len() < 100,
1245                "transfer description too long: {} chars",
1246                desc.len()
1247            );
1248        }
1249        if let Some(desc) = read_file.description {
1250            assert!(
1251                desc.len() < 100,
1252                "read-file description too long: {} chars",
1253                desc.len()
1254            );
1255        }
1256        if let Some(desc) = write_file.description {
1257            assert!(
1258                desc.len() < 100,
1259                "write-file description too long: {} chars",
1260                desc.len()
1261            );
1262        }
1263        if let Some(desc) = replace_in_file.description {
1264            assert!(
1265                desc.len() < 100,
1266                "replace-in-file description too long: {} chars",
1267                desc.len()
1268            );
1269        }
1270    }
1271}