chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,15 @@
|
||||
# Introduction
|
||||
This file has a short description of each mcp server in `mcp-servers` directory
|
||||
|
||||
# MCP servers
|
||||
|
||||
## mcp-bash-server
|
||||
|
||||
This mcp server provide the ability to execute shell command or ASCII text executable
|
||||
script in the machine running this server.
|
||||
|
||||
Use a simple Oauth2.0 for authentication.
|
||||
|
||||
## mcp-log-server
|
||||
|
||||
This mcp server provide the ability to query log from greptime db.
|
||||
@@ -0,0 +1,2 @@
|
||||
/target
|
||||
logs/
|
||||
Generated
+2813
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,56 @@
|
||||
# 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.
|
||||
|
||||
[package]
|
||||
name = "mcp-bash-server"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
tracing-appender = "0.2"
|
||||
anyhow = "1.0.98"
|
||||
rmcp = { version = "0.4.0", features = ["server", "transport-sse-server", "transport-io", "transport-streamable-http-server", "auth"]}
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
tokio = { version = "1.0", features = ["full"] }
|
||||
tokio-stream = { version = "0.1", features = ["net"] }
|
||||
tokio-util = { version = "0.7.15", features = ["full"] }
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = [
|
||||
"env-filter",
|
||||
"std",
|
||||
"fmt",
|
||||
] }
|
||||
axum = { version = "0.8", features = ["macros"] }
|
||||
chrono = "0.4"
|
||||
tower-http = { version = "0.6", features = ["cors"] }
|
||||
askama = { version = "0.14"}
|
||||
rand = { version = "0.8", features = ["std"] }
|
||||
uuid = { version = "1.6", features = ["v4", "serde"] }
|
||||
serde_urlencoded = "0.7"
|
||||
oauth2 = "5.0"
|
||||
toml = "0.8"
|
||||
regex = "1.11.1"
|
||||
shlex = "1.3"
|
||||
|
||||
[dev-dependencies]
|
||||
tokio-test = "0.4"
|
||||
tempfile = "3.0"
|
||||
tower = "0.5"
|
||||
hyper = "1.0"
|
||||
axum-test = "14.1"
|
||||
|
||||
[[bin]]
|
||||
name = "mcp-bash-server"
|
||||
path = "src/main.rs"
|
||||
@@ -0,0 +1,97 @@
|
||||
# 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.
|
||||
|
||||
# Multi-stage build for mcp-bash-server
|
||||
# Stage 1: Build stage
|
||||
FROM ubuntu:noble AS builder
|
||||
|
||||
# Set environment variables
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV RUST_VERSION=stable
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
# Set proxy environment variables via build arguments
|
||||
ARG HTTPS_PROXY
|
||||
ARG HTTP_PROXY
|
||||
ENV https_proxy=${HTTPS_PROXY}
|
||||
ENV http_proxy=${HTTP_PROXY}
|
||||
|
||||
# Install build dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
build-essential \
|
||||
pkg-config \
|
||||
libssl-dev \
|
||||
ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Rust
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain ${RUST_VERSION}
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /app
|
||||
|
||||
# Copy the mcp-bash-server project
|
||||
COPY mcp-bash-server/ .
|
||||
|
||||
# Build the project in release mode
|
||||
RUN cargo build --release
|
||||
|
||||
# Stage 2: Runtime stage
|
||||
FROM ubuntu:noble AS runtime
|
||||
|
||||
# Set environment variables
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install runtime dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
ca-certificates \
|
||||
libssl3 \
|
||||
bash \
|
||||
coreutils \
|
||||
procps \
|
||||
curl \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Create a non-root user for security
|
||||
RUN useradd -m -s /bin/bash mcpuser
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /app
|
||||
|
||||
# Copy the compiled binary from builder stage
|
||||
COPY --from=builder /app/target/release/mcp-bash-server /app/mcp-bash-server
|
||||
|
||||
# Copy configuration and templates
|
||||
COPY --from=builder /app/config.toml /app/config.toml
|
||||
COPY --from=builder /app/templates/ /app/templates/
|
||||
|
||||
# Create logs directory and set permissions
|
||||
RUN mkdir -p logs && chown -R mcpuser:mcpuser /app
|
||||
|
||||
# Switch to non-root user
|
||||
USER mcpuser
|
||||
|
||||
# Expose the port
|
||||
EXPOSE 4000
|
||||
|
||||
# Health check
|
||||
HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
|
||||
CMD curl -f http://localhost:4000/ || exit 1
|
||||
|
||||
# Set the entrypoint to run the mcp-bash-server
|
||||
CMD ["./mcp-bash-server"]
|
||||
@@ -0,0 +1,333 @@
|
||||
# mcp-bash-server
|
||||
|
||||
A HertzBeat MCP server for running scripts with security command blacklist and logging capabilities
|
||||
|
||||
## Dependencies
|
||||
|
||||
### Rust
|
||||
|
||||
If you need to deploy this MCP Server locally, you will need a Rust runtime environment.
|
||||
|
||||
Visit [rust-lang.org](https://www.rust-lang.org/tools/install) to learn how to install the Rust runtime environment.
|
||||
|
||||
Version `1.88.0` can absolutely work, and we recommend using the latest version of Rust.
|
||||
|
||||
## Deployment
|
||||
|
||||
### Local Deployment
|
||||
|
||||
If you want to run this MCP server locally using the default settings provided by the project, simply run the following command in the project root directory:
|
||||
|
||||
```Rust
|
||||
cargo run
|
||||
```
|
||||
|
||||
This MCP server will be deployed at `http://127.0.0.1:4000/mcp`, and you can use the `modelcontextprotocol/inspector` tool to connect to and use this MCP server.
|
||||
|
||||
For information on how to use the modelcontextprotocol/inspector tool, refer to the [inspector documentation](https://github.com/modelcontextprotocol/inspector).
|
||||
|
||||
If you encounter any issues while using Inspector, it is recommended to use version `v0.16.2`. Other versions may also work.
|
||||
|
||||
### Container Deployment
|
||||
|
||||
Using container deployment for this MCP Server is an excellent way to try out the tools provided by the server without polluting your machine, as all MCP Server operations are completed within the container.
|
||||
|
||||
Refer to the `Dockerfile` in the code repository and create your own Dockerfile. After creating it, run the following build command:
|
||||
|
||||
```shell
|
||||
docker build -t apache/hertzbeat-mcp-bash-server:latest .
|
||||
```
|
||||
|
||||
You can use the proxy in docker build
|
||||
|
||||
```shell
|
||||
docker build --build-arg HTTPS_PROXY=<your https_proxy> --build-arg HTTP_PROXY=<your http_proxy> -t apache/hertzbeat-mcp-bash-server:latest .
|
||||
```
|
||||
|
||||
After building, use the following command to run it:
|
||||
|
||||
```shell
|
||||
docker run -d --name mcp-bash-server -p 4000:4000 --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest
|
||||
```
|
||||
|
||||
The MCP Server inside the container runs on 0.0.0.0:4000. On the host machine, use the inspector with URL `http://localhost:4000/mcp` to connect to the MCP Server inside the container.
|
||||
|
||||
#### Use custom config in container
|
||||
|
||||
Container's workdir is `/app` and it will run the `/app/mcp-bash-server` when it start, this program will read the `config.toml` at the same directory, so you can put the `config.toml` in the `/app` directory to cover the default config in image. Use the command below to do it.
|
||||
|
||||
```shell
|
||||
docker run -d --name mcp-bash-server -p 4000:4000 -v `pwd`/config.toml:/app/config.toml --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest
|
||||
```
|
||||
|
||||
If you are using SELinux, you may need to run the command instead to let the container access the file in host.
|
||||
|
||||
```shell
|
||||
docker run -d --name mcp-bash-server -p 4000:4000 -v `pwd`/config.toml:/app/config.toml:Z --restart unless-stopped apache/hertzbeat-mcp-bash-server:latest
|
||||
```
|
||||
|
||||
To check if the config.toml is used, do this
|
||||
|
||||
```shell
|
||||
docker logs mcp-bash-server
|
||||
```
|
||||
|
||||
or check the config.toml in container
|
||||
|
||||
```shell
|
||||
docker exec mcp-bash-server ls /app
|
||||
docker exec mcp-bash-server cat /app/config.toml
|
||||
```
|
||||
|
||||
## Use it by Agents
|
||||
|
||||
### Vscode Copilot
|
||||
|
||||
Start the MCP Server in daemon mode, then add the settings to your Vscode Copilot mcp config. Note that you can use modelcontextprotocol/inspector to finish the OAuth flow and get the access token, put it in config and restart server, then Copilot can connect to the server and use the tools.
|
||||
|
||||
```json
|
||||
{
|
||||
"servers": {
|
||||
"bash-server": {
|
||||
"url": "http://localhost:4000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer <your-token>"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Currently Vscode MCP OAuth can not automatically authorize this bash-server**
|
||||
The vscode mcp OAuth flow is:
|
||||
|
||||
1. GET /.well-known/oauth-authorization-server
|
||||
2. GET /authorize with query-params
|
||||
3. ...
|
||||
|
||||
But we requires the client registration before accessing endpoint `/authorize` with query-params that contains invalid client-id. So we can only set the token manually now.
|
||||
|
||||
## Configuration
|
||||
|
||||
The `config.toml` file contains the settings for this MCP server. The configuration items currently supported by the configuration file are as follows:
|
||||
|
||||
```toml
|
||||
# This is the configuration file for mcp-bash-server
|
||||
|
||||
[settings]
|
||||
# Port for MCP Server deployment
|
||||
port = 4000
|
||||
# IP for MCP Server deployment
|
||||
host = "127.0.0.1"
|
||||
# Usage environment for MCP Server, can be either development or production environment.
|
||||
# Set env to "development" or "production"; production uses OAuth 2.0.
|
||||
env = "development"
|
||||
|
||||
[whitelist]
|
||||
# Whitelist of allowed commands (exact match), a string list.
|
||||
# Only commands where the complete command string exactly matches an item in this list will be allowed
|
||||
commands = [
|
||||
"echo hello",
|
||||
"ls -la",
|
||||
"pwd",
|
||||
# Add your allowed commands here
|
||||
]
|
||||
# Whitelist of allowed command regex patterns, a list of regex expression strings
|
||||
# Commands where the complete command string matches any of these regex patterns will be allowed
|
||||
regex = [
|
||||
'^echo [a-zA-Z0-9 ]+$',
|
||||
'^ls [a-zA-Z0-9 /-]*$',
|
||||
# Add your whitelist regex patterns here
|
||||
]
|
||||
|
||||
[blacklist]
|
||||
# Blacklist of forbidden commands (exact match), a string list
|
||||
# Blacklist has higher priority than whitelist. If the complete command string exactly matches an item in this list, the entire command will be blocked, even if it would be allowed by the whitelist
|
||||
# Note: Only exact matches are blocked. For example, if "rm" is blacklisted, only the exact command "rm" is blocked, not commands like "rm -rf /tmp/test"
|
||||
commands = [
|
||||
# Dangerous file operations
|
||||
"rm -rf /",
|
||||
"shutdown",
|
||||
# Add your forbidden commands here
|
||||
]
|
||||
# Blacklist of forbidden command regex patterns, a list of regex expression strings
|
||||
# Blacklist has higher priority than whitelist. If the complete command string matches any of these regex patterns, it will be blocked, even if it would be allowed by the whitelist
|
||||
regex = [
|
||||
# Block any command with dangerous operators
|
||||
'.*[|&;`$()><].*',
|
||||
# Block commands that try to write to system directories
|
||||
'.*/etc/.*',
|
||||
'.*/root/.*',
|
||||
# Block commands with sudo or su
|
||||
'^sudo .*',
|
||||
'^su .*',
|
||||
# Add your blacklist regex patterns here
|
||||
]
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
### Unit Tests
|
||||
|
||||
Run `cargo test` to execute all unit tests.
|
||||
|
||||
The unit tests include:
|
||||
|
||||
#### Configuration Tests (`config.rs`)
|
||||
|
||||
- **Config Parsing**: Tests parsing of TOML configuration files with various settings
|
||||
- **Invalid File Handling**: Tests error handling for non-existent configuration files
|
||||
- **Invalid TOML Handling**: Tests error handling for malformed TOML syntax
|
||||
- **Whitelist Creation**: Tests creation and validation of command whitelists
|
||||
- **Blacklist Creation**: Tests creation and validation of command blacklists
|
||||
- **Settings Creation**: Tests creation of server settings with port, host, and environment configurations
|
||||
|
||||
#### Validator Tests (`validator.rs`)
|
||||
|
||||
- **Validator Creation**: Tests creation of command validators with blacklist and whitelist rules (both commands and regex)
|
||||
- **Whitelisted Command Validation**: Tests allowing commands that exactly match whitelist entries
|
||||
- **Whitelisted Regex Validation**: Tests allowing commands that match whitelist regex patterns
|
||||
- **Blacklisted Command Blocking**: Tests blocking of commands that exactly match blacklist entries (complete command string)
|
||||
- **Blacklisted Regex Blocking**: Tests blocking of commands matching blacklist regex patterns
|
||||
- **Priority Testing**: Tests that blacklist has higher priority than whitelist
|
||||
- **Exact vs Partial Match Testing**: Tests that blacklist only blocks exact command matches, not partial matches
|
||||
- **Default Deny Behavior**: Tests that commands not in whitelist are denied by default
|
||||
- **Empty Command Handling**: Tests that empty commands are denied by default
|
||||
- **Case Sensitivity Testing**: Tests case-sensitive command matching behavior
|
||||
- **Invalid Regex Handling**: Tests error handling for invalid regex patterns
|
||||
- **Complex Command Testing**: Tests validation of multi-argument commands
|
||||
- **Nested Shell Command Testing**: Tests validation of commands within shell invocations
|
||||
|
||||
#### Bash Server Tests (`bash_server.rs`)
|
||||
|
||||
- **Response Success/Failure**: Tests command execution response handling for both successful and failed commands
|
||||
- **Command Stringification**: Tests conversion of Command objects to readable string representations
|
||||
- **Command Execution with Timeout**: Tests command execution with timeout protection
|
||||
- **Server Creation**: Tests BashServer instantiation with and without configuration
|
||||
- **Request/Response Serialization**: Tests JSON serialization and deserialization of command requests and responses
|
||||
- **Environment Variable Handling**: Tests setting and using environment variables in commands
|
||||
- **Working Directory Handling**: Tests command execution in specific working directories
|
||||
- **Debug and Clone Traits**: Tests that BashServer properly implements Debug and Clone traits
|
||||
- **Timeout Behavior**: Tests that commands properly timeout when exceeding time limits
|
||||
- **Unix script and Python Execution**: Tests tool `execute_script` and `execute_python`
|
||||
|
||||
### Manual Testing with Inspector
|
||||
|
||||
Use the official MCP debugging tool modelcontextprotocol/inspector. For usage instructions, refer to the [inspector documentation](https://github.com/modelcontextprotocol/inspector).
|
||||
|
||||
If you deploy the MCP Server locally using the default method, run the inspector debugging tool after the server starts.
|
||||
|
||||
```shell
|
||||
npx @modelcontextprotocol/inspector
|
||||
```
|
||||
|
||||
Then set the connection method to `Streamable HTTP` and set the URL to `http://127.0.0.1:4000/mcp`.
|
||||
|
||||
### Using OAuth in Inspector (Optional)
|
||||
|
||||
If your running mode is `development`, you don't need to go through `Open Auth Settings` and can directly click Connect to connect to the MCP Server. If your running mode is `production`, you need to complete the `Open Auth Settings` verification.
|
||||
|
||||
The method for using OAuth verification and connection is as follows:
|
||||
|
||||
1. Click `Open Auth Settings`
|
||||
2. Click `Quick OAuth Flow`
|
||||
3. Click `Approve` on the pop-up webpage
|
||||
4. Return to the MCP inspector, click on the Access Tokens under `Authentication Complete` in `OAuth Flow Progress`. Copy the `access_token` from there
|
||||
5. Click `Authentication`, paste the previously copied token into the `Bearer Token` field, then click Connect
|
||||
|
||||
After connecting, you can test the various tools provided by the MCP Server in the inspector.
|
||||
|
||||
## Tools
|
||||
|
||||
The MCP Server currently provides the following tools:
|
||||
|
||||
### Command Execution Tools
|
||||
|
||||
1. **all_execute_via_default_shell**
|
||||
- Description: Execute commands using the default shell on all types of operating systems
|
||||
- Parameters:
|
||||
- `command`: The bash command or script to execute
|
||||
- `working_dir`: Working directory for the command (optional)
|
||||
- `env_vars`: Environment variables (optional)
|
||||
- `timeout_seconds`: Timeout in seconds (default: 30)
|
||||
- Returns: Execution result with stdout, stderr, exit code, and success status
|
||||
|
||||
2. **execute_python**
|
||||
- Description: Execute a Python script
|
||||
- Parameters: Same as above, but command should be Python code
|
||||
- Returns: Python script execution result
|
||||
|
||||
3. **execute_script**
|
||||
- Description: Execute a Unix script by writing it to a temporary file
|
||||
- Parameters: Same as above, command should be shell script content
|
||||
- Returns: Script execution result
|
||||
|
||||
### System Information Tools
|
||||
|
||||
1. **unix_get_system_info_via_default_shell**
|
||||
- Description: Get comprehensive system information including OS, kernel, memory, and disk usage
|
||||
- Parameters: None
|
||||
- Returns: Detailed system information including hostname, OS, kernel, architecture, uptime, memory usage, and disk usage
|
||||
|
||||
2. **unix_preset_get_system_info_via_default_shell**
|
||||
- Description: Get formatted system information in a structured format
|
||||
- Parameters: None
|
||||
- Returns: Version, hostname, and uptime in a single line format
|
||||
|
||||
3. **unix_get_available_shell**
|
||||
- Description: Get the available shells on Unix-like operating systems
|
||||
- Parameters: None
|
||||
- Returns: List of available shells from /etc/shells
|
||||
|
||||
### Performance Monitoring Tools
|
||||
|
||||
1. **unix_preset_get_nic_info_via_default_shell**
|
||||
- Description: Get network interface statistics including receive and transmit bytes
|
||||
- Parameters: None
|
||||
- Returns: Interface names and traffic data in tabular format
|
||||
|
||||
2. **unix_preset_get_cpu_info_via_default_shell**
|
||||
- Description: Get detailed CPU information including model, core count, load averages, and performance metrics
|
||||
- Parameters: None
|
||||
- Returns: CPU model, core count, load averages, idle percentage, and other performance metrics
|
||||
|
||||
3. **unix_preset_get_disk_free_info_via_default_shell**
|
||||
- Description: Get disk usage information for all mounted filesystems
|
||||
- Parameters: None
|
||||
- Returns: Filesystem usage data with columns for filesystem, used, available, usage percentage, and mount point
|
||||
|
||||
4. **unix_preset_get_top10_cpu_processes_via_default_shell**
|
||||
- Description: Get top 10 processes consuming the most CPU
|
||||
- Parameters: None
|
||||
- Returns: Process information sorted by CPU usage in descending order
|
||||
|
||||
5. **unix_preset_get_top10_mem_processes_via_default_shell**
|
||||
- Description: Get top 10 processes consuming the most memory
|
||||
- Parameters: None
|
||||
- Returns: Process information sorted by memory usage in descending order
|
||||
|
||||
### Security Features
|
||||
|
||||
All command execution tools support comprehensive security validation:
|
||||
|
||||
- **Default Deny Policy**: All commands are denied by default unless explicitly allowed
|
||||
- **Blacklist Priority**: Blacklist rules have higher priority than whitelist rules
|
||||
- **Dual Validation Modes**: Both exact string matching and regex pattern matching for blacklist and whitelist
|
||||
- **Command Validation Logic**:
|
||||
1. Empty commands are denied by default
|
||||
2. If the **complete command string** exactly matches any blacklist command entry, deny immediately
|
||||
3. If the **complete command string** matches any blacklist regex pattern, deny immediately
|
||||
4. If command doesn't match blacklist and the **complete command string** exactly matches any whitelist command entry, allow
|
||||
5. If command doesn't match blacklist and the **complete command string** matches any whitelist regex pattern, allow
|
||||
6. Otherwise, deny by default
|
||||
- **Important**: Blacklist validation checks the **entire command string**, not individual arguments. For example, if `"rm"` is blacklisted, only the exact command `"rm"` is blocked, not commands like `"rm -rf /tmp/test"` (which would be evaluated separately against the whitelist)
|
||||
- **Timeout Protection**: Commands have configurable timeout limits to prevent hanging
|
||||
- **Environment Isolation**: Optional working directory and environment variable settings
|
||||
- **Logging**: All command executions and validation results are logged for audit purposes
|
||||
- **Nested Shell Protection**: Validates commands within shell invocations (e.g., `bash -c "command"`)
|
||||
|
||||
### Tool Categories
|
||||
|
||||
- **General Execution**: Tools 1-3 for executing arbitrary commands, Python scripts, and shell scripts
|
||||
- **System Monitoring**: Tools 1-3 for gathering basic system information
|
||||
- **Performance Analysis**: Tools 1-5 for detailed performance monitoring and process analysis
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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.
|
||||
|
||||
# This is the configuration file for mcp-bash-server
|
||||
|
||||
[settings]
|
||||
# Port for MCP Server deployment
|
||||
port = 4000
|
||||
# IP for MCP Server deployment
|
||||
host = "127.0.0.1"
|
||||
# Usage environment for MCP Server, can be development or production environment.
|
||||
# Set env to "development" or "production", production uses oauth2.0.
|
||||
env = "production"
|
||||
|
||||
[whitelist]
|
||||
# Whitelist of allowed commands (exact match), a string list.
|
||||
# Only commands where the complete command string exactly matches an item in this list will be allowed
|
||||
commands = [
|
||||
"echo hello",
|
||||
"ls -la",
|
||||
"pwd",
|
||||
# Add your allowed commands here
|
||||
]
|
||||
# Whitelist of allowed command regex patterns, a regex expression string list
|
||||
# Commands where the complete command string matches any of these regex patterns will be allowed
|
||||
regex = [
|
||||
'^echo [a-zA-Z0-9 ]+$',
|
||||
'^ls [a-zA-Z0-9 /-]*$',
|
||||
# Add your whitelist regex patterns here
|
||||
]
|
||||
|
||||
[blacklist]
|
||||
# Blacklist of forbidden commands (exact match), a string list
|
||||
# Blacklist has higher priority than whitelist. If the complete command string exactly matches an item in this list, the entire command will be blocked, even if it would be allowed by the whitelist
|
||||
# Note: Only exact matches are blocked. For example, if "rm" is blacklisted, only the exact command "rm" is blocked, not commands like "rm -rf /tmp/test"
|
||||
commands = [
|
||||
# Dangerous file operations
|
||||
"rm -rf /",
|
||||
"shutdown",
|
||||
# Add your forbidden commands here
|
||||
]
|
||||
# Blacklist of forbidden command regex patterns, a regex expression string list
|
||||
# Blacklist has higher priority than whitelist. If the complete command string matches any of these regex patterns, it will be blocked, even if it would be allowed by the whitelist
|
||||
regex = [
|
||||
# Block any command with dangerous operators
|
||||
'.*[|&;`$()><].*',
|
||||
# Block commands that try to write to system directories
|
||||
'.*/etc/.*',
|
||||
'.*/root/.*',
|
||||
# Block commands with sudo or su
|
||||
'^sudo .*',
|
||||
'^su .*',
|
||||
# Add your blacklist regex patterns here
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,190 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
//! Configuration management for the MCP Bash Server
|
||||
//!
|
||||
//! This module handles reading and parsing configuration from TOML files,
|
||||
//! including server settings and security blacklists for command validation.
|
||||
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct Config {
|
||||
pub settings: Settings,
|
||||
pub blacklist: Blacklist,
|
||||
pub whitelist: Whitelist,
|
||||
}
|
||||
|
||||
/// Security whitelist configuration for command validation
|
||||
/// Contains lists of allowed commands and operations
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Whitelist {
|
||||
/// List of command names that are allowed to be executed
|
||||
pub commands: Vec<String>,
|
||||
/// List of regular expressions for commands that are allowed
|
||||
/// These patterns are matched against the full command line
|
||||
pub regex: Vec<String>,
|
||||
}
|
||||
|
||||
/// Security blacklist configuration for command validation
|
||||
/// Contains lists of forbidden commands and regex patterns
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Blacklist {
|
||||
/// List of command names that are not allowed to be executed
|
||||
pub commands: Vec<String>,
|
||||
/// List of regular expressions for commands that are not allowed
|
||||
/// These patterns are matched against the full command line
|
||||
pub regex: Vec<String>,
|
||||
}
|
||||
|
||||
/// Server runtime settings including network configuration
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct Settings {
|
||||
pub port: u16,
|
||||
pub host: String,
|
||||
pub env: Option<String>, // "development" or "production"
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Read and parse configuration from a TOML file
|
||||
/// Returns parsed Config structure or error if file cannot be read/parsed
|
||||
pub fn read_config(file: &str) -> Result<Config> {
|
||||
let toml_str = fs::read_to_string(file)?;
|
||||
let config: Config = toml::from_str(&toml_str)?;
|
||||
Ok(config)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
#[test]
|
||||
fn test_config_parsing() {
|
||||
let config_content = r#"
|
||||
[settings]
|
||||
port = 4000
|
||||
host = "127.0.0.1"
|
||||
env = "development"
|
||||
|
||||
[whitelist]
|
||||
commands = ["echo hello", "ls -la"]
|
||||
regex = ["echo.*", "ls.*"]
|
||||
|
||||
[blacklist]
|
||||
commands = ["rm", "dd"]
|
||||
regex = [".*[|&].*", "^sudo .*"]
|
||||
"#;
|
||||
|
||||
let mut temp_file = NamedTempFile::new().unwrap();
|
||||
write!(temp_file, "{config_content}").unwrap();
|
||||
let temp_path = temp_file.path().to_str().unwrap();
|
||||
|
||||
let config = Config::read_config(temp_path).unwrap();
|
||||
|
||||
assert_eq!(config.settings.port, 4000);
|
||||
assert_eq!(config.settings.host, "127.0.0.1");
|
||||
assert_eq!(config.settings.env, Some("development".to_string()));
|
||||
|
||||
assert_eq!(config.whitelist.commands.len(), 2);
|
||||
assert!(
|
||||
config
|
||||
.whitelist
|
||||
.commands
|
||||
.contains(&"echo hello".to_string())
|
||||
);
|
||||
assert!(config.whitelist.commands.contains(&"ls -la".to_string()));
|
||||
|
||||
assert_eq!(config.whitelist.regex.len(), 2);
|
||||
assert!(config.whitelist.regex.contains(&"echo.*".to_string()));
|
||||
assert!(config.whitelist.regex.contains(&"ls.*".to_string()));
|
||||
|
||||
assert_eq!(config.blacklist.commands.len(), 2);
|
||||
assert!(config.blacklist.commands.contains(&"rm".to_string()));
|
||||
assert!(config.blacklist.commands.contains(&"dd".to_string()));
|
||||
|
||||
assert_eq!(config.blacklist.regex.len(), 2);
|
||||
assert!(config.blacklist.regex.contains(&".*[|&].*".to_string()));
|
||||
assert!(config.blacklist.regex.contains(&"^sudo .*".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_invalid_file() {
|
||||
let result = Config::read_config("non_existent_file.toml");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_invalid_toml() {
|
||||
let invalid_config = r#"
|
||||
[settings
|
||||
port = 4000
|
||||
"#;
|
||||
|
||||
let mut temp_file = NamedTempFile::new().unwrap();
|
||||
write!(temp_file, "{invalid_config}").unwrap();
|
||||
let temp_path = temp_file.path().to_str().unwrap();
|
||||
|
||||
let result = Config::read_config(temp_path);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_whitelist_creation() {
|
||||
let commands = vec!["echo".to_string(), "ls".to_string()];
|
||||
let regex = vec!["test.*".to_string()];
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: commands.clone(),
|
||||
regex: regex.clone(),
|
||||
};
|
||||
|
||||
assert_eq!(whitelist.commands, commands);
|
||||
assert_eq!(whitelist.regex, regex);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blacklist_creation() {
|
||||
let commands = vec!["rm".to_string(), "dd".to_string()];
|
||||
let regex = vec![".*[|&].*".to_string(), "^sudo .*".to_string()];
|
||||
|
||||
let blacklist = Blacklist {
|
||||
commands: commands.clone(),
|
||||
regex: regex.clone(),
|
||||
};
|
||||
|
||||
assert_eq!(blacklist.commands, commands);
|
||||
assert_eq!(blacklist.regex, regex);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_settings_creation() {
|
||||
let settings = Settings {
|
||||
port: 8080,
|
||||
host: "localhost".to_string(),
|
||||
env: Some("production".to_string()),
|
||||
};
|
||||
|
||||
assert_eq!(settings.port, 8080);
|
||||
assert_eq!(settings.host, "localhost");
|
||||
assert_eq!(settings.env, Some("production".to_string()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
//! Common utilities and shared components for the MCP Bash Server
|
||||
//!
|
||||
//! This module contains the core functionality including:
|
||||
//! - Bash command execution server implementation
|
||||
//! - Configuration management
|
||||
//! - OAuth2 authentication system
|
||||
//! - Command validation and security enforcement
|
||||
|
||||
pub mod bash_server;
|
||||
pub mod config;
|
||||
pub mod oauth;
|
||||
pub mod validator;
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,410 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
//! Command validation module for security enforcement
|
||||
//!
|
||||
//! This module provides command validation functionality to prevent execution
|
||||
//! of dangerous commands and operations based on configurable blacklists.
|
||||
|
||||
use tracing::debug;
|
||||
|
||||
use crate::{common::config::Whitelist, config::Blacklist};
|
||||
use rmcp::model::ErrorData;
|
||||
use tracing::error;
|
||||
|
||||
/// Command validator that checks commands against security blacklists
|
||||
/// Prevents execution of dangerous commands and operations
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Validator {
|
||||
/// Security blacklist configuration containing forbidden commands and operations
|
||||
blacklist: Blacklist,
|
||||
/// Security whitelist configuration containing secure commands
|
||||
whitelist: Whitelist,
|
||||
}
|
||||
|
||||
impl Validator {
|
||||
/// Create a new validator with the specified blacklist configuration
|
||||
pub fn new(blacklist: Blacklist, whitelist: Whitelist) -> Self {
|
||||
Validator {
|
||||
blacklist,
|
||||
whitelist,
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate a command against the security blacklist and whitelist
|
||||
/// Returns Ok(()) if command is safe, Err(ErrorData) if command is blocked
|
||||
///
|
||||
/// Security Logic:
|
||||
/// 1. All commands are denied by default
|
||||
/// 2. Blacklist has higher priority than whitelist
|
||||
/// 3. If command matches any blacklist pattern (exact or regex), deny
|
||||
/// 4. If command doesn't match blacklist and matches whitelist (exact or regex), allow
|
||||
/// 5. Otherwise, deny
|
||||
pub fn is_unsafe_command(&self, args: &str) -> Result<(), ErrorData> {
|
||||
// Handle empty command - deny by default
|
||||
if args.is_empty() {
|
||||
return Err(ErrorData::invalid_request(
|
||||
"Empty command not allowed".to_string(),
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
let full_cmd = args.to_string();
|
||||
|
||||
debug!("Validating command: {}", full_cmd);
|
||||
|
||||
// First check blacklist - if any match, deny immediately
|
||||
// Check blacklist exact commands - match against the full command
|
||||
for blacklisted_cmd in &self.blacklist.commands {
|
||||
if &full_cmd == blacklisted_cmd {
|
||||
error!(
|
||||
"Command blocked by blacklist exact match: {}, full command: {}",
|
||||
blacklisted_cmd, full_cmd
|
||||
);
|
||||
return Err(ErrorData::invalid_request(
|
||||
format!("Command blocked by blacklist: {blacklisted_cmd}"),
|
||||
None,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Check blacklist regex patterns
|
||||
for pattern in &self.blacklist.regex {
|
||||
match regex::Regex::new(pattern) {
|
||||
Ok(regex) => {
|
||||
if regex.is_match(&full_cmd) {
|
||||
error!(
|
||||
"Command blocked by blacklist regex pattern: {}, full command: {}",
|
||||
pattern, full_cmd
|
||||
);
|
||||
return Err(ErrorData::invalid_request(
|
||||
format!("Command blocked by blacklist pattern: {pattern}"),
|
||||
None,
|
||||
));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Invalid regex pattern in blacklist: {}, error: {}",
|
||||
pattern, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Now check whitelist - if any match, allow
|
||||
// Check whitelist exact commands
|
||||
if self.whitelist.commands.contains(&full_cmd) {
|
||||
debug!("Command allowed by whitelist exact match: {}", full_cmd);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Check whitelist regex patterns
|
||||
for pattern in &self.whitelist.regex {
|
||||
match regex::Regex::new(pattern) {
|
||||
Ok(regex) => {
|
||||
if regex.is_match(&full_cmd) {
|
||||
debug!("Command allowed by whitelist regex pattern: {}", pattern);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Invalid regex pattern in whitelist: {}, error: {}",
|
||||
pattern, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Default deny - command didn't match whitelist
|
||||
error!("Command denied - not in whitelist: {full_cmd}");
|
||||
Err(ErrorData::invalid_request(
|
||||
format!("Command not allowed: {full_cmd}"),
|
||||
None,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn create_test_validator() -> Validator {
|
||||
let blacklist = Blacklist {
|
||||
commands: vec![
|
||||
"rm".to_string(),
|
||||
"dd".to_string(),
|
||||
"shutdown".to_string(),
|
||||
"kill".to_string(),
|
||||
],
|
||||
regex: vec![
|
||||
".*[|&;><].*".to_string(), // Block commands with dangerous operators
|
||||
"^sudo .*".to_string(), // Block sudo commands
|
||||
".*/etc/.*".to_string(), // Block access to /etc
|
||||
".*passwd.*".to_string(), // Block password-related commands
|
||||
],
|
||||
};
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: vec![
|
||||
"echo hello".to_string(),
|
||||
"ls -la".to_string(),
|
||||
"pwd".to_string(),
|
||||
],
|
||||
regex: vec![
|
||||
"^echo [a-zA-Z0-9 ]+$".to_string(), // Only simple echo commands
|
||||
"^ls [a-zA-Z0-9 /-]*$".to_string(), // Only simple ls commands
|
||||
"^(pwd|whoami|date|uptime)$".to_string(), // Basic system info commands
|
||||
],
|
||||
};
|
||||
|
||||
Validator::new(blacklist, whitelist)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validator_creation() {
|
||||
let validator = create_test_validator();
|
||||
assert_eq!(validator.blacklist.commands.len(), 4);
|
||||
assert_eq!(validator.blacklist.regex.len(), 4);
|
||||
assert_eq!(validator.whitelist.commands.len(), 3);
|
||||
assert_eq!(validator.whitelist.regex.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_command() {
|
||||
let validator = create_test_validator();
|
||||
let result = validator.is_unsafe_command("pwd");
|
||||
assert!(result.is_ok()); // pwd is whitelisted
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_whitelisted_command() {
|
||||
let validator = create_test_validator();
|
||||
let result = validator.is_unsafe_command("echo hello");
|
||||
assert!(result.is_ok()); // "echo hello" is exactly whitelisted
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_whitelisted_regex_command() {
|
||||
let validator = create_test_validator();
|
||||
let result = validator.is_unsafe_command("echo test123");
|
||||
assert!(result.is_ok()); // Matches whitelist regex "^echo [a-zA-Z0-9 ]+$"
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blacklisted_command() {
|
||||
let validator = create_test_validator();
|
||||
// Test a command that contains a blacklisted command but isn't exact match
|
||||
let result = validator.is_unsafe_command("rm -rf /");
|
||||
// Should be denied because "rm -rf /" is not in whitelist (default deny)
|
||||
// not because "rm" is blacklisted (since we need exact command match)
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blacklisted_exact_command() {
|
||||
let validator = create_test_validator();
|
||||
// Test exact match against blacklist
|
||||
let result = validator.is_unsafe_command("rm");
|
||||
assert!(result.is_err()); // "rm" exactly matches blacklist
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blacklisted_regex_operation() {
|
||||
let validator = create_test_validator();
|
||||
// Test pipe operation which should be blocked by regex
|
||||
let result = validator.is_unsafe_command("echo test | cat");
|
||||
assert!(result.is_err()); // Contains "|" which matches ".*[|&;><].*" regex
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_blacklisted_commands() {
|
||||
let validator = create_test_validator();
|
||||
|
||||
// Test commands that are exactly in the blacklist
|
||||
let dangerous_commands = vec!["rm", "dd", "shutdown", "kill"];
|
||||
for cmd in dangerous_commands {
|
||||
let result = validator.is_unsafe_command(cmd);
|
||||
assert!(result.is_err(), "Command '{cmd}' should be blocked");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_blacklisted_regex_operations() {
|
||||
let validator = create_test_validator();
|
||||
|
||||
let dangerous_ops = vec!["|", "&", ";", ">"];
|
||||
for op in dangerous_ops {
|
||||
let cmd = format!("echo test {op}");
|
||||
let result = validator.is_unsafe_command(&cmd);
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Operation '{op}' should be blocked by regex"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_command() {
|
||||
let validator = create_test_validator();
|
||||
let result = validator.is_unsafe_command("");
|
||||
assert!(result.is_err()); // Empty commands are now denied by default
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_complex_safe_command() {
|
||||
let validator = create_test_validator();
|
||||
let result = validator.is_unsafe_command("find /tmp -name *.txt");
|
||||
assert!(result.is_err()); // Not in whitelist, so denied by default
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_regex_in_whitelist() {
|
||||
let blacklist = Blacklist {
|
||||
commands: vec![],
|
||||
regex: vec![],
|
||||
};
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: vec![],
|
||||
regex: vec!["[invalid_regex".to_string()], // Invalid regex
|
||||
};
|
||||
|
||||
let validator = Validator::new(blacklist, whitelist);
|
||||
let result = validator.is_unsafe_command("test");
|
||||
// Should be denied because command is not in whitelist and default is deny
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_case_sensitivity() {
|
||||
let validator = create_test_validator();
|
||||
|
||||
// Test uppercase version of blacklisted command
|
||||
let result = validator.is_unsafe_command("RM");
|
||||
// Should be denied because not in whitelist (default deny)
|
||||
assert!(result.is_err());
|
||||
|
||||
// Test exact case match for blacklisted command
|
||||
let result = validator.is_unsafe_command("rm");
|
||||
assert!(result.is_err()); // Should be blocked by blacklist
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_command_match() {
|
||||
let validator = create_test_validator();
|
||||
|
||||
// Test command that contains blacklisted word but isn't exact match
|
||||
let result = validator.is_unsafe_command("remove"); // contains "rm" but shouldn't match
|
||||
assert!(result.is_err()); // Should be denied because not in whitelist (default deny)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blacklist_priority_over_whitelist() {
|
||||
// Test that blacklist has higher priority than whitelist
|
||||
let blacklist = Blacklist {
|
||||
commands: vec!["echo hello".to_string()], // Changed to exact match
|
||||
regex: vec![],
|
||||
};
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: vec!["echo hello".to_string()],
|
||||
regex: vec!["^echo .*".to_string()],
|
||||
};
|
||||
|
||||
let validator = Validator::new(blacklist, whitelist);
|
||||
let result = validator.is_unsafe_command("echo hello");
|
||||
assert!(result.is_err()); // Should be blocked because "echo hello" is in blacklist
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blacklist_regex_priority() {
|
||||
// Test that blacklist regex has higher priority than whitelist
|
||||
let blacklist = Blacklist {
|
||||
commands: vec![],
|
||||
regex: vec![".*sudo.*".to_string()],
|
||||
};
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: vec!["sudo ls".to_string()],
|
||||
regex: vec!["^sudo .*".to_string()],
|
||||
};
|
||||
|
||||
let validator = Validator::new(blacklist, whitelist);
|
||||
let result = validator.is_unsafe_command("sudo ls");
|
||||
assert!(result.is_err()); // Should be blocked by blacklist regex
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_deny_behavior() {
|
||||
// Test that commands not in whitelist are denied by default
|
||||
let blacklist = Blacklist {
|
||||
commands: vec![],
|
||||
regex: vec![],
|
||||
};
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: vec!["pwd".to_string()],
|
||||
regex: vec![],
|
||||
};
|
||||
|
||||
let validator = Validator::new(blacklist, whitelist);
|
||||
|
||||
// Command in whitelist should be allowed
|
||||
let result = validator.is_unsafe_command("pwd");
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Command not in whitelist should be denied
|
||||
let result = validator.is_unsafe_command("ls");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_vs_partial_blacklist_matching() {
|
||||
// Test the fix: blacklist should match exactly, not partially
|
||||
let blacklist = Blacklist {
|
||||
commands: vec!["rm".to_string()], // Only "rm" exactly
|
||||
regex: vec![],
|
||||
};
|
||||
|
||||
let whitelist = Whitelist {
|
||||
commands: vec!["rm -rf /tmp/test".to_string()], // Allow this specific command
|
||||
regex: vec![],
|
||||
};
|
||||
|
||||
let validator = Validator::new(blacklist, whitelist);
|
||||
|
||||
// "rm" exactly should be blocked by blacklist
|
||||
let result = validator.is_unsafe_command("rm");
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Exact match 'rm' should be blocked by blacklist"
|
||||
);
|
||||
|
||||
// "rm -rf /tmp/test" should be allowed by whitelist (not blocked by blacklist)
|
||||
let result = validator.is_unsafe_command("rm -rf /tmp/test");
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"'rm -rf /tmp/test' should be allowed by whitelist"
|
||||
);
|
||||
|
||||
// "rm -rf /" should be denied (not in whitelist, not exact blacklist match)
|
||||
let result = validator.is_unsafe_command("rm -rf /");
|
||||
assert!(result.is_err(), "'rm -rf /' should be denied by default");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
<!--
|
||||
~ 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.
|
||||
-->
|
||||
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>MCP OAuth Server</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; margin: 40px auto; max-width: 800px; line-height: 1.6; }
|
||||
h1, h2 { color: #333; }
|
||||
code { background: #f4f4f4; padding: 2px 5px; border-radius: 3px; }
|
||||
.endpoint { background: #f9f9f9; padding: 15px; border-radius: 5px; margin-bottom: 15px; }
|
||||
.flow { background: #e8f5e9; padding: 15px; border-radius: 5px; margin-bottom: 15px; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>MCP OAuth Server</h1>
|
||||
<p>This is an MCP server with OAuth 2.0 integration to a third-party authorization server.</p>
|
||||
|
||||
<h2>Available Endpoints:</h2>
|
||||
|
||||
<div class="endpoint">
|
||||
<h3>Authorization Endpoint</h3>
|
||||
<p><code>GET /authorize</code></p>
|
||||
<p>Parameters:</p>
|
||||
<ul>
|
||||
<li><code>response_type</code> - Must be "code"</li>
|
||||
<li><code>client_id</code> - Client identifier (e.g., "mcp-client")</li>
|
||||
<li><code>redirect_uri</code> - URI to redirect after authorization</li>
|
||||
<li><code>scope</code> - Optional requested scope</li>
|
||||
<li><code>state</code> - Optional state value for CSRF prevention</li>
|
||||
</ul>
|
||||
</div>
|
||||
|
||||
<div class="endpoint">
|
||||
<h3>Token Endpoint</h3>
|
||||
<p><code>POST /token</code></p>
|
||||
<p>Parameters:</p>
|
||||
<ul>
|
||||
<li><code>grant_type</code> - Must be "authorization_code"</li>
|
||||
<li><code>code</code> - The authorization code</li>
|
||||
<li><code>client_id</code> - Client identifier</li>
|
||||
<li><code>client_secret</code> - Client secret</li>
|
||||
<li><code>redirect_uri</code> - Redirect URI used in authorization request</li>
|
||||
</ul>
|
||||
</div>
|
||||
|
||||
<div class="endpoint">
|
||||
<h3>MCP streamablehttp Endpoints</h3>
|
||||
<p><code>/mcp</code> - Streamablehttp connection endpoint (requires OAuth token)</p>
|
||||
</div>
|
||||
|
||||
<div class="flow">
|
||||
<h2>OAuth Flow:</h2>
|
||||
<ol>
|
||||
<li>MCP Client initiates OAuth flow with this MCP Server</li>
|
||||
<li>MCP Server redirects to Third-Party OAuth Server</li>
|
||||
<li>User authenticates with Third-Party Server</li>
|
||||
<li>Third-Party Server redirects back to MCP Server with auth code</li>
|
||||
<li>MCP Server exchanges the code for a third-party access token</li>
|
||||
<li>MCP Server generates its own token bound to the third-party session</li>
|
||||
<li>MCP Server completes the OAuth flow with the MCP Client</li>
|
||||
</ol>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,618 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
//! MCP Bash Server - A Model Context Protocol server for executing bash commands
|
||||
//!
|
||||
//! This server provides secure bash command execution capabilities through the MCP protocol.
|
||||
//! It includes OAuth2 authentication, command validation, and cross-platform support.
|
||||
//!
|
||||
//! Features:
|
||||
//! - Secure command execution with blacklist validation
|
||||
//! - OAuth2 authentication for client authorization
|
||||
//! - Cross-platform shell support (Linux, Windows, macOS)
|
||||
//! - Built-in system information tools
|
||||
//! - Configurable timeout and environment settings
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::{net::SocketAddr, sync::Arc};
|
||||
|
||||
use anyhow::Result;
|
||||
use axum::{
|
||||
Router,
|
||||
body::Body,
|
||||
http::{HeaderMap, Request},
|
||||
middleware::{self, Next},
|
||||
response::{Html, IntoResponse, Response},
|
||||
routing::{get, post},
|
||||
};
|
||||
use rmcp::transport::streamable_http_server::{
|
||||
StreamableHttpService, session::local::LocalSessionManager,
|
||||
};
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tracing::info;
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||
|
||||
// Import modules
|
||||
mod common;
|
||||
use common::bash_server::BashServer;
|
||||
use common::config;
|
||||
use common::oauth::{
|
||||
McpOAuthStore, oauth_approve, oauth_authorization_server, oauth_authorize, oauth_register,
|
||||
oauth_token, validate_token_middleware,
|
||||
};
|
||||
|
||||
const INDEX_HTML: &str = include_str!("html/mcp_oauth_index.html");
|
||||
|
||||
/// Global storage for server bind address, initialized once at startup
|
||||
// Init once from environment variable BIND_ADDRESS
|
||||
pub static BIND_ADDRESS: OnceLock<String> = OnceLock::new();
|
||||
|
||||
/// Root path handler
|
||||
/// Serves the main OAuth authorization index page
|
||||
async fn index() -> Html<&'static str> {
|
||||
Html(INDEX_HTML)
|
||||
}
|
||||
|
||||
/// Wrapper function for oauth_authorization_server to handle BIND_ADDRESS
|
||||
async fn oauth_authorization_server_handler(headers: HeaderMap) -> impl IntoResponse {
|
||||
let bind_address = BIND_ADDRESS
|
||||
.get()
|
||||
.expect("BIND_ADDRESS must be initialized in main()");
|
||||
oauth_authorization_server(bind_address, headers).await
|
||||
}
|
||||
|
||||
/// HTTP request logging middleware
|
||||
/// Logs all incoming requests including method, URI, headers and response status
|
||||
async fn log_request(request: Request<Body>, next: Next) -> Response {
|
||||
let method = request.method().clone();
|
||||
let uri = request.uri().clone();
|
||||
let version = request.version();
|
||||
|
||||
// Log headers
|
||||
let headers = request.headers().clone();
|
||||
let mut header_log = String::new();
|
||||
for (key, value) in headers.iter() {
|
||||
let value_str = value.to_str().unwrap_or("<binary>");
|
||||
header_log.push_str(&format!("\n {key}: {value_str}"));
|
||||
}
|
||||
|
||||
// Try to get request body for form submissions
|
||||
let content_type = headers
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("");
|
||||
|
||||
let request_info = if content_type.contains("application/x-www-form-urlencoded")
|
||||
|| content_type.contains("application/json")
|
||||
{
|
||||
format!("{method} {uri} {version:?}{header_log}\nContent-Type: {content_type}")
|
||||
} else {
|
||||
format!("{method} {uri} {version:?}{header_log}")
|
||||
};
|
||||
|
||||
info!("REQUEST: {}", request_info);
|
||||
|
||||
// Call the actual handler
|
||||
let response = next.run(request).await;
|
||||
|
||||
// Log response status
|
||||
let status = response.status();
|
||||
info!("RESPONSE: {} for {} {}", status, method, uri);
|
||||
|
||||
response
|
||||
}
|
||||
|
||||
/// Main application entry point
|
||||
/// Sets up logging, OAuth store, HTTP server, and starts the MCP bash server
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
// Initialize logging
|
||||
let logs = tracing_appender::rolling::daily("logs", "mcp.log");
|
||||
let (non_blocking, _guard) = tracing_appender::non_blocking(logs);
|
||||
let log_setting = tracing_subscriber::fmt::layer().with_writer(non_blocking);
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "debug".to_string().into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.with(log_setting)
|
||||
.init();
|
||||
|
||||
// Read environment mode from config file, default to "production"
|
||||
let config = config::Config::read_config("config.toml")?;
|
||||
let env_mode = config
|
||||
.settings
|
||||
.env
|
||||
.clone()
|
||||
.unwrap_or_else(|| "production".to_string());
|
||||
let is_dev = env_mode == "development";
|
||||
|
||||
// Create the OAuth store
|
||||
let oauth_store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
let host = config.settings.host.clone();
|
||||
let port = config.settings.port;
|
||||
let bind_address = format!("{host}:{port}");
|
||||
|
||||
let addr = bind_address.parse::<SocketAddr>()?;
|
||||
let _ = BIND_ADDRESS.set(bind_address);
|
||||
|
||||
// Create StreamableHttpServer
|
||||
let service = StreamableHttpService::new(
|
||||
|| Ok(BashServer::new()),
|
||||
LocalSessionManager::default().into(),
|
||||
Default::default(),
|
||||
);
|
||||
|
||||
let server_router = Router::new().nest_service("/mcp", service);
|
||||
|
||||
// Add OAuth authentication middleware only if not in development mode
|
||||
let protected_server_router = if is_dev {
|
||||
server_router
|
||||
} else {
|
||||
server_router.layer(middleware::from_fn_with_state(
|
||||
oauth_store.clone(),
|
||||
validate_token_middleware,
|
||||
))
|
||||
};
|
||||
|
||||
// Create CORS layer for the oauth authorization server endpoint
|
||||
let cors_layer = CorsLayer::new()
|
||||
.allow_origin(Any)
|
||||
.allow_methods(Any)
|
||||
.allow_headers(Any);
|
||||
|
||||
// Create a sub-router for the oauth authorization server endpoint with CORS
|
||||
let oauth_server_router = Router::new()
|
||||
.route(
|
||||
"/.well-known/oauth-authorization-server",
|
||||
get(oauth_authorization_server_handler).options(oauth_authorization_server_handler),
|
||||
)
|
||||
.route("/token", post(oauth_token).options(oauth_token))
|
||||
.route("/register", post(oauth_register).options(oauth_register))
|
||||
.layer(cors_layer)
|
||||
.with_state(oauth_store.clone());
|
||||
|
||||
// Create HTTP router with request logging middleware
|
||||
let app = Router::new()
|
||||
.route("/", get(index))
|
||||
.route("/authorize", get(oauth_authorize))
|
||||
.route("/approve", post(oauth_approve))
|
||||
.merge(oauth_server_router) // Merge the CORS-enabled oauth server router
|
||||
.merge(protected_server_router)
|
||||
.with_state(oauth_store.clone())
|
||||
.layer(middleware::from_fn(log_request));
|
||||
|
||||
// Start HTTP server
|
||||
info!("MCP OAuth Server started on {}", addr);
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() })
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::body::Body;
|
||||
use axum::http::Method;
|
||||
use axum::http::Request;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_index_handler() {
|
||||
let response = index().await;
|
||||
let html_content = response.0;
|
||||
|
||||
// Verify it returns the expected HTML content
|
||||
assert_eq!(html_content, INDEX_HTML);
|
||||
assert!(html_content.contains("OAuth"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_authorization_server_handler() {
|
||||
use axum::http::HeaderMap;
|
||||
|
||||
// Set up BIND_ADDRESS for testing
|
||||
let _ = BIND_ADDRESS.set("localhost:8080".to_string());
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("host", "localhost:8080".parse().unwrap());
|
||||
|
||||
let response = oauth_authorization_server_handler(headers).await;
|
||||
|
||||
// Test that the handler returns a response
|
||||
// We can't easily test the exact content without mocking, but we can verify it doesn't panic
|
||||
let _response_body = response.into_response();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bind_address_initialization() {
|
||||
// Create a new OnceLock for testing to avoid conflicts
|
||||
let test_bind_address: OnceLock<String> = OnceLock::new();
|
||||
|
||||
// Test that we can set the value once
|
||||
let result = test_bind_address.set("127.0.0.1:9090".to_string());
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Test that we can get the value
|
||||
let value = test_bind_address.get();
|
||||
assert!(value.is_some());
|
||||
assert_eq!(value.unwrap(), "127.0.0.1:9090");
|
||||
|
||||
// Test that we can't set it again
|
||||
let result2 = test_bind_address.set("different:port".to_string());
|
||||
assert!(result2.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_index_html_constant() {
|
||||
// Test that INDEX_HTML is not empty and contains expected content
|
||||
assert!(INDEX_HTML.contains("html") || INDEX_HTML.contains("HTML"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_log_request_middleware_functionality() {
|
||||
// Test basic properties of log_request function
|
||||
// Since it requires complex setup with actual middleware,
|
||||
// we focus on testing the types and structure
|
||||
|
||||
let request = Request::builder()
|
||||
.method(Method::GET)
|
||||
.uri("/test")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
// Verify request properties that log_request would process
|
||||
assert_eq!(request.method(), Method::GET);
|
||||
assert_eq!(request.uri().path(), "/test");
|
||||
assert!(request.headers().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_module_imports() {
|
||||
// Test that our modules are properly imported and accessible
|
||||
let _server = BashServer::new();
|
||||
let _store = McpOAuthStore::new();
|
||||
|
||||
// Test config module
|
||||
let config_result = config::Config::read_config("nonexistent.toml");
|
||||
assert!(config_result.is_err()); // Should fail gracefully
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_error_handling_types() {
|
||||
// Test that Result type is properly used
|
||||
let test_result: Result<String> = Ok("test".to_string());
|
||||
assert!(test_result.is_ok());
|
||||
|
||||
let test_error: Result<String> = Err(anyhow::anyhow!("test error"));
|
||||
assert!(test_error.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dependencies_availability() {
|
||||
// Test that critical dependencies are available
|
||||
use std::sync::Arc;
|
||||
|
||||
let _arc_store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Test that we can create basic types
|
||||
let _socket_addr: Result<SocketAddr, _> = "127.0.0.1:8080".parse();
|
||||
}
|
||||
|
||||
// ========== OAuth Mock Tests ==========
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_store_functionality() {
|
||||
use common::oauth::{AuthToken, McpOAuthStore};
|
||||
use oauth2::{AccessToken, EmptyExtraTokenFields};
|
||||
|
||||
let store = McpOAuthStore::new();
|
||||
|
||||
// Test client validation
|
||||
let client = store
|
||||
.validate_client("mcp-client", "http://localhost:8080/callback")
|
||||
.await;
|
||||
assert!(client.is_some());
|
||||
|
||||
let invalid_client = store
|
||||
.validate_client("invalid-client", "http://localhost:8080/callback")
|
||||
.await;
|
||||
assert!(invalid_client.is_none());
|
||||
|
||||
// Test auth session creation
|
||||
let session_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile email".to_string()),
|
||||
Some("test-state".to_string()),
|
||||
"test-session-123".to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(session_id, "test-session-123");
|
||||
|
||||
// Test token update and MCP token creation
|
||||
let auth_token = AuthToken::new(
|
||||
AccessToken::new("mock-third-party-token".to_string()),
|
||||
oauth2::basic::BasicTokenType::Bearer,
|
||||
EmptyExtraTokenFields {},
|
||||
);
|
||||
|
||||
let update_result = store
|
||||
.update_auth_session_token(&session_id, auth_token)
|
||||
.await;
|
||||
assert!(update_result.is_ok());
|
||||
|
||||
let mcp_token = store.create_mcp_token(&session_id).await;
|
||||
assert!(mcp_token.is_ok());
|
||||
|
||||
let token = mcp_token.unwrap();
|
||||
assert!(token.access_token.starts_with("mcp-token-"));
|
||||
assert_eq!(token.client_id, "mcp-client");
|
||||
|
||||
// Test token validation
|
||||
let validated = store.validate_token(&token.access_token).await;
|
||||
assert!(validated.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_authorization_flow_mock() {
|
||||
use common::oauth::{AuthorizeQuery, McpOAuthStore};
|
||||
use std::sync::Arc;
|
||||
|
||||
let store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Mock authorization request
|
||||
let auth_query = AuthorizeQuery {
|
||||
response_type: "code".to_string(),
|
||||
client_id: "mcp-client".to_string(),
|
||||
redirect_uri: "http://localhost:8080/callback".to_string(),
|
||||
scope: Some("profile email".to_string()),
|
||||
state: Some("test-state-456".to_string()),
|
||||
};
|
||||
|
||||
// Test that oauth_authorize function can be called
|
||||
// Note: In a real test, we'd use test frameworks like tower::ServiceExt
|
||||
// but here we're testing the basic functionality
|
||||
let store_clone = store.clone();
|
||||
let sessions_before = store_clone.auth_sessions.read().await.len();
|
||||
|
||||
// Verify store is accessible and functional
|
||||
assert_eq!(sessions_before, 0);
|
||||
|
||||
// Test client validation within the flow
|
||||
let client_validation = store
|
||||
.validate_client(&auth_query.client_id, &auth_query.redirect_uri)
|
||||
.await;
|
||||
assert!(client_validation.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_token_exchange_mock() {
|
||||
use common::oauth::AuthToken;
|
||||
use common::oauth::{McpOAuthStore, TokenRequest};
|
||||
use oauth2::{AccessToken, EmptyExtraTokenFields};
|
||||
use std::sync::Arc;
|
||||
|
||||
let store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Create a session and add auth token (simulating successful OAuth flow)
|
||||
let session_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile".to_string()),
|
||||
Some("test-state".to_string()),
|
||||
"token-exchange-session".to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_token = AuthToken::new(
|
||||
AccessToken::new("mock-external-token".to_string()),
|
||||
oauth2::basic::BasicTokenType::Bearer,
|
||||
EmptyExtraTokenFields {},
|
||||
);
|
||||
|
||||
store
|
||||
.update_auth_session_token(&session_id, auth_token)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Mock token request (just for structure validation)
|
||||
let _token_request = TokenRequest {
|
||||
grant_type: "authorization_code".to_string(),
|
||||
code: "mock-auth-code".to_string(),
|
||||
client_id: "mcp-client".to_string(),
|
||||
client_secret: "mcp-client-secret".to_string(),
|
||||
redirect_uri: "http://localhost:8080/callback".to_string(),
|
||||
code_verifier: None,
|
||||
refresh_token: "".to_string(),
|
||||
};
|
||||
|
||||
// Test token creation
|
||||
let mcp_token_result = store.create_mcp_token(&session_id).await;
|
||||
assert!(mcp_token_result.is_ok());
|
||||
|
||||
let mcp_token = mcp_token_result.unwrap();
|
||||
assert_eq!(mcp_token.token_type, "bearer");
|
||||
assert_eq!(mcp_token.expires_in, Some(3600));
|
||||
assert!(mcp_token.refresh_token.is_some());
|
||||
|
||||
// Verify token can be validated
|
||||
let validation_result = store.validate_token(&mcp_token.access_token).await;
|
||||
assert!(validation_result.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_middleware_functionality() {
|
||||
use common::oauth::AuthToken;
|
||||
use common::oauth::McpOAuthStore;
|
||||
use oauth2::{AccessToken, EmptyExtraTokenFields};
|
||||
use std::sync::Arc;
|
||||
|
||||
let store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Create a valid token for middleware testing
|
||||
let session_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile".to_string()),
|
||||
None,
|
||||
"middleware-test-session".to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_token = AuthToken::new(
|
||||
AccessToken::new("middleware-test-token".to_string()),
|
||||
oauth2::basic::BasicTokenType::Bearer,
|
||||
EmptyExtraTokenFields {},
|
||||
);
|
||||
|
||||
store
|
||||
.update_auth_session_token(&session_id, auth_token)
|
||||
.await
|
||||
.unwrap();
|
||||
let mcp_token = store.create_mcp_token(&session_id).await.unwrap();
|
||||
|
||||
// Test token validation (simulating middleware behavior)
|
||||
let valid_token_check = store.validate_token(&mcp_token.access_token).await;
|
||||
assert!(valid_token_check.is_some());
|
||||
|
||||
// Test invalid token
|
||||
let invalid_token_check = store.validate_token("invalid-token-12345").await;
|
||||
assert!(invalid_token_check.is_none());
|
||||
|
||||
// Test empty token
|
||||
let empty_token_check = store.validate_token("").await;
|
||||
assert!(empty_token_check.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_error_handling() {
|
||||
use common::oauth::McpOAuthStore;
|
||||
use std::sync::Arc;
|
||||
|
||||
let store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Test creating MCP token without session
|
||||
let no_session_result = store.create_mcp_token("nonexistent-session").await;
|
||||
assert!(no_session_result.is_err());
|
||||
assert_eq!(no_session_result.unwrap_err(), "Session not found");
|
||||
|
||||
// Test creating MCP token without auth token in session
|
||||
let session_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile".to_string()),
|
||||
None,
|
||||
"no-auth-token-session".to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let no_auth_token_result = store.create_mcp_token(&session_id).await;
|
||||
assert!(no_auth_token_result.is_err());
|
||||
assert_eq!(
|
||||
no_auth_token_result.unwrap_err(),
|
||||
"No third-party token available for session"
|
||||
);
|
||||
|
||||
// Test updating nonexistent session
|
||||
let auth_token = oauth2::StandardTokenResponse::new(
|
||||
oauth2::AccessToken::new("test-token".to_string()),
|
||||
oauth2::basic::BasicTokenType::Bearer,
|
||||
oauth2::EmptyExtraTokenFields {},
|
||||
);
|
||||
|
||||
let update_nonexistent = store
|
||||
.update_auth_session_token("nonexistent", auth_token)
|
||||
.await;
|
||||
assert!(update_nonexistent.is_err());
|
||||
assert_eq!(update_nonexistent.unwrap_err(), "Session not found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_oauth_security_validations() {
|
||||
use common::oauth::McpOAuthStore;
|
||||
use std::sync::Arc;
|
||||
|
||||
let store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Test invalid client ID
|
||||
let invalid_client = store
|
||||
.validate_client("malicious-client", "http://localhost:8080/callback")
|
||||
.await;
|
||||
assert!(invalid_client.is_none());
|
||||
|
||||
// Test invalid redirect URI (potential open redirect attack)
|
||||
let malicious_redirect = store
|
||||
.validate_client("mcp-client", "http://evil.com/steal-tokens")
|
||||
.await;
|
||||
assert!(malicious_redirect.is_none());
|
||||
|
||||
// Test valid client with valid redirect URI
|
||||
let valid_client = store
|
||||
.validate_client("mcp-client", "http://localhost:8080/callback")
|
||||
.await;
|
||||
assert!(valid_client.is_some());
|
||||
|
||||
// Test that tokens are properly random and unique
|
||||
let session1_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile".to_string()),
|
||||
None,
|
||||
"security-test-1".to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let session2_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile".to_string()),
|
||||
None,
|
||||
"security-test-2".to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Add auth tokens to both sessions
|
||||
for (i, session_id) in [&session1_id, &session2_id].iter().enumerate() {
|
||||
let auth_token = oauth2::StandardTokenResponse::new(
|
||||
oauth2::AccessToken::new(format!("security-token-{i}")),
|
||||
oauth2::basic::BasicTokenType::Bearer,
|
||||
oauth2::EmptyExtraTokenFields {},
|
||||
);
|
||||
store
|
||||
.update_auth_session_token(session_id, auth_token)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let token1 = store.create_mcp_token(&session1_id).await.unwrap();
|
||||
let token2 = store.create_mcp_token(&session2_id).await.unwrap();
|
||||
|
||||
// Tokens should be different
|
||||
assert_ne!(token1.access_token, token2.access_token);
|
||||
assert_ne!(token1.refresh_token, token2.refresh_token);
|
||||
|
||||
// Both should be valid
|
||||
assert!(store.validate_token(&token1.access_token).await.is_some());
|
||||
assert!(store.validate_token(&token2.access_token).await.is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
<!--
|
||||
~ 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.
|
||||
-->
|
||||
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>MCP OAuth</title>
|
||||
<style>
|
||||
:root {
|
||||
--primary-color: #4285f4;
|
||||
--secondary-color: #f1f1f1;
|
||||
--text-color: #333;
|
||||
--border-color: #ddd;
|
||||
}
|
||||
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
background-color: #f8f9fa;
|
||||
color: var(--text-color);
|
||||
}
|
||||
|
||||
.container {
|
||||
background: white;
|
||||
padding: 2rem;
|
||||
border-radius: 12px;
|
||||
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
|
||||
max-width: 600px;
|
||||
width: 90%;
|
||||
margin: 1rem;
|
||||
}
|
||||
|
||||
h1 {
|
||||
margin: 0 0 1.5rem 0;
|
||||
font-size: 1.8rem;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
.client-info {
|
||||
background: var(--secondary-color);
|
||||
padding: 1rem;
|
||||
border-radius: 8px;
|
||||
margin-bottom: 1.5rem;
|
||||
}
|
||||
|
||||
.btn-group {
|
||||
display: flex;
|
||||
gap: 1rem;
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
.btn {
|
||||
padding: 0.75rem 1.5rem;
|
||||
border-radius: 6px;
|
||||
cursor: pointer;
|
||||
font-size: 1rem;
|
||||
transition: all 0.2s;
|
||||
border: none;
|
||||
}
|
||||
|
||||
.btn-primary {
|
||||
background-color: var(--primary-color);
|
||||
color: white;
|
||||
}
|
||||
|
||||
.btn-primary:hover {
|
||||
background-color: #3367d6;
|
||||
}
|
||||
|
||||
.btn-secondary {
|
||||
background-color: var(--secondary-color);
|
||||
color: var(--text-color);
|
||||
border: 1px solid var(--border-color);
|
||||
}
|
||||
|
||||
.btn-secondary:hover {
|
||||
background-color: #e0e0e0;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<h1>MCP OAuth</h1>
|
||||
<div class="client-info">
|
||||
<p><strong>{{ client_id }}</strong> requests access to your account.</p>
|
||||
<p>requested scopes: {{ scopes }}</p>
|
||||
</div>
|
||||
|
||||
<form action="/approve" method="post">
|
||||
<input type="hidden" name="client_id" value="{{ client_id }}">
|
||||
<input type="hidden" name="redirect_uri" value="{{ redirect_uri }}">
|
||||
<input type="hidden" name="scope" value="{{ scope }}">
|
||||
<input type="hidden" name="state" value="{{ state }}">
|
||||
|
||||
<div class="btn-group">
|
||||
<button type="submit" name="approved" value="true" class="btn btn-primary">Approve</button>
|
||||
<button type="submit" name="approved" value="false" class="btn btn-secondary">Reject</button>
|
||||
</div>
|
||||
</form>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,74 @@
|
||||
# HertzBeat 日志 MCP
|
||||
|
||||
该 MCP 服务基于 GreptimeDB 查询 HertzBeat 自身运行日志。使用前需要启用 GreptimeDB 日志写入。
|
||||
|
||||
## 结构化只读查询
|
||||
|
||||
服务只暴露 `query_logs`,不再接收或执行调用方提供的 SQL。服务端生成的查询固定为:
|
||||
|
||||
```sql
|
||||
SELECT timestamp, severity_text, body
|
||||
FROM hzb_logs
|
||||
WHERE <结构化过滤条件>
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT <1-100>
|
||||
```
|
||||
|
||||
`query_logs` 支持以下可选参数:
|
||||
|
||||
| 参数 | 类型 | 说明 |
|
||||
| --- | --- | --- |
|
||||
| `severity` | 字符串 | `TRACE`、`DEBUG`、`INFO`、`WARN`、`ERROR` 或 `FATAL` |
|
||||
| `keyword` | 字符串 | 日志正文关键词,最长 256 个字符 |
|
||||
| `startTime` | 整数 | 开始时间,Unix 毫秒时间戳 |
|
||||
| `endTime` | 整数 | 结束时间,Unix 毫秒时间戳 |
|
||||
| `limit` | 整数 | 返回条数,默认为 20,范围为 1~100 |
|
||||
|
||||
调用示例:
|
||||
|
||||
```json
|
||||
{
|
||||
"severity": "ERROR",
|
||||
"keyword": "connection refused",
|
||||
"startTime": 1783785600000,
|
||||
"endTime": 1783872000000,
|
||||
"limit": 20
|
||||
}
|
||||
```
|
||||
|
||||
当前 `hzb_logs` 表没有 `monitorId` 字段,因此本接口不提供无效的监控 ID 过滤。如果后续需要该能力,应先在 OpenTelemetry 日志写入链中定义并提取统一的监控 ID 字段。
|
||||
|
||||
## GreptimeDB 账号
|
||||
|
||||
服务支持通过 `greptime.username` 和 `greptime.password` 发送 HTTP Basic Authentication。用户名和密码必须同时配置;建议配合 HTTPS 或可信内网使用。
|
||||
|
||||
项目当前 Docker Compose 使用 GreptimeDB `v0.14.3`,该版本只提供身份认证,不能限制用户为只读权限。因此当前真正生效的安全边界是“删除原始 SQL参数并固定生成单条 `SELECT`”。使用 GreptimeDB 1.0 及以上版本时,应为该 MCP 配置独立的 `ro`/`readonly` 账号。
|
||||
|
||||
## Claude Desktop 集成(stdio)
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"hertzbeat-mcp": {
|
||||
"command": "java",
|
||||
"args": [
|
||||
"-Dspring.ai.mcp.server.stdio=true",
|
||||
"-Dspring.main.web-application-type=none",
|
||||
"-Dlogging.pattern.console=",
|
||||
"-jar",
|
||||
"${PATH}/hertzbeat-mcp-2.0-SNAPSHOT.jar"
|
||||
],
|
||||
"env": {
|
||||
"GREPTIME_URL": "http://${IP}:4000",
|
||||
"GREPTIME_DATABASE": "public",
|
||||
"GREPTIME_USERNAME": "${READ_ONLY_USERNAME}",
|
||||
"GREPTIME_PASSWORD": "${READ_ONLY_PASSWORD}"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
不启用 GreptimeDB 认证时,可以省略 `GREPTIME_USERNAME` 和 `GREPTIME_PASSWORD`;不能只配置其中一个。
|
||||
|
||||
旧的 `getHertzbeatLog({"querySql": "..."})` Tool 已直接移除,不保留兼容入口,以免继续暴露任意 SQL执行能力。
|
||||
@@ -0,0 +1,99 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!--
|
||||
~ 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.
|
||||
-->
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>org.apache.hertzbeat</groupId>
|
||||
<artifactId>hertzbeat</artifactId>
|
||||
<version>2.0-SNAPSHOT</version>
|
||||
</parent>
|
||||
|
||||
<artifactId>hertzbeat-mcp</artifactId>
|
||||
|
||||
<properties>
|
||||
<maven.compiler.source>17</maven.compiler.source>
|
||||
<maven.compiler.target>17</maven.compiler.target>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<spring-ai.version>1.0.0-M6</spring-ai.version>
|
||||
</properties>
|
||||
|
||||
<dependencyManagement>
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-dependencies</artifactId>
|
||||
<version>3.4.2</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bom</artifactId>
|
||||
<version>${spring-ai.version}</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
</dependencyManagement>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-mcp-server-webflux-spring-boot-starter</artifactId>
|
||||
<version>${spring-ai.version}</version>
|
||||
</dependency>
|
||||
<!-- json path parser-->
|
||||
<dependency>
|
||||
<groupId>com.jayway.jsonpath</groupId>
|
||||
<artifactId>json-path</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework</groupId>
|
||||
<artifactId>spring-web</artifactId>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-maven-plugin</artifactId>
|
||||
<executions>
|
||||
<execution>
|
||||
<goals>
|
||||
<goal>repackage</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-compiler-plugin</artifactId>
|
||||
<version>${maven-compiler-plugin.version}</version>
|
||||
<configuration>
|
||||
<release>${java.version}</release>
|
||||
<compilerArgs>
|
||||
<compilerArg>-parameters</compilerArg>
|
||||
</compilerArgs>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
package org.apache.hertzbeat.mcp.server;
|
||||
|
||||
import org.apache.hertzbeat.mcp.server.service.LogService;
|
||||
import org.springframework.ai.tool.ToolCallbackProvider;
|
||||
import org.springframework.ai.tool.method.MethodToolCallbackProvider;
|
||||
import org.springframework.boot.SpringApplication;
|
||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
|
||||
/**
|
||||
* MCP Server Application
|
||||
*/
|
||||
@SpringBootApplication
|
||||
public class McpServerApplication {
|
||||
|
||||
public static void main(String[] args) {
|
||||
SpringApplication.run(McpServerApplication.class, args);
|
||||
}
|
||||
|
||||
@Bean
|
||||
public ToolCallbackProvider tools(
|
||||
LogService logService) {
|
||||
return MethodToolCallbackProvider.builder()
|
||||
.toolObjects(logService)
|
||||
.build();
|
||||
}
|
||||
}
|
||||
+299
@@ -0,0 +1,299 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
package org.apache.hertzbeat.mcp.server.service;
|
||||
|
||||
import com.jayway.jsonpath.JsonPath;
|
||||
import com.jayway.jsonpath.ReadContext;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.springframework.ai.tool.annotation.Tool;
|
||||
import org.springframework.ai.tool.annotation.ToolParam;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.http.HttpHeaders;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.util.LinkedMultiValueMap;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.client.RestClient;
|
||||
|
||||
import java.time.Instant;
|
||||
import java.time.LocalDateTime;
|
||||
import java.time.ZoneId;
|
||||
import java.time.format.DateTimeFormatter;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* Log query service.
|
||||
*/
|
||||
@Service
|
||||
@Slf4j
|
||||
public class LogService {
|
||||
|
||||
private static final String BASE_QUERY = "SELECT timestamp, severity_text, body FROM hzb_logs";
|
||||
private static final String TIMESTAMP_COLUMN = "timestamp";
|
||||
private static final String SEVERITY_TEXT_COLUMN = "severity_text";
|
||||
private static final String BODY_COLUMN = "body";
|
||||
private static final int DEFAULT_LIMIT = 20;
|
||||
private static final int MAX_LIMIT = 100;
|
||||
private static final int MAX_KEYWORD_LENGTH = 256;
|
||||
private static final long NANOS_PER_MILLISECOND = 1_000_000L;
|
||||
private static final Set<String> SUPPORTED_SEVERITIES =
|
||||
Set.of("TRACE", "DEBUG", "INFO", "WARN", "ERROR", "FATAL");
|
||||
private static final DateTimeFormatter DATE_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
|
||||
|
||||
private final RestClient restClient;
|
||||
private final String database;
|
||||
|
||||
/**
|
||||
* Creates a log query service.
|
||||
*
|
||||
* @param greptimeUrl GreptimeDB URL
|
||||
* @param database GreptimeDB database name
|
||||
* @param username GreptimeDB username
|
||||
* @param password GreptimeDB password
|
||||
*/
|
||||
@Autowired
|
||||
public LogService(@Value("${greptime.url}") String greptimeUrl,
|
||||
@Value("${greptime.database:public}") String database,
|
||||
@Value("${greptime.username:}") String username,
|
||||
@Value("${greptime.password:}") String password) {
|
||||
this(RestClient.builder(), greptimeUrl, database, username, password);
|
||||
}
|
||||
|
||||
LogService(RestClient.Builder restClientBuilder, String greptimeUrl, String database,
|
||||
String username, String password) {
|
||||
boolean hasUsername = username != null && !username.isBlank();
|
||||
boolean hasPassword = password != null && !password.isBlank();
|
||||
if (hasUsername != hasPassword) {
|
||||
throw new IllegalArgumentException("GreptimeDB username and password must be configured together");
|
||||
}
|
||||
|
||||
RestClient.Builder builder = restClientBuilder
|
||||
.baseUrl(greptimeUrl)
|
||||
.defaultHeader(HttpHeaders.ACCEPT, MediaType.APPLICATION_JSON_VALUE)
|
||||
.defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_FORM_URLENCODED_VALUE);
|
||||
if (hasUsername) {
|
||||
builder.defaultHeaders(headers -> headers.setBasicAuth(username, password));
|
||||
}
|
||||
this.restClient = builder.build();
|
||||
this.database = database == null || database.isBlank() ? "public" : database;
|
||||
}
|
||||
|
||||
/**
|
||||
* Queries HertzBeat system logs using structured filters.
|
||||
*
|
||||
* @param severity log severity
|
||||
* @param keyword log body keyword
|
||||
* @param startTime start time as a Unix timestamp in milliseconds
|
||||
* @param endTime end time as a Unix timestamp in milliseconds
|
||||
* @param limit maximum number of records to return
|
||||
* @return formatted log query results
|
||||
*/
|
||||
@Tool(name = "query_logs", description = "Query HertzBeat system logs with structured read-only filters")
|
||||
public String queryLogs(
|
||||
@ToolParam(required = false,
|
||||
description = "Log severity; supported values: TRACE, DEBUG, INFO, WARN, ERROR, FATAL")
|
||||
String severity,
|
||||
@ToolParam(required = false, description = "Keyword to search in the log body; maximum 256 characters")
|
||||
String keyword,
|
||||
@ToolParam(required = false, description = "Start time as a Unix timestamp in milliseconds")
|
||||
Long startTime,
|
||||
@ToolParam(required = false, description = "End time as a Unix timestamp in milliseconds")
|
||||
Long endTime,
|
||||
@ToolParam(required = false,
|
||||
description = "Maximum number of records to return; defaults to 20 and cannot exceed 100")
|
||||
Integer limit) {
|
||||
try {
|
||||
String response = executeQuery(buildQuery(severity, keyword, startTime, endTime, limit));
|
||||
return formatQueryResults(response);
|
||||
} catch (IllegalArgumentException e) {
|
||||
return "Invalid query parameters: " + e.getMessage();
|
||||
} catch (Exception e) {
|
||||
log.error("Failed to query logs", e);
|
||||
return "Failed to query logs";
|
||||
}
|
||||
}
|
||||
|
||||
static String buildQuery(String severity, String keyword, Long startTime, Long endTime, Integer limit) {
|
||||
if (startTime != null && startTime < 0) {
|
||||
throw new IllegalArgumentException("Start time must not be negative");
|
||||
}
|
||||
if (endTime != null && endTime < 0) {
|
||||
throw new IllegalArgumentException("End time must not be negative");
|
||||
}
|
||||
if (startTime != null && endTime != null && startTime > endTime) {
|
||||
throw new IllegalArgumentException("Start time must not be later than end time");
|
||||
}
|
||||
|
||||
int queryLimit = limit == null ? DEFAULT_LIMIT : limit;
|
||||
if (queryLimit < 1 || queryLimit > MAX_LIMIT) {
|
||||
throw new IllegalArgumentException("Limit must be between 1 and 100");
|
||||
}
|
||||
|
||||
List<String> conditions = new ArrayList<>(4);
|
||||
if (startTime != null) {
|
||||
conditions.add("timestamp >= " + toNanoseconds(startTime, "Start time"));
|
||||
}
|
||||
if (endTime != null) {
|
||||
conditions.add("timestamp <= " + toNanoseconds(endTime, "End time"));
|
||||
}
|
||||
|
||||
String normalizedSeverity = normalizeSeverity(severity);
|
||||
if (normalizedSeverity != null) {
|
||||
conditions.add("severity_text = '" + normalizedSeverity + "'");
|
||||
}
|
||||
|
||||
if (keyword != null && !keyword.isBlank()) {
|
||||
String normalizedKeyword = keyword.strip();
|
||||
if (normalizedKeyword.length() > MAX_KEYWORD_LENGTH) {
|
||||
throw new IllegalArgumentException("Log keyword must not exceed 256 characters");
|
||||
}
|
||||
conditions.add("matches_term(body, '" + escapeSqlLiteral(normalizedKeyword) + "')");
|
||||
}
|
||||
|
||||
StringBuilder query = new StringBuilder(BASE_QUERY);
|
||||
if (!conditions.isEmpty()) {
|
||||
query.append(" WHERE ").append(String.join(" AND ", conditions));
|
||||
}
|
||||
return query.append(" ORDER BY timestamp DESC LIMIT ").append(queryLimit).toString();
|
||||
}
|
||||
|
||||
private static String normalizeSeverity(String severity) {
|
||||
if (severity == null || severity.isBlank()) {
|
||||
return null;
|
||||
}
|
||||
String normalizedSeverity = severity.strip().toUpperCase(Locale.ROOT);
|
||||
if (!SUPPORTED_SEVERITIES.contains(normalizedSeverity)) {
|
||||
throw new IllegalArgumentException("Unsupported log severity");
|
||||
}
|
||||
return normalizedSeverity;
|
||||
}
|
||||
|
||||
private static long toNanoseconds(long epochMilliseconds, String fieldName) {
|
||||
try {
|
||||
return Math.multiplyExact(epochMilliseconds, NANOS_PER_MILLISECOND);
|
||||
} catch (ArithmeticException e) {
|
||||
throw new IllegalArgumentException(fieldName + " is out of the supported range", e);
|
||||
}
|
||||
}
|
||||
|
||||
private static String escapeSqlLiteral(String value) {
|
||||
return value.replace("'", "''");
|
||||
}
|
||||
|
||||
private String executeQuery(String sql) {
|
||||
MultiValueMap<String, String> formData = new LinkedMultiValueMap<>();
|
||||
formData.add("sql", sql);
|
||||
log.debug("Executing structured log query");
|
||||
|
||||
return restClient.post()
|
||||
.uri(uriBuilder -> uriBuilder.path("/v1/sql").queryParam("db", database).build())
|
||||
.contentType(MediaType.APPLICATION_FORM_URLENCODED)
|
||||
.body(formData)
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
}
|
||||
|
||||
private String formatQueryResults(String response) {
|
||||
ReadContext ctx = JsonPath.parse(response);
|
||||
List<Map<String, Object>> columnSchemas = ctx.read("$.output[0].records.schema.column_schemas");
|
||||
List<List<Object>> rows = ctx.read("$.output[0].records.rows");
|
||||
int totalRows = ctx.read("$.output[0].records.total_rows");
|
||||
|
||||
ColumnIndices indices = findColumnIndices(columnSchemas);
|
||||
StringBuilder result = new StringBuilder()
|
||||
.append("Query Results:\n\n")
|
||||
.append("Log Time\t\t\tLog Level\tLog Content\n")
|
||||
.append("----------------------------------------------------\n");
|
||||
|
||||
if (rows != null && !rows.isEmpty()) {
|
||||
formatRows(rows, indices, result);
|
||||
result.append("\nTotal ").append(totalRows).append(" records");
|
||||
} else {
|
||||
result.append("No data");
|
||||
}
|
||||
|
||||
return result.toString();
|
||||
}
|
||||
|
||||
private record ColumnIndices(int timestamp, int severityText, int body) {}
|
||||
|
||||
private ColumnIndices findColumnIndices(List<Map<String, Object>> columnSchemas) {
|
||||
int timestampIndex = -1;
|
||||
int severityTextIndex = -1;
|
||||
int bodyIndex = -1;
|
||||
|
||||
for (int i = 0; i < columnSchemas.size(); i++) {
|
||||
String columnName = (String) columnSchemas.get(i).get("name");
|
||||
switch (columnName) {
|
||||
case TIMESTAMP_COLUMN -> timestampIndex = i;
|
||||
case SEVERITY_TEXT_COLUMN -> severityTextIndex = i;
|
||||
case BODY_COLUMN -> bodyIndex = i;
|
||||
default -> {
|
||||
// Ignore other columns
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return new ColumnIndices(timestampIndex, severityTextIndex, bodyIndex);
|
||||
}
|
||||
|
||||
private void formatRows(List<List<Object>> rows, ColumnIndices indices, StringBuilder result) {
|
||||
for (List<Object> row : rows) {
|
||||
appendTimestamp(row, indices.timestamp(), result);
|
||||
appendSeverity(row, indices.severityText(), result);
|
||||
appendBody(row, indices.body(), result);
|
||||
result.append("\n");
|
||||
}
|
||||
}
|
||||
|
||||
private void appendTimestamp(List<Object> row, int index, StringBuilder result) {
|
||||
if (index >= 0 && index < row.size()) {
|
||||
Object value = row.get(index);
|
||||
if (value instanceof Number) {
|
||||
long timestamp = ((Number) value).longValue();
|
||||
LocalDateTime dateTime = LocalDateTime.ofInstant(
|
||||
Instant.ofEpochMilli(timestamp / 1_000_000),
|
||||
ZoneId.systemDefault());
|
||||
result.append(DATE_FORMATTER.format(dateTime)).append("\t");
|
||||
return;
|
||||
}
|
||||
}
|
||||
result.append("Unknown time\t");
|
||||
}
|
||||
|
||||
private void appendSeverity(List<Object> row, int index, StringBuilder result) {
|
||||
if (index >= 0 && index < row.size()) {
|
||||
result.append(row.get(index)).append("\t");
|
||||
} else {
|
||||
result.append("Unknown\t");
|
||||
}
|
||||
}
|
||||
|
||||
private void appendBody(List<Object> row, int index, StringBuilder result) {
|
||||
if (index >= 0 && index < row.size()) {
|
||||
result.append(row.get(index));
|
||||
} else {
|
||||
result.append("No content");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
# 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.
|
||||
|
||||
spring:
|
||||
main:
|
||||
banner-mode: off
|
||||
ai:
|
||||
mcp:
|
||||
server:
|
||||
name: hertzbeat-log-analysis-server
|
||||
version: 0.1
|
||||
|
||||
#
|
||||
#logging:
|
||||
# file:
|
||||
# name:
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
package org.apache.hertzbeat.mcp.server.service;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
|
||||
import static org.springframework.test.web.client.match.MockRestRequestMatchers.content;
|
||||
import static org.springframework.test.web.client.match.MockRestRequestMatchers.header;
|
||||
import static org.springframework.test.web.client.match.MockRestRequestMatchers.method;
|
||||
import static org.springframework.test.web.client.match.MockRestRequestMatchers.requestTo;
|
||||
import static org.springframework.test.web.client.response.MockRestResponseCreators.withServerError;
|
||||
import static org.springframework.test.web.client.response.MockRestResponseCreators.withSuccess;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.Base64;
|
||||
import java.util.Map;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.tool.ToolCallback;
|
||||
import org.springframework.ai.tool.method.MethodToolCallbackProvider;
|
||||
import org.springframework.http.HttpHeaders;
|
||||
import org.springframework.http.HttpMethod;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.util.LinkedMultiValueMap;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.client.RestClient;
|
||||
import org.springframework.test.web.client.MockRestServiceServer;
|
||||
|
||||
/**
|
||||
* Security tests for {@link LogService} queries.
|
||||
*/
|
||||
class LogServiceTest {
|
||||
|
||||
private static final String BASE_QUERY = "SELECT timestamp, severity_text, body FROM hzb_logs";
|
||||
|
||||
@Test
|
||||
void shouldBuildDefaultReadOnlyQuery() {
|
||||
assertThat(LogService.buildQuery(null, null, null, null, null))
|
||||
.isEqualTo(BASE_QUERY + " ORDER BY timestamp DESC LIMIT 20");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldExposeOnlyStructuredToolParameters() {
|
||||
LogService service = new LogService(
|
||||
RestClient.builder(), "http://localhost:4000", "public", "", "");
|
||||
ToolCallback[] callbacks = MethodToolCallbackProvider.builder().toolObjects(service).build().getToolCallbacks();
|
||||
|
||||
assertThat(callbacks).singleElement().satisfies(callback -> {
|
||||
assertThat(callback.getToolDefinition().name()).isEqualTo("query_logs");
|
||||
assertThat(callback.getToolDefinition().inputSchema())
|
||||
.contains("severity", "keyword", "startTime", "endTime", "limit")
|
||||
.doesNotContain("querySql");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldBuildQueryFromValidatedFilters() {
|
||||
String query = LogService.buildQuery(
|
||||
" error ", " x'); DELETE FROM hzb_logs; -- ", 1L, 2L, 10);
|
||||
|
||||
assertThat(query).isEqualTo(BASE_QUERY
|
||||
+ " WHERE timestamp >= 1000000"
|
||||
+ " AND timestamp <= 2000000"
|
||||
+ " AND severity_text = 'ERROR'"
|
||||
+ " AND matches_term(body, 'x''); DELETE FROM hzb_logs; --')"
|
||||
+ " ORDER BY timestamp DESC LIMIT 10");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRejectInvalidFilters() {
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery("NOTICE", null, null, null, null));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, null, -1L, null, null));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, null, null, -1L, null));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, null, 2L, 1L, null));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, null, Long.MAX_VALUE, null, null));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, "x".repeat(257), null, null, null));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, null, null, null, 0));
|
||||
assertThatIllegalArgumentException()
|
||||
.isThrownBy(() -> LogService.buildQuery(null, null, null, null, 101));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldSendGeneratedSelectWithBasicAuthentication() {
|
||||
RestClient.Builder restClientBuilder = RestClient.builder();
|
||||
MockRestServiceServer server = MockRestServiceServer.bindTo(restClientBuilder).build();
|
||||
LogService service = new LogService(
|
||||
restClientBuilder, "http://localhost:4000", "observability", "reader", "secret");
|
||||
String query = BASE_QUERY + " ORDER BY timestamp DESC LIMIT 20";
|
||||
MultiValueMap<String, String> expectedForm = new LinkedMultiValueMap<>();
|
||||
expectedForm.setAll(Map.of("sql", query));
|
||||
String credentials = Base64.getEncoder()
|
||||
.encodeToString("reader:secret".getBytes(StandardCharsets.UTF_8));
|
||||
|
||||
server.expect(requestTo("http://localhost:4000/v1/sql?db=observability"))
|
||||
.andExpect(method(HttpMethod.POST))
|
||||
.andExpect(header(HttpHeaders.AUTHORIZATION, "Basic " + credentials))
|
||||
.andExpect(content().formData(expectedForm))
|
||||
.andRespond(withSuccess(emptyResult(), MediaType.APPLICATION_JSON));
|
||||
|
||||
assertThat(service.queryLogs(null, null, null, null, null)).contains("No data");
|
||||
server.verify();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldFormatReturnedLogs() {
|
||||
RestClient.Builder restClientBuilder = RestClient.builder();
|
||||
MockRestServiceServer server = MockRestServiceServer.bindTo(restClientBuilder).build();
|
||||
LogService service = new LogService(restClientBuilder, "http://localhost:4000", "public", "", "");
|
||||
server.expect(requestTo("http://localhost:4000/v1/sql?db=public"))
|
||||
.andRespond(withSuccess(resultWithOneLog(), MediaType.APPLICATION_JSON));
|
||||
|
||||
assertThat(service.queryLogs("ERROR", "failure", null, null, 1))
|
||||
.contains("ERROR", "failure", "Total 1 records");
|
||||
server.verify();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldReturnSafeErrorMessages() {
|
||||
RestClient.Builder restClientBuilder = RestClient.builder();
|
||||
MockRestServiceServer server = MockRestServiceServer.bindTo(restClientBuilder).build();
|
||||
LogService service = new LogService(restClientBuilder, "http://localhost:4000", "public", "", "");
|
||||
|
||||
assertThat(service.queryLogs("NOTICE", null, null, null, null))
|
||||
.isEqualTo("Invalid query parameters: Unsupported log severity");
|
||||
|
||||
server.expect(requestTo("http://localhost:4000/v1/sql?db=public"))
|
||||
.andRespond(withServerError());
|
||||
assertThat(service.queryLogs(null, null, null, null, null)).isEqualTo("Failed to query logs");
|
||||
server.verify();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldRequireCompleteCredentials() {
|
||||
assertThatIllegalArgumentException().isThrownBy(() -> new LogService(
|
||||
RestClient.builder(), "http://localhost:4000", "public", "reader", ""));
|
||||
assertThatIllegalArgumentException().isThrownBy(() -> new LogService(
|
||||
RestClient.builder(), "http://localhost:4000", "public", "", "secret"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldCreateWithProductionConstructor() {
|
||||
assertThat(new LogService("http://localhost:4000", "public", "", "")).isNotNull();
|
||||
}
|
||||
|
||||
private static String emptyResult() {
|
||||
return """
|
||||
{
|
||||
"output": [{
|
||||
"records": {
|
||||
"schema": {"column_schemas": [
|
||||
{"name": "timestamp"},
|
||||
{"name": "severity_text"},
|
||||
{"name": "body"}
|
||||
]},
|
||||
"rows": [],
|
||||
"total_rows": 0
|
||||
}
|
||||
}]
|
||||
}
|
||||
""";
|
||||
}
|
||||
|
||||
private static String resultWithOneLog() {
|
||||
return """
|
||||
{
|
||||
"output": [{
|
||||
"records": {
|
||||
"schema": {"column_schemas": [
|
||||
{"name": "timestamp"},
|
||||
{"name": "severity_text"},
|
||||
{"name": "body"}
|
||||
]},
|
||||
"rows": [[0, "ERROR", "failure"]],
|
||||
"total_rows": 1
|
||||
}
|
||||
}]
|
||||
}
|
||||
""";
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user