yexuanyang commented on code in PR #3547: URL: https://github.com/apache/hertzbeat/pull/3547#discussion_r2283956856
########## mcp-servers/mcp-bash-server/src/common/bash_server.rs: ########## @@ -0,0 +1,1203 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Bash Server Implementation for MCP (Model Context Protocol) +//! +//! This module provides a bash command execution server that can run shell commands +//! safely with validation and timeout controls. It supports multiple operating systems +//! and provides various execution methods for different use cases. + +#![allow(unused_imports, unused_variables, dead_code)] +use rmcp::model::Content; +use rmcp::{ + RoleServer, ServerHandler, + handler::server::{ + router::tool::ToolRouter, + tool::{IntoCallToolResult, Parameters}, + }, + model::*, + schemars, + serde_json::Value, + service::RequestContext, + tool, +}; +use rmcp::{serde_json, tool_handler, tool_router}; +use serde::{Deserialize, Serialize}; +use std::ffi::OsStr; +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; + +use std::process::Output; +use std::{ + borrow::Cow, + env, + fs::{self, File}, + io::Write, + process::Command, +}; +use tracing::error; +use tracing::info; +use uuid::Uuid; + +use crate::common::config::Config; +use crate::common::validator::Validator; + +/// Request structure for executing bash commands +/// Contains all necessary parameters for command execution including +/// working directory, environment variables, and timeout settings +#[derive(Debug, Deserialize, schemars::JsonSchema)] +pub struct DefaultExecuteRequest { + #[schemars(description = "The bash command or script to execute")] + pub command: String, + #[schemars(description = "Working directory for the command (optional)")] + pub working_dir: Option<String>, + #[schemars(description = "Environment variables (optional)")] + pub env_vars: Option<std::collections::HashMap<String, String>>, + #[schemars(description = "Timeout in seconds (default: 30)")] + pub timeout_seconds: Option<u64>, +} + +/// Response structure for command execution results +/// Contains stdout, stderr, exit code, success status and any parsed data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DefaultExecuteResponse { + /// Standard output from the executed command + pub stdout: String, + /// Standard error output from the executed command + pub stderr: String, + /// Exit code returned by the command (0 for success, non-zero for failure) + pub exit_code: i32, + /// Whether the command executed successfully (exit code 0) + pub success: bool, + /// Additional parsed data from command output (JSON format) + pub parsed_data: serde_json::Value, +} + +impl IntoCallToolResult for DefaultExecuteResponse { + fn into_call_tool_result(self) -> Result<CallToolResult, ErrorData> { + let content = if self.success { + format!( + "Command executed successfully (exit code: {})\n\nSTDOUT:\n{}\n\nSTDERR:\n{}", + self.exit_code, self.stdout, self.stderr + ) + } else { + format!( + "Command failed (exit code: {})\n\nSTDOUT:\n{}\n\nSTDERR:\n{}", + self.exit_code, self.stdout, self.stderr + ) + }; + + Ok(CallToolResult { + content: Some(vec![Content::text(content)]), + structured_content: None, + is_error: Some(!self.success), + }) + } +} + +/// Main bash server implementation that handles command execution +/// with optional validation and timeout controls +#[derive(Debug, Clone)] +pub struct BashServer { + validator: Option<Validator>, + tool_router: ToolRouter<BashServer>, +} + +/// Trait for command execution utilities +/// Provides common functionality for command formatting and execution +pub trait CommandRunner { + /// Convert a Command instance to a readable string representation + /// Handles proper quoting of arguments containing spaces + fn stringify_command(cmd: &Command) -> String { + let program = cmd.get_program().to_string_lossy(); + let args = cmd + .get_args() + .map(|arg| { + let s = arg.to_string_lossy(); + if s.contains(' ') || s.contains('"') { + format!("{s:?}") + } else { + s.to_string() + } + }) + .collect::<Vec<String>>() + .join(" "); + + format!("{program} {args}") + } + + /// Execute a command with timeout protection + /// Returns command output or timeout error + async fn execute_command_with_timeout( + timeout: std::time::Duration, + mut cmd: Command, + ) -> Result<Output, ErrorData> { + let cmd_str = Self::stringify_command(&cmd); + // Execute command with timeout + let output = tokio::time::timeout(timeout, async { + tokio::task::spawn_blocking(move || cmd.output()).await + }) + .await + .map_err(|_| ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: Cow::Owned("Command execution timed out".to_string()), + data: None, + })? + .map_err(|e| ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: Cow::Owned(format!("Failed to spawn command: {e}")), + data: None, + })? + .map_err(|e| ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: Cow::Owned(format!("Command execution failed: {e}")), + data: None, + })?; + + // log execution of command + info!("Execute command: {cmd_str}"); + Ok(output) + } +} + +impl CommandRunner for BashServer {} + +impl BashServer { + /// Create a new BashServer instance + /// Attempts to load configuration from config.toml, creates validator if successful + pub fn new() -> Self { + let tool_router = Self::tool_router(); + if let Ok(config) = Config::read_config("config.toml") + .inspect_err(|e| eprintln!("read config.toml fail, error: {e}")) + { + let blacklist = config.blacklist; + let whitelist = config.whitelist; + BashServer { + validator: Some(Validator::new(blacklist, whitelist)), + tool_router, + } + } else { + Self { + validator: None, + tool_router, + } + } + } + + /// Internal method for executing commands via default shell + /// Supports validation toggle and handles cross-platform shell differences + async fn _all_execute_via_default_shell( + &self, + need_validate: bool, + request: DefaultExecuteRequest, + ) -> Result<CallToolResult, ErrorData> { + let timeout_duration = + std::time::Duration::from_secs(request.timeout_seconds.unwrap_or(30)); + + let mut cmd = if cfg!(target_os = "linux") { + let mut cmd = Command::new("bash"); + cmd.arg("-c"); + cmd + } else if cfg!(target_os = "windows") { + let mut cmd = Command::new("powershell"); + cmd.arg("-c"); + cmd + } else { + let mut cmd = Command::new("sh"); + cmd.arg("-c"); + cmd + }; + cmd.arg(&request.command); + + // Set working directory if provided + if let Some(working_dir) = &request.working_dir { + cmd.current_dir(working_dir); + } + + // Set environment variables if provided + if let Some(env_vars) = &request.env_vars { + for (key, value) in env_vars { + cmd.env(key, value); + } + } + + // Validate the commands + if let Some(validator) = &self.validator Review Comment: Do you mean validate the request.command instead of full_args? For example, validate the `echo 'Hello Test'` instead of `powershell -c "echo 'Hello Test'"`. ########## mcp-servers/mcp-bash-server/src/common/bash_server.rs: ########## @@ -0,0 +1,1203 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Bash Server Implementation for MCP (Model Context Protocol) +//! +//! This module provides a bash command execution server that can run shell commands +//! safely with validation and timeout controls. It supports multiple operating systems +//! and provides various execution methods for different use cases. + +#![allow(unused_imports, unused_variables, dead_code)] +use rmcp::model::Content; +use rmcp::{ + RoleServer, ServerHandler, + handler::server::{ + router::tool::ToolRouter, + tool::{IntoCallToolResult, Parameters}, + }, + model::*, + schemars, + serde_json::Value, + service::RequestContext, + tool, +}; +use rmcp::{serde_json, tool_handler, tool_router}; +use serde::{Deserialize, Serialize}; +use std::ffi::OsStr; +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; + +use std::process::Output; +use std::{ + borrow::Cow, + env, + fs::{self, File}, + io::Write, + process::Command, +}; +use tracing::error; +use tracing::info; +use uuid::Uuid; + +use crate::common::config::Config; +use crate::common::validator::Validator; + +/// Request structure for executing bash commands +/// Contains all necessary parameters for command execution including +/// working directory, environment variables, and timeout settings +#[derive(Debug, Deserialize, schemars::JsonSchema)] +pub struct DefaultExecuteRequest { + #[schemars(description = "The bash command or script to execute")] + pub command: String, + #[schemars(description = "Working directory for the command (optional)")] + pub working_dir: Option<String>, + #[schemars(description = "Environment variables (optional)")] + pub env_vars: Option<std::collections::HashMap<String, String>>, + #[schemars(description = "Timeout in seconds (default: 30)")] + pub timeout_seconds: Option<u64>, +} + +/// Response structure for command execution results +/// Contains stdout, stderr, exit code, success status and any parsed data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DefaultExecuteResponse { + /// Standard output from the executed command + pub stdout: String, + /// Standard error output from the executed command + pub stderr: String, + /// Exit code returned by the command (0 for success, non-zero for failure) + pub exit_code: i32, + /// Whether the command executed successfully (exit code 0) + pub success: bool, + /// Additional parsed data from command output (JSON format) + pub parsed_data: serde_json::Value, +} + +impl IntoCallToolResult for DefaultExecuteResponse { + fn into_call_tool_result(self) -> Result<CallToolResult, ErrorData> { + let content = if self.success { + format!( + "Command executed successfully (exit code: {})\n\nSTDOUT:\n{}\n\nSTDERR:\n{}", + self.exit_code, self.stdout, self.stderr + ) + } else { + format!( + "Command failed (exit code: {})\n\nSTDOUT:\n{}\n\nSTDERR:\n{}", + self.exit_code, self.stdout, self.stderr + ) + }; + + Ok(CallToolResult { + content: Some(vec![Content::text(content)]), + structured_content: None, + is_error: Some(!self.success), + }) + } +} + +/// Main bash server implementation that handles command execution +/// with optional validation and timeout controls +#[derive(Debug, Clone)] +pub struct BashServer { + validator: Option<Validator>, + tool_router: ToolRouter<BashServer>, +} + +/// Trait for command execution utilities +/// Provides common functionality for command formatting and execution +pub trait CommandRunner { + /// Convert a Command instance to a readable string representation + /// Handles proper quoting of arguments containing spaces + fn stringify_command(cmd: &Command) -> String { + let program = cmd.get_program().to_string_lossy(); + let args = cmd + .get_args() + .map(|arg| { + let s = arg.to_string_lossy(); + if s.contains(' ') || s.contains('"') { + format!("{s:?}") + } else { + s.to_string() + } + }) + .collect::<Vec<String>>() + .join(" "); + + format!("{program} {args}") + } + + /// Execute a command with timeout protection + /// Returns command output or timeout error + async fn execute_command_with_timeout( + timeout: std::time::Duration, + mut cmd: Command, + ) -> Result<Output, ErrorData> { + let cmd_str = Self::stringify_command(&cmd); + // Execute command with timeout + let output = tokio::time::timeout(timeout, async { + tokio::task::spawn_blocking(move || cmd.output()).await + }) + .await + .map_err(|_| ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: Cow::Owned("Command execution timed out".to_string()), + data: None, + })? + .map_err(|e| ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: Cow::Owned(format!("Failed to spawn command: {e}")), + data: None, + })? + .map_err(|e| ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: Cow::Owned(format!("Command execution failed: {e}")), + data: None, + })?; + + // log execution of command + info!("Execute command: {cmd_str}"); + Ok(output) + } +} + +impl CommandRunner for BashServer {} + +impl BashServer { + /// Create a new BashServer instance + /// Attempts to load configuration from config.toml, creates validator if successful + pub fn new() -> Self { + let tool_router = Self::tool_router(); + if let Ok(config) = Config::read_config("config.toml") + .inspect_err(|e| eprintln!("read config.toml fail, error: {e}")) + { + let blacklist = config.blacklist; + let whitelist = config.whitelist; + BashServer { + validator: Some(Validator::new(blacklist, whitelist)), + tool_router, + } + } else { + Self { + validator: None, + tool_router, + } + } + } + + /// Internal method for executing commands via default shell + /// Supports validation toggle and handles cross-platform shell differences + async fn _all_execute_via_default_shell( + &self, + need_validate: bool, + request: DefaultExecuteRequest, + ) -> Result<CallToolResult, ErrorData> { + let timeout_duration = + std::time::Duration::from_secs(request.timeout_seconds.unwrap_or(30)); + + let mut cmd = if cfg!(target_os = "linux") { + let mut cmd = Command::new("bash"); + cmd.arg("-c"); + cmd + } else if cfg!(target_os = "windows") { + let mut cmd = Command::new("powershell"); + cmd.arg("-c"); + cmd + } else { + let mut cmd = Command::new("sh"); + cmd.arg("-c"); + cmd + }; + cmd.arg(&request.command); + + // Set working directory if provided + if let Some(working_dir) = &request.working_dir { + cmd.current_dir(working_dir); + } + + // Set environment variables if provided + if let Some(env_vars) = &request.env_vars { + for (key, value) in env_vars { + cmd.env(key, value); + } + } + + // Validate the commands + if let Some(validator) = &self.validator + && need_validate + { + let program = cmd.get_program(); + let mut full_args: Vec<&OsStr> = vec![program]; + let args: Vec<&OsStr> = cmd.get_args().collect(); + full_args.extend(args); + validator.is_unsafe_command(full_args)?; + } + Review Comment: ```suggestion if let Some(validator) = &self.validator && need_validate { validator.is_unsafe_command(&request.command)?; } ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: notifications-unsubscr...@hertzbeat.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: notifications-unsubscr...@hertzbeat.apache.org For additional commands, e-mail: notifications-h...@hertzbeat.apache.org
