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 `shell` and `sudo_shell` tools.
5
6use std::path::{Path, PathBuf};
7use std::sync::Arc;
8use std::sync::atomic::{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 tokio::task::JoinHandle;
20use tokio_util::sync::CancellationToken;
21use tracing::{debug, error, info, warn};
22
23use crate::background::job::NewRunningJob;
24use crate::background::transfer::{SharedTransferJob, TransferJobRegistry};
25use crate::background::{JobRegistry, JobState, LocalLogSpooler, SharedJobState};
26use crate::config::Config;
27use crate::error::{Result, SshMcpError};
28#[cfg(unix)]
29use crate::platform::O_NOFOLLOW_FLAG;
30use crate::server::handlers::file_edit_common::{FileEditFaultInjection, FileEditPrivilege};
31#[cfg(test)]
32use crate::server::validation::validate_background_log_path;
33use crate::ssh::{
34    CommandOutput, SshConfig, SshConnectionManager, sanitize_command, wrap_sudo_command,
35};
36use crate::tools::ApplyPatchParams;
37use crate::transfer::{
38    TransferEngine, TransferEventSink, TransferParams, TransferResponse, TransferRunContext,
39    TransferSshOptions,
40};
41
42mod args;
43mod exec;
44mod handlers;
45mod testing;
46mod tools;
47mod validation;
48
49const BACKGROUND_START_TIMEOUT: Duration = Duration::from_secs(20);
50
51const JOB_COMPLETED_RETENTION: Duration = Duration::from_secs(60 * 60);
52
53static JOB_COUNTER: AtomicU64 = AtomicU64::new(0);
54
55fn make_job_id() -> String {
56    let counter = JOB_COUNTER.fetch_add(1, Ordering::Relaxed);
57    let epoch_ms = SystemTime::now()
58        .duration_since(UNIX_EPOCH)
59        .map(|d| d.as_millis())
60        .unwrap_or(0);
61    format!("{}-{}", epoch_ms, counter)
62}
63
64/// SeSSHion SSH MCP server
65///
66/// The main server implementation that provides MCP tools for remote SSH
67/// command execution.
68#[derive(Clone)]
69pub struct SshMcpServer {
70    /// Server configuration
71    config: Config,
72
73    /// SSH connection manager
74    connection: Arc<SshConnectionManager>,
75
76    /// Command execution timeout
77    timeout: Duration,
78
79    /// Maximum command length
80    max_chars: Option<usize>,
81
82    spooler: Arc<LocalLogSpooler>,
83    job_registry: Arc<JobRegistry>,
84
85    transfer_job_registry: Arc<TransferJobRegistry>,
86    transfer_tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
87    transfer_shutdown: CancellationToken,
88
89    transfer: TransferEngine,
90}
91
92struct TransferTerminalGuard {
93    job: SharedTransferJob,
94    fallback: Option<TransferResponse>,
95}
96
97impl TransferTerminalGuard {
98    fn new(job: SharedTransferJob, fallback: TransferResponse) -> Self {
99        Self {
100            job,
101            fallback: Some(fallback),
102        }
103    }
104
105    fn finish(mut self, response: &TransferResponse) {
106        if let Ok(mut job) = self.job.lock() {
107            job.finish(response);
108        }
109        self.fallback = None;
110    }
111}
112
113impl Drop for TransferTerminalGuard {
114    fn drop(&mut self) {
115        let Some(fallback) = self.fallback.take() else {
116            return;
117        };
118        if let Ok(mut job) = self.job.lock()
119            && !job.is_terminal()
120        {
121            job.finish(&fallback);
122        }
123    }
124}
125
126impl SshMcpServer {
127    /// Create a new SeSSHion server
128    ///
129    /// This sets up the SSH connection manager based on the provided configuration.
130    /// Connection is not established until a tool is actually used.
131    pub async fn new(config: Config) -> Result<Self> {
132        Self::new_with_spool_dir(config, None).await
133    }
134
135    /// Create a new SeSSHion server with an optional local spool directory.
136    pub async fn new_with_spool_dir(config: Config, spool_dir: Option<PathBuf>) -> Result<Self> {
137        let local_root = std::env::current_dir()?;
138
139        let spooler = Arc::new(resolve_local_spooler(spool_dir)?);
140        spooler.ensure_dir().await.map_err(|e| {
141            SshMcpError::Config(format!(
142                "failed to initialize local log spool dir {}: {e}",
143                spooler.base_dir().display()
144            ))
145        })?;
146        let job_registry = Arc::new(JobRegistry::new(JOB_COMPLETED_RETENTION));
147        let transfer_job_registry = Arc::new(TransferJobRegistry::new(JOB_COMPLETED_RETENTION));
148
149        // Build SSH configuration
150        let mut ssh_config = SshConfig::new(&config.host, &config.user).with_port(config.port);
151
152        // Add authentication
153        if let Some(ref password) = config.password {
154            ssh_config = ssh_config.with_password(password);
155        }
156
157        if let Some(ref key_path) = config.key {
158            // Read the key file
159            let key_content = tokio::fs::read_to_string(key_path)
160                .await
161                .map_err(SshMcpError::Io)?;
162            ssh_config = ssh_config.with_private_key(&key_content);
163        }
164
165        // Add elevation passwords if provided
166        if let Some(ref su_password) = config.su_password {
167            ssh_config = ssh_config.with_su_password(su_password);
168        }
169
170        if let Some(ref sudo_password) = config.sudo_password {
171            ssh_config = ssh_config.with_sudo_password(sudo_password);
172        }
173
174        // Add keepalive settings for human-like connection persistence
175        ssh_config = ssh_config
176            .with_keepalive_interval(config.keepalive_interval)
177            .with_keepalive_max(config.keepalive_max);
178
179        // Add reconnect and health probe settings
180        ssh_config = ssh_config
181            .with_reconnect_retries(config.reconnect_retries)
182            .with_reconnect_backoff_ms(config.reconnect_backoff_ms)
183            .with_health_probe_timeout_ms(config.health_probe_timeout_ms);
184
185        // Add host key verification settings
186        ssh_config = ssh_config
187            .with_host_key_checking(config.strict_host_key_checking)
188            .with_known_hosts(config.known_hosts.clone());
189
190        // Add output token limit for OOM protection
191        ssh_config = ssh_config.with_max_output_tokens(config.max_output_tokens);
192
193        // Create connection manager
194        let connection = Arc::new(SshConnectionManager::new(ssh_config).await);
195
196        let timeout = Duration::from_millis(config.timeout_ms);
197        let max_chars = config.max_chars;
198
199        Ok(Self {
200            config,
201            connection,
202            timeout,
203            max_chars,
204            spooler,
205            job_registry,
206            transfer_job_registry,
207            transfer_tasks: Arc::new(Mutex::new(Vec::new())),
208            transfer_shutdown: CancellationToken::new(),
209            transfer: TransferEngine::new(local_root),
210        })
211    }
212
213    fn connection_id(&self) -> String {
214        format!(
215            "{}@{}:{}",
216            self.config.user, self.config.host, self.config.port
217        )
218    }
219
220    fn default_local_log_path(
221        &self,
222        job_id: &str,
223    ) -> std::result::Result<(PathBuf, String), String> {
224        let path = self
225            .spooler
226            .log_path_for(job_id)
227            .map_err(|e| format!("failed to generate local log path for job_id='{job_id}': {e}"))?;
228        let path_str = path.to_string_lossy().to_string();
229        Ok((path, path_str))
230    }
231
232    async fn ensure_local_log_file(&self, log_path: &Path) -> std::result::Result<(), SshMcpError> {
233        self.spooler.ensure_dir().await.map_err(|e| {
234            SshMcpError::Config(format!(
235                "failed to ensure local log spool dir {}: {e}",
236                self.spooler.base_dir().display()
237            ))
238        })?;
239
240        if log_path.parent() != Some(self.spooler.base_dir()) {
241            return Err(SshMcpError::InvalidParams(format!(
242                "log_path must be directly under {}",
243                self.spooler.base_dir().display()
244            )));
245        }
246
247        match tokio::fs::symlink_metadata(log_path).await {
248            Ok(meta) => {
249                let ft = meta.file_type();
250                if ft.is_symlink() {
251                    return Err(SshMcpError::invalid_params(
252                        "log_path is a symlink (refusing to follow it)",
253                    ));
254                }
255                if !ft.is_file() {
256                    return Err(SshMcpError::invalid_params(
257                        "log_path exists but is not a regular file",
258                    ));
259                }
260            }
261            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
262            Err(e) => return Err(SshMcpError::Io(e)),
263        }
264
265        let mut opts = tokio::fs::OpenOptions::new();
266        opts.write(true).create(true).truncate(true);
267
268        #[cfg(unix)]
269        {
270            opts.custom_flags(O_NOFOLLOW_FLAG);
271        }
272
273        let file = match opts.open(log_path).await {
274            Ok(f) => f,
275            Err(e) => {
276                if let Ok(meta) = tokio::fs::symlink_metadata(log_path).await
277                    && meta.file_type().is_symlink()
278                {
279                    return Err(SshMcpError::invalid_params(
280                        "log_path is a symlink (refusing to follow it)",
281                    ));
282                }
283                return Err(SshMcpError::Io(e));
284            }
285        };
286
287        file.sync_all().await.map_err(SshMcpError::Io)
288    }
289
290    async fn register_running_job(
291        &self,
292        job_id: &str,
293        pid: u32,
294        log_path: PathBuf,
295        command: &str,
296    ) -> SharedJobState {
297        let job = Arc::new(Mutex::new(JobState::new_running(NewRunningJob {
298            job_id: job_id.to_string(),
299            pid,
300            log_path,
301            command: command.to_string(),
302            connection_id: self.connection_id(),
303        })));
304
305        self.job_registry
306            .insert(job_id.to_string(), Arc::clone(&job))
307            .await;
308
309        let persisted = {
310            let guard = job.lock().await;
311            guard.clone()
312        };
313        if let Err(e) = self.spooler.persist_job_state(&persisted).await {
314            warn!(job_id = ?job_id, error = ?e, "failed to persist running job state");
315        }
316
317        job
318    }
319
320    /// Get a reference to the SSH connection manager
321    pub fn connection(&self) -> &Arc<SshConnectionManager> {
322        &self.connection
323    }
324
325    /// Close the server and cleanup resources
326    pub async fn shutdown(&self) {
327        info!("Shutting down SeSSHion...");
328        self.transfer_shutdown.cancel();
329        let tasks = {
330            let mut tasks = self.transfer_tasks.lock().await;
331            std::mem::take(&mut *tasks)
332        };
333        let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
334        for mut task in tasks {
335            if tokio::time::timeout_at(deadline, &mut task).await.is_err() {
336                task.abort();
337                let _ = task.await;
338            }
339        }
340        self.connection.close().await;
341    }
342
343    /// Execute a command (used by shell tool)
344    async fn execute_command_with_timeout(
345        &self,
346        command: &str,
347        timeout: Duration,
348    ) -> std::result::Result<CallToolResult, McpError> {
349        debug!(
350            "shell tool called: cmd_len={}, background=false, sudo=false, timeout_ms={}",
351            command.len(),
352            timeout.as_millis()
353        );
354
355        // Sanitize the command
356        let sanitized = match self.sanitize_or_tool_error(command) {
357            Ok(cmd) => cmd,
358            Err(result) => return Ok(result),
359        };
360
361        // Foreground execution is detachable-by-design:
362        // - Start the command on a dedicated SSH channel
363        // - Stream remote stdout/stderr into a local spool file
364        // - If timeout elapses, return JSON with job_id/pid/log_path while the stream continues
365
366        let requires_elevation = self.connection.get_su_password().is_some();
367        if requires_elevation {
368            if let Err(e) = self.connection.ensure_connected().await {
369                error!(error = ?e, "Failed to ensure SSH connection");
370                return Ok(CallToolResult::error(vec![ContentBlock::text(
371                    e.to_string(),
372                )]));
373            }
374
375            if let Err(e) = self.connection.ensure_elevated().await {
376                debug!(error = ?e, "Elevation failed, will run as normal user");
377            }
378        }
379
380        // Ensure connection is established for detached foreground execution path.
381        if !requires_elevation && let Err(e) = self.connection.ensure_connected().await {
382            error!(error = ?e, "Failed to ensure SSH connection");
383            return Ok(CallToolResult::error(vec![ContentBlock::text(
384                e.to_string(),
385            )]));
386        }
387
388        self.execute_detachable_foreground_impl(&sanitized, &sanitized, timeout)
389            .await
390    }
391
392    async fn execute_command(
393        &self,
394        command: &str,
395    ) -> std::result::Result<CallToolResult, McpError> {
396        self.execute_command_with_timeout(command, self.timeout)
397            .await
398    }
399
400    async fn execute_background_command(
401        &self,
402        command: &str,
403        log_path: Option<&str>,
404    ) -> std::result::Result<CallToolResult, McpError> {
405        self.execute_background_impl(command, log_path, exec::BackgroundPrivilege::Normal)
406            .await
407    }
408
409    /// Execute a command with sudo (used by sudo_shell tool)
410    async fn execute_sudo_command_with_timeout(
411        &self,
412        command: &str,
413        timeout: Duration,
414    ) -> std::result::Result<CallToolResult, McpError> {
415        debug!(
416            "sudo_shell tool called: cmd_len={}, background=false, sudo=true, timeout_ms={}",
417            command.len(),
418            timeout.as_millis()
419        );
420
421        // Sanitize the command
422        let sanitized = match self.sanitize_or_tool_error(command) {
423            Ok(cmd) => cmd,
424            Err(result) => return Ok(result),
425        };
426
427        // Wrap the command with sudo
428        let sudo_password = self.connection.get_sudo_password();
429        let wrapped_command = wrap_sudo_command(&sanitized, sudo_password);
430        debug!(
431            "Wrapped sudo command (password hidden): sudo -n sh -c '...' or printf '...' | sudo ..."
432        );
433
434        if let Err(e) = self.connection.ensure_connected().await {
435            error!(error = ?e, "Failed to ensure SSH connection");
436            return Ok(CallToolResult::error(vec![ContentBlock::text(
437                e.to_string(),
438            )]));
439        }
440
441        self.execute_detachable_foreground_impl(
442            &wrapped_command,
443            &format!("sudo {sanitized}"),
444            timeout,
445        )
446        .await
447    }
448
449    async fn execute_sudo_command(
450        &self,
451        command: &str,
452    ) -> std::result::Result<CallToolResult, McpError> {
453        self.execute_sudo_command_with_timeout(command, self.timeout)
454            .await
455    }
456
457    async fn execute_background_sudo_command(
458        &self,
459        command: &str,
460        log_path: Option<&str>,
461    ) -> std::result::Result<CallToolResult, McpError> {
462        let sudo_password = self.connection.get_sudo_password();
463        self.execute_background_impl(
464            command,
465            log_path,
466            exec::BackgroundPrivilege::Sudo {
467                password: sudo_password,
468            },
469        )
470        .await
471    }
472
473    fn sanitize_or_tool_error(&self, command: &str) -> std::result::Result<String, CallToolResult> {
474        sanitize_command(command, self.max_chars).map_err(|e| {
475            error!(error = ?e, "Command sanitization failed");
476            CallToolResult::error(vec![ContentBlock::text(format!("Error: {}", e))])
477        })
478    }
479
480    fn calltool_from_command_output(output: CommandOutput) -> CallToolResult {
481        // Combine stdout and stderr for the response
482        let mut result_text = output.stdout;
483        if !output.stderr.is_empty() {
484            if !result_text.is_empty() {
485                result_text.push_str("\n--- stderr ---\n");
486            }
487            result_text.push_str(&output.stderr);
488        }
489
490        // Check for error exit code.
491        // exit_code=None means the SSH channel was torn down without delivering
492        // an exit status or exit signal — treat as error, not success.
493        if output.exit_code.map(|code| code != 0).unwrap_or(true) {
494            CallToolResult::error(vec![ContentBlock::text(result_text)])
495        } else {
496            CallToolResult::success(vec![ContentBlock::text(result_text)])
497        }
498    }
499
500    /// Build shell tool definition (compact)
501    fn shell_tool() -> Tool {
502        tools::shell_tool()
503    }
504
505    /// Build sudo_shell tool definition (compact)
506    fn sudo_shell_tool() -> Tool {
507        tools::sudo_shell_tool()
508    }
509
510    /// Build transfer tool definition (compact)
511    fn transfer_tool() -> Tool {
512        tools::transfer_tool()
513    }
514
515    /// Build check_process tool definition
516    fn check_process_tool() -> Tool {
517        tools::check_process_tool()
518    }
519
520    /// Build apply_patch tool definition
521    fn apply_patch_tool() -> Tool {
522        tools::apply_patch_tool()
523    }
524
525    /// Build sudo_apply_patch tool definition
526    fn sudo_apply_patch_tool() -> Tool {
527        tools::sudo_apply_patch_tool()
528    }
529
530    /// Resolve timeout duration from optional milliseconds, falling back to server default.
531    fn resolve_timeout(&self, timeout_ms: Option<u64>) -> Duration {
532        timeout_ms
533            .map(Duration::from_millis)
534            .unwrap_or(self.timeout)
535    }
536
537    /// Parse tool parameters from JSON with standardized error handling.
538    fn parse_tool_params<T: serde::de::DeserializeOwned>(
539        &self,
540        args: serde_json::Map<String, serde_json::Value>,
541        tool_name: &str,
542    ) -> std::result::Result<T, McpError> {
543        serde_json::from_value(serde_json::Value::Object(args))
544            .map_err(|e| McpError::invalid_params(format!("invalid {tool_name} params: {e}"), None))
545    }
546
547    async fn run_transfer_response(
548        &self,
549        params: TransferParams,
550        cancellation: CancellationToken,
551        progress: Option<TransferEventSink>,
552    ) -> TransferResponse {
553        let timeout = self.resolve_timeout(params.timeout_ms);
554        let key_path = self.config.key.clone();
555
556        self.transfer
557            .run_controlled(
558                &self.connection,
559                params,
560                TransferRunContext {
561                    timeout,
562                    ssh: TransferSshOptions {
563                        host: self.config.host.clone(),
564                        port: self.config.port,
565                        user: self.config.user.clone(),
566                        key_path,
567                        host_key_checking: self.config.strict_host_key_checking,
568                        known_hosts: self.config.known_hosts.clone(),
569                    },
570                },
571                cancellation,
572                progress,
573            )
574            .await
575    }
576
577    /// Execute a foreground transfer and serialize its response.
578    async fn execute_transfer(
579        &self,
580        params: TransferParams,
581        verbose: bool,
582        cancellation: CancellationToken,
583    ) -> std::result::Result<CallToolResult, McpError> {
584        let resp = self.run_transfer_response(params, cancellation, None).await;
585        let body = resp
586            .to_json(verbose)
587            .unwrap_or_else(|_| "{\"ok\":false,\"error\":\"serialization_error\"}".to_string());
588        Ok(CallToolResult::success(vec![ContentBlock::text(body)]))
589    }
590
591    async fn execute_background_transfer(
592        &self,
593        params: TransferParams,
594    ) -> std::result::Result<CallToolResult, McpError> {
595        let job_id = format!("transfer-{}", make_job_id());
596        let job = self.transfer_job_registry.register(job_id.clone());
597        let progress_job = Arc::clone(&job);
598        let progress = TransferEventSink::new(move |event| {
599            if let Ok(mut job) = progress_job.lock() {
600                job.apply_event(event);
601            }
602        });
603
604        let server = self.clone();
605        let task_params = params.clone();
606        let fallback = TransferResponse::error(
607            params,
608            self.transfer.local_root(),
609            "background transfer stopped unexpectedly",
610        );
611        let body = serde_json::json!({
612            "ok": true,
613            "background": true,
614            "job_id": job_id,
615            "job_type": "transfer",
616            "state": "running",
617            "phase": "queued",
618        });
619
620        let mut tasks = self.transfer_tasks.lock().await;
621        tasks.retain(|task| !task.is_finished());
622        let task = tokio::spawn(async move {
623            tokio::task::yield_now().await;
624            let terminal = TransferTerminalGuard::new(job, fallback);
625            let response = server
626                .run_transfer_response(
627                    task_params,
628                    server.transfer_shutdown.child_token(),
629                    Some(progress),
630                )
631                .await;
632            terminal.finish(&response);
633        });
634        tasks.push(task);
635        drop(tasks);
636
637        Ok(CallToolResult::success(vec![ContentBlock::text(
638            body.to_string(),
639        )]))
640    }
641}
642
643fn resolve_local_spooler(spool_dir: Option<PathBuf>) -> Result<LocalLogSpooler> {
644    match spool_dir {
645        Some(path) if !path.is_absolute() => Err(SshMcpError::Config(format!(
646            "spool directory must be absolute: {}",
647            path.display()
648        ))),
649        Some(path) => Ok(LocalLogSpooler::new(path)),
650        None => Ok(LocalLogSpooler::new_default()),
651    }
652}
653
654fn server_implementation() -> Implementation {
655    Implementation::new("ssh-mcp", env!("CARGO_PKG_VERSION"))
656        .with_title("SeSSHion")
657        .with_description(env!("CARGO_PKG_DESCRIPTION"))
658        .with_website_url("https://github.com/0FL01/SeSSHion")
659}
660
661impl ServerHandler for SshMcpServer {
662    /// Return server information
663    fn get_info(&self) -> ServerInfo {
664        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
665            .with_protocol_version(ProtocolVersion::LATEST)
666            .with_server_info(server_implementation())
667            .with_instructions(format!(
668                "SeSSHion v{} - SSH MCP server for {}@{}:{}",
669                env!("CARGO_PKG_VERSION"),
670                self.config.user,
671                self.config.host,
672                self.config.port,
673            ))
674    }
675
676    /// List available tools
677    async fn list_tools(
678        &self,
679        _request: Option<PaginatedRequestParams>,
680        _context: RequestContext<RoleServer>,
681    ) -> std::result::Result<ListToolsResult, McpError> {
682        debug!("list_tools called");
683
684        let mut tools = vec![Self::shell_tool()];
685
686        // Docs/expected order: shell, optional sudo tools, check_process, transfer, apply_patch.
687        if !self.config.disable_sudo {
688            tools.push(Self::sudo_shell_tool());
689            tools.push(Self::sudo_apply_patch_tool());
690        }
691        tools.push(Self::check_process_tool());
692        tools.push(Self::transfer_tool());
693        tools.push(Self::apply_patch_tool());
694
695        Ok(ListToolsResult {
696            tools,
697            next_cursor: None,
698            meta: Default::default(),
699        })
700    }
701
702    /// Call a tool
703    async fn call_tool(
704        &self,
705        request: CallToolRequestParams,
706        context: RequestContext<RoleServer>,
707    ) -> std::result::Result<CallToolResult, McpError> {
708        let tool_name: &str = request.name.as_ref();
709        debug!("call_tool called: {:?}", tool_name);
710
711        let args = request.arguments.unwrap_or_default();
712
713        // Route to the appropriate tool
714        match tool_name {
715            "shell" => {
716                let parsed = self.parse_common_tool_args(&args)?;
717                let timeout = self.resolve_timeout(parsed.timeout_ms);
718
719                if parsed.background {
720                    self.execute_background_command(&parsed.command, parsed.log_path.as_deref())
721                        .await
722                } else {
723                    self.execute_command_with_timeout(&parsed.command, timeout)
724                        .await
725                }
726            }
727            "sudo_shell" => {
728                if self.config.disable_sudo {
729                    return Err(McpError::invalid_params(
730                        "sudo_shell tool is disabled",
731                        None,
732                    ));
733                }
734
735                let parsed = self.parse_common_tool_args(&args)?;
736                let timeout = self.resolve_timeout(parsed.timeout_ms);
737
738                if parsed.background {
739                    self.execute_background_sudo_command(
740                        &parsed.command,
741                        parsed.log_path.as_deref(),
742                    )
743                    .await
744                } else {
745                    self.execute_sudo_command_with_timeout(&parsed.command, timeout)
746                        .await
747                }
748            }
749            "transfer" => {
750                let params: TransferParams = self.parse_tool_params(args, "transfer")?;
751                let verbose = params.verbose;
752                if params.background {
753                    self.execute_background_transfer(params).await
754                } else {
755                    self.execute_transfer(params, verbose, context.ct.clone())
756                        .await
757                }
758            }
759            "check_process" => {
760                let params: args::CheckProcessToolArgs =
761                    self.parse_tool_params(args, "check_process")?;
762                self.execute_check_process(params.check, params.wait_for, context.ct.cancelled())
763                    .await
764            }
765            "apply_patch" => {
766                let params: ApplyPatchParams = self.parse_tool_params(args, "apply_patch")?;
767                self.execute_apply_patch(
768                    params,
769                    FileEditFaultInjection::None,
770                    FileEditPrivilege::User,
771                )
772                .await
773            }
774            "sudo_apply_patch" => {
775                if self.config.disable_sudo {
776                    return Err(McpError::invalid_params(
777                        "sudo_apply_patch tool is disabled",
778                        None,
779                    ));
780                }
781
782                let params: ApplyPatchParams = self.parse_tool_params(args, "sudo_apply_patch")?;
783                self.execute_apply_patch(
784                    params,
785                    FileEditFaultInjection::None,
786                    FileEditPrivilege::Sudo,
787                )
788                .await
789            }
790            _ => Err(McpError::invalid_params(
791                format!("Unknown tool: {}", tool_name),
792                None,
793            )),
794        }
795    }
796}
797
798#[cfg(test)]
799mod tests {
800    use super::*;
801    use crate::background::response::{
802        BACKGROUND_JSON_SNIPPET_LIMIT_CHARS, background_json_err, background_json_timeout,
803    };
804    use crate::background::wrapper::{build_background_wrapper_script, remote_job_log_path};
805
806    fn extract_text_from_result(result: &CallToolResult) -> String {
807        result
808            .content
809            .iter()
810            .filter_map(|c| c.as_text().map(|text| text.text.clone()))
811            .collect::<Vec<_>>()
812            .join("\n")
813    }
814
815    #[test]
816    fn test_server_info() {
817        let implementation = server_implementation();
818
819        assert_eq!(implementation.name, "ssh-mcp");
820        assert_eq!(implementation.title.as_deref(), Some("SeSSHion"));
821        assert_eq!(implementation.version, env!("CARGO_PKG_VERSION"));
822        assert_eq!(
823            implementation.website_url.as_deref(),
824            Some("https://github.com/0FL01/SeSSHion")
825        );
826        assert_eq!(
827            implementation.description.as_deref(),
828            Some(env!("CARGO_PKG_DESCRIPTION"))
829        );
830    }
831
832    #[test]
833    fn test_resolve_local_spooler_rejects_relative_override() {
834        let error = resolve_local_spooler(Some(PathBuf::from("relative/spool")))
835            .expect_err("relative spool directory must be rejected");
836
837        assert!(matches!(
838            error,
839            SshMcpError::Config(message) if message.contains("must be absolute")
840        ));
841    }
842
843    #[test]
844    fn test_shell_tool_definition() {
845        let tool = SshMcpServer::shell_tool();
846        assert_eq!(tool.name.as_ref(), "shell");
847        assert!(tool.description.is_some());
848    }
849
850    #[test]
851    fn test_sudo_shell_tool_definition() {
852        let tool = SshMcpServer::sudo_shell_tool();
853        assert_eq!(tool.name.as_ref(), "sudo_shell");
854        assert!(tool.description.is_some());
855    }
856
857    #[test]
858    fn test_apply_patch_tool_definition() {
859        let tool = SshMcpServer::apply_patch_tool();
860        assert_eq!(tool.name.as_ref(), "apply_patch");
861        assert!(tool.description.is_some());
862    }
863
864    #[test]
865    fn test_sudo_apply_patch_tool_definition() {
866        let tool = SshMcpServer::sudo_apply_patch_tool();
867        assert_eq!(tool.name.as_ref(), "sudo_apply_patch");
868        assert!(tool.description.is_some());
869    }
870
871    #[test]
872    fn test_build_background_wrapper_escapes_single_quotes_in_user_command() {
873        let remote_log = remote_job_log_path("job-1");
874        let script = build_background_wrapper_script("job-1", "echo 'hello world'", &remote_log);
875        assert!(script.contains("exec sh -c 'set +m; echo '\"'\"'hello world'\"'\"''"));
876    }
877
878    #[test]
879    fn test_build_background_wrapper_is_busybox_friendly() {
880        let remote_log = remote_job_log_path("job-1");
881        let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
882        assert!(!script.contains("dirname --"));
883        assert!(!script.contains("mkdir -p --"));
884        assert!(!script.contains("sh -lc"));
885        assert!(script.contains("exec sh -c"));
886        assert!(!script.contains("nohup"));
887    }
888
889    #[test]
890    fn test_background_wrapper_emits_markers_and_exec() {
891        let remote_log = remote_job_log_path("job-1");
892        let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
893        assert!(script.contains("__SSH_MCP_JOB_ID=job-1"));
894        assert!(script.contains("__SSH_MCP_PID=$$"));
895        assert!(script.contains("__SSH_MCP_LOG=$LOG"));
896        assert!(script.contains("exec sh -c"));
897    }
898
899    #[test]
900    fn test_background_wrapper_does_not_redirect_remote_output() {
901        let remote_log = remote_job_log_path("job-1");
902        let script = build_background_wrapper_script("job-1", "echo test", &remote_log);
903        assert!(!script.contains(">$LOG"));
904        assert!(!script.contains("2>&1"));
905        assert!(!script.contains("$EXIT"));
906        assert!(!script.contains("nohup"));
907    }
908
909    #[test]
910    fn test_validate_background_log_path_rejects_leading_dash() {
911        let err =
912            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "-not-a-path").unwrap_err();
913        assert!(err.contains("start with '-'") || err.contains("start with"));
914    }
915
916    #[test]
917    fn test_validate_background_log_path_rejects_newlines() {
918        assert!(
919            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\nrm -rf /").is_err()
920        );
921        assert!(
922            validate_background_log_path(Path::new("/tmp/ssh-mcp"), "/tmp/x\rrm -rf /").is_err()
923        );
924    }
925
926    #[test]
927    fn test_background_json_err_omits_unregistered_job_fields() {
928        let long_error = "e".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
929        let long_stderr = "s".repeat(BACKGROUND_JSON_SNIPPET_LIMIT_CHARS + 10);
930
931        let result = background_json_err(&long_error, &long_stderr);
932        let text = extract_text_from_result(&result);
933
934        let value: serde_json::Value =
935            serde_json::from_str(text.trim()).expect("background_json_err should return JSON");
936
937        assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
938        assert_eq!(
939            value.get("background").and_then(|v| v.as_bool()),
940            Some(true)
941        );
942        assert_eq!(value.get("truncated").and_then(|v| v.as_bool()), Some(true));
943        assert!(value.get("job_id").is_none());
944        assert!(value.get("log_path").is_none());
945        assert!(value.get("hint").is_none());
946
947        let fields = value
948            .get("truncated_fields")
949            .expect("expected truncated_fields");
950        assert_eq!(fields.get("error").and_then(|v| v.as_bool()), Some(true));
951        assert_eq!(fields.get("stderr").and_then(|v| v.as_bool()), Some(true));
952
953        let error_snippet = value
954            .get("error")
955            .and_then(|v| v.as_str())
956            .expect("expected error field");
957        assert_eq!(
958            error_snippet.chars().count(),
959            BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
960        );
961        let stderr_snippet = value
962            .get("stderr")
963            .and_then(|v| v.as_str())
964            .expect("expected stderr field");
965        assert_eq!(
966            stderr_snippet.chars().count(),
967            BACKGROUND_JSON_SNIPPET_LIMIT_CHARS
968        );
969    }
970
971    #[test]
972    fn test_background_json_timeout_hint_contains_pid_and_check_process_tool() {
973        let result = background_json_timeout(
974            "job-42",
975            4242,
976            "/tmp/ssh-mcp/local.log",
977            &crate::background::response::BackgroundTimeoutSnapshot {
978                state: "running",
979                still_running: true,
980                exit_code: None,
981                state_reason: None,
982                elapsed_time: "00:01",
983                log_exists: true,
984                log_tail: "tail line",
985                tail_lines_used: 50,
986            },
987        );
988        let text = extract_text_from_result(&result);
989
990        let value: serde_json::Value =
991            serde_json::from_str(text.trim()).expect("background_json_timeout should return JSON");
992
993        assert_eq!(value.get("ok").and_then(|v| v.as_bool()), Some(false));
994        assert_eq!(value.get("timeout").and_then(|v| v.as_bool()), Some(true));
995        assert_eq!(
996            value.get("background").and_then(|v| v.as_bool()),
997            Some(true)
998        );
999        assert_eq!(
1000            value.get("still_running").and_then(|v| v.as_bool()),
1001            Some(true)
1002        );
1003        assert_eq!(value.get("state").and_then(|v| v.as_str()), Some("running"));
1004        assert_eq!(
1005            value.get("tail_lines_used").and_then(|v| v.as_u64()),
1006            Some(50)
1007        );
1008        assert_eq!(
1009            value.get("elapsed_time").and_then(|v| v.as_str()),
1010            Some("00:01")
1011        );
1012        assert_eq!(
1013            value.get("log_tail").and_then(|v| v.as_str()),
1014            Some("tail line")
1015        );
1016
1017        let hint = value
1018            .get("hint")
1019            .and_then(|v| v.as_str())
1020            .expect("expected hint field");
1021
1022        // Hint should contain the actual job_id value
1023        assert!(
1024            hint.contains("job_id=job-42"),
1025            "hint should contain the actual job_id value; got: '{hint}'"
1026        );
1027        // Hint should mention the check_process tool
1028        assert!(
1029            hint.contains("check_process"),
1030            "hint should mention check_process tool; got: '{hint}'"
1031        );
1032        // Hint should warn against restarting
1033        assert!(
1034            hint.contains("DO NOT restart"),
1035            "hint should warn against restarting; got: '{hint}'"
1036        );
1037        // Hint should use TIMEOUT_RECOVERY prefix
1038        assert!(
1039            hint.contains("TIMEOUT_RECOVERY"),
1040            "hint should start with TIMEOUT_RECOVERY; got: '{hint}'"
1041        );
1042        assert!(
1043            hint.contains("MCP client deadlines may be shorter than timeout_ms"),
1044            "hint should distinguish the client deadline from timeout_ms; got: '{hint}'"
1045        );
1046        assert!(
1047            hint.contains("background=true"),
1048            "hint should recommend explicit background mode; got: '{hint}'"
1049        );
1050        // Hint should NOT contain old placeholders
1051        assert!(
1052            !hint.contains("<pid>"),
1053            "hint should not contain <pid> placeholder; got: '{hint}'"
1054        );
1055        assert!(
1056            !hint.contains("<log_path>"),
1057            "hint should not contain <log_path> placeholder; got: '{hint}'"
1058        );
1059    }
1060}