1use 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#[derive(Clone)]
98pub struct SshMcpServer {
99 config: Config,
101
102 connection: Arc<SshConnectionManager>,
104
105 timeout: Duration,
107
108 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 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 let mut ssh_config = SshConfig::new(&config.host, &config.user).with_port(config.port);
140
141 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 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 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 ssh_config = ssh_config
165 .with_keepalive_interval(config.keepalive_interval)
166 .with_keepalive_max(config.keepalive_max);
167
168 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 ssh_config = ssh_config
176 .with_host_key_checking(config.strict_host_key_checking)
177 .with_known_hosts(config.known_hosts.clone());
178
179 ssh_config = ssh_config.with_max_output_tokens(config.max_output_tokens);
181
182 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 pub fn connection(&self) -> &Arc<SshConnectionManager> {
311 &self.connection
312 }
313
314 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 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 let sanitized = match self.sanitize_or_tool_error(command) {
361 Ok(cmd) => cmd,
362 Err(result) => return Ok(result),
363 };
364
365 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 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 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 let sanitized = match self.sanitize_or_tool_error(command) {
444 Ok(cmd) => cmd,
445 Err(result) => return Ok(result),
446 };
447
448 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 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 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 fn exec_tool() -> Tool {
547 tools::exec_tool()
548 }
549
550 fn sudo_exec_tool() -> Tool {
552 tools::sudo_exec_tool()
553 }
554
555 fn transfer_tool() -> Tool {
557 tools::transfer_tool()
558 }
559
560 fn check_process_tool() -> Tool {
562 tools::check_process_tool()
563 }
564
565 fn read_file_tool() -> Tool {
567 tools::read_file_tool()
568 }
569
570 fn write_file_tool() -> Tool {
572 tools::write_file_tool()
573 }
574
575 fn replace_in_file_tool() -> Tool {
577 tools::replace_in_file_tool()
578 }
579
580 pub fn get_tool_documentation(tool_name: &str) -> Option<&'static str> {
585 tools::get_tool_documentation(tool_name)
586 }
587
588 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 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 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 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 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 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 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 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 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 #[test]
802 fn test_server_info() {
803 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 assert!(
1124 hint.contains("job_id=job-42"),
1125 "hint should contain the actual job_id value; got: '{hint}'"
1126 );
1127 assert!(
1129 hint.contains("check-process"),
1130 "hint should mention check-process tool; got: '{hint}'"
1131 );
1132 assert!(
1134 hint.contains("DO NOT restart"),
1135 "hint should warn against restarting; got: '{hint}'"
1136 );
1137 assert!(
1139 hint.contains("TIMEOUT_RECOVERY"),
1140 "hint should start with TIMEOUT_RECOVERY; got: '{hint}'"
1141 );
1142 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 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 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 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}