mirror of
https://github.com/dromara/hertzbeat.git
synced 2026-09-17 09:40:58 +00:00
maintenance: validate MCP OAuth exchanges (#4279)
Co-authored-by: Duansg <siguoduan@gmail.com>
This commit is contained in:
Logic
Duansg
@@ -44,10 +44,12 @@ jobs:
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@4be7066ada62dd38de10e7b70166bc74ed198c30 # stable
|
||||
with:
|
||||
toolchain: ${{ env.RUST_VERSION }}
|
||||
components: rustfmt, clippy
|
||||
run: |
|
||||
rustup toolchain install "$RUST_VERSION" \
|
||||
--profile minimal \
|
||||
--component rustfmt \
|
||||
--component clippy
|
||||
rustup default "$RUST_VERSION"
|
||||
|
||||
- name: Cache cargo registry
|
||||
uses: actions/cache@v4
|
||||
|
||||
Generated
+5
-1
@@ -1062,15 +1062,18 @@ dependencies = [
|
||||
"askama",
|
||||
"axum 0.8.4",
|
||||
"axum-test",
|
||||
"base64",
|
||||
"chrono",
|
||||
"hyper",
|
||||
"oauth2",
|
||||
"ipnet",
|
||||
"rand 0.8.5",
|
||||
"regex",
|
||||
"rmcp",
|
||||
"serde",
|
||||
"serde_urlencoded",
|
||||
"sha2",
|
||||
"shlex",
|
||||
"subtle",
|
||||
"tempfile",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
@@ -1082,6 +1085,7 @@ dependencies = [
|
||||
"tracing",
|
||||
"tracing-appender",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
|
||||
@@ -36,13 +36,17 @@ axum = { version = "0.8", features = ["macros"] }
|
||||
chrono = "0.4"
|
||||
tower-http = { version = "0.6", features = ["cors"] }
|
||||
askama = { version = "0.14"}
|
||||
base64 = "0.22"
|
||||
rand = { version = "0.8", features = ["std"] }
|
||||
uuid = { version = "1.6", features = ["v4", "serde"] }
|
||||
serde_urlencoded = "0.7"
|
||||
oauth2 = "5.0"
|
||||
sha2 = "0.10"
|
||||
toml = "0.8"
|
||||
regex = "1.11.1"
|
||||
shlex = "1.3"
|
||||
subtle = "2.6"
|
||||
url = "2.5"
|
||||
ipnet = "2.11"
|
||||
|
||||
[dev-dependencies]
|
||||
tokio-test = "0.4"
|
||||
|
||||
@@ -18,12 +18,42 @@ Version `1.88.0` can absolutely work, and we recommend using the latest version
|
||||
|
||||
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
|
||||
```shell
|
||||
export MCP_OAUTH_APPROVAL_SECRET="<at-least-32-random-characters>"
|
||||
export MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com"
|
||||
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.
|
||||
|
||||
Production mode requires `MCP_OAUTH_APPROVAL_SECRET`. The operator enters this
|
||||
secret on the OAuth approval page; it is separate from dynamically registered
|
||||
client credentials. Production also requires an explicit HTTPS
|
||||
`MCP_OAUTH_PUBLIC_BASE_URL`; OAuth metadata never derives public endpoints from
|
||||
the request `Host` header. Development mode generates a temporary approval
|
||||
secret and may derive a loopback HTTP URL from the local bind address.
|
||||
|
||||
Dynamic registration supports public clients (`token_endpoint_auth_method:
|
||||
none`) and confidential clients (`client_secret_post`). Authorization requests
|
||||
must use PKCE S256. Authorization transactions and codes are one-time and
|
||||
short-lived; access tokens expire after one hour, and refresh tokens expire
|
||||
after one day and rotate on every use. A client that holds a refresh token
|
||||
remains registered for at least that token's full lifetime. Open client
|
||||
registration is limited to 16 successful registrations per minute for each
|
||||
TCP peer, so one source cannot consume every caller's admission window. The
|
||||
rate-limit source table is bounded. Deployments behind a reverse proxy can set
|
||||
`MCP_OAUTH_TRUSTED_PROXY_CIDRS` to a comma-separated list of the exact proxy
|
||||
networks. Only connections from those networks may supply `X-Forwarded-For`;
|
||||
untrusted peers cannot spoof the limiter identity, and `/0` networks are
|
||||
rejected. The proxy must overwrite or safely append that header. An unused
|
||||
registered client expires after one hour; expired clients are removed before
|
||||
the 1,024-client capacity check.
|
||||
When anonymous registrations fill the remaining capacity, the oldest client
|
||||
that has never received a refresh token is reclaimed; clients with live
|
||||
refresh credentials are not evicted. A client that receives `invalid_client`
|
||||
after an unused registration expires must dynamically register again. OAuth
|
||||
form and JSON bodies are limited to 16 KiB.
|
||||
|
||||
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.
|
||||
@@ -47,23 +77,38 @@ docker build --build-arg HTTPS_PROXY=<your https_proxy> --build-arg HTTP_PROXY=<
|
||||
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
|
||||
docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \
|
||||
-e MCP_OAUTH_APPROVAL_SECRET="<at-least-32-random-characters>" \
|
||||
-e MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" \
|
||||
--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.
|
||||
The MCP Server inside the container runs on 0.0.0.0:4000, while the example
|
||||
publishes it only on the host loopback interface. On the host machine, use the
|
||||
inspector with URL `http://localhost:4000/mcp` to connect to the MCP Server
|
||||
inside the container. Remote deployments must terminate TLS before forwarding
|
||||
OAuth endpoints to 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
|
||||
docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \
|
||||
-e MCP_OAUTH_APPROVAL_SECRET="<at-least-32-random-characters>" \
|
||||
-e MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" \
|
||||
-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
|
||||
docker run -d --name mcp-bash-server -p 127.0.0.1:4000:4000 \
|
||||
-e MCP_OAUTH_APPROVAL_SECRET="<at-least-32-random-characters>" \
|
||||
-e MCP_OAUTH_PUBLIC_BASE_URL="https://mcp.example.com" \
|
||||
-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
|
||||
@@ -98,14 +143,9 @@ Start the MCP Server in daemon mode, then add the settings to your Vscode Copilo
|
||||
}
|
||||
```
|
||||
|
||||
**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.
|
||||
OAuth-capable MCP clients can discover the authorization metadata, dynamically
|
||||
register a public client, and complete the PKCE flow. A manually configured
|
||||
bearer token remains available for clients that do not implement MCP OAuth.
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -231,7 +271,8 @@ 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
|
||||
3. Enter the server operator's `MCP_OAUTH_APPROVAL_SECRET`, then 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
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -31,7 +31,7 @@
|
||||
</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>
|
||||
<p>This MCP server uses an OAuth 2.0 authorization-code flow with PKCE.</p>
|
||||
|
||||
<h2>Available Endpoints:</h2>
|
||||
|
||||
@@ -41,10 +41,12 @@
|
||||
<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>
|
||||
<li><code>client_id</code> - Dynamically registered client identifier</li>
|
||||
<li><code>redirect_uri</code> - Exact registered redirect URI</li>
|
||||
<li><code>scope</code> - Optional supported scopes</li>
|
||||
<li><code>state</code> - Recommended client transaction state</li>
|
||||
<li><code>code_challenge</code> - PKCE S256 challenge</li>
|
||||
<li><code>code_challenge_method</code> - Must be "S256"</li>
|
||||
</ul>
|
||||
</div>
|
||||
|
||||
@@ -53,11 +55,13 @@
|
||||
<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>grant_type</code> - "authorization_code" or "refresh_token"</li>
|
||||
<li><code>code</code> - One-time authorization code</li>
|
||||
<li><code>client_id</code> - Client identifier</li>
|
||||
<li><code>client_secret</code> - Client secret</li>
|
||||
<li><code>client_secret</code> - Required only for confidential clients</li>
|
||||
<li><code>redirect_uri</code> - Redirect URI used in authorization request</li>
|
||||
<li><code>code_verifier</code> - PKCE verifier for authorization-code exchange</li>
|
||||
<li><code>refresh_token</code> - Rotating token for refresh grants</li>
|
||||
</ul>
|
||||
</div>
|
||||
|
||||
@@ -69,13 +73,12 @@
|
||||
<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>
|
||||
<li>The MCP client discovers metadata and dynamically registers.</li>
|
||||
<li>The client starts authorization with an S256 PKCE challenge.</li>
|
||||
<li>The resource owner authenticates and approves the bound transaction.</li>
|
||||
<li>The server returns a short-lived, one-time authorization code.</li>
|
||||
<li>The client exchanges the code and PKCE verifier for expiring tokens.</li>
|
||||
<li>The client sends the bearer access token to the MCP endpoint.</li>
|
||||
</ol>
|
||||
</div>
|
||||
</body>
|
||||
|
||||
@@ -1,42 +1,33 @@
|
||||
/*
|
||||
* 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
|
||||
* 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.
|
||||
* 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
|
||||
//! MCP Bash Server.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::{net::SocketAddr, sync::Arc};
|
||||
|
||||
use anyhow::Result;
|
||||
use anyhow::{Context, Result, bail};
|
||||
use axum::{
|
||||
Router,
|
||||
body::Body,
|
||||
http::{HeaderMap, Request},
|
||||
http::Request,
|
||||
middleware::{self, Next},
|
||||
response::{Html, IntoResponse, Response},
|
||||
response::{Html, Response},
|
||||
routing::{get, post},
|
||||
};
|
||||
use rmcp::transport::streamable_http_server::{
|
||||
@@ -46,81 +37,73 @@ 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,
|
||||
McpOAuthStore, generate_random_string, oauth_approve, oauth_authorization_server,
|
||||
oauth_authorize, oauth_register, oauth_token, parse_trusted_proxy_cidrs,
|
||||
validate_public_base_url, 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
|
||||
/// Log request metadata without reading form bodies or emitting credentials.
|
||||
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}"));
|
||||
for (key, value) in &headers {
|
||||
let value = if key == "authorization" || key == "cookie" {
|
||||
"<redacted>"
|
||||
} else {
|
||||
value.to_str().unwrap_or("<binary>")
|
||||
};
|
||||
header_log.push_str(&format!("\n {key}: {value}"));
|
||||
}
|
||||
|
||||
// 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
|
||||
info!("REQUEST: {method} {uri} {version:?}{header_log}");
|
||||
let response = next.run(request).await;
|
||||
|
||||
// Log response status
|
||||
let status = response.status();
|
||||
info!("RESPONSE: {} for {} {}", status, method, uri);
|
||||
|
||||
info!("RESPONSE: {} for {} {}", response.status(), method, uri);
|
||||
response
|
||||
}
|
||||
|
||||
/// Main application entry point
|
||||
/// Sets up logging, OAuth store, HTTP server, and starts the MCP bash server
|
||||
fn approval_secret_for_mode(is_dev: bool, configured: Option<String>) -> Result<String> {
|
||||
if is_dev {
|
||||
return Ok(configured.unwrap_or_else(|| generate_random_string(32)));
|
||||
}
|
||||
let approval_secret = configured
|
||||
.context("MCP_OAUTH_APPROVAL_SECRET must be set when the server runs in production mode")?;
|
||||
if approval_secret.len() < 32 {
|
||||
bail!("MCP_OAUTH_APPROVAL_SECRET must contain at least 32 characters");
|
||||
}
|
||||
Ok(approval_secret)
|
||||
}
|
||||
|
||||
fn public_base_url_for_mode(
|
||||
is_dev: bool,
|
||||
configured: Option<String>,
|
||||
bind_address: &str,
|
||||
) -> Result<url::Url> {
|
||||
if let Some(configured) = configured {
|
||||
return validate_public_base_url(&configured, !is_dev)
|
||||
.map_err(|message| anyhow::anyhow!(message));
|
||||
}
|
||||
if !is_dev {
|
||||
bail!("MCP_OAUTH_PUBLIC_BASE_URL must be set when the server runs in production mode");
|
||||
}
|
||||
let local_address = bind_address.replacen("0.0.0.0", "127.0.0.1", 1);
|
||||
validate_public_base_url(&format!("http://{local_address}"), false)
|
||||
.map_err(|message| anyhow::anyhow!(message))
|
||||
}
|
||||
|
||||
#[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);
|
||||
@@ -133,7 +116,6 @@ async fn main() -> Result<()> {
|
||||
.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
|
||||
@@ -141,27 +123,33 @@ async fn main() -> Result<()> {
|
||||
.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 approval_secret =
|
||||
approval_secret_for_mode(is_dev, std::env::var("MCP_OAUTH_APPROVAL_SECRET").ok())?;
|
||||
let public_base_url = public_base_url_for_mode(
|
||||
is_dev,
|
||||
std::env::var("MCP_OAUTH_PUBLIC_BASE_URL").ok(),
|
||||
&bind_address,
|
||||
)?;
|
||||
let trusted_proxy_config = std::env::var("MCP_OAUTH_TRUSTED_PROXY_CIDRS").ok();
|
||||
let trusted_proxies = parse_trusted_proxy_cidrs(trusted_proxy_config.as_deref())
|
||||
.map_err(|message| anyhow::anyhow!(message))?;
|
||||
let oauth_store = Arc::new(McpOAuthStore::with_trusted_proxies(
|
||||
approval_secret,
|
||||
public_base_url,
|
||||
trusted_proxies,
|
||||
));
|
||||
|
||||
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 {
|
||||
@@ -171,448 +159,92 @@ async fn main() -> Result<()> {
|
||||
))
|
||||
};
|
||||
|
||||
// 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),
|
||||
get(oauth_authorization_server).options(oauth_authorization_server),
|
||||
)
|
||||
.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(oauth_server_router)
|
||||
.merge(protected_server_router)
|
||||
.with_state(oauth_store.clone())
|
||||
.with_state(oauth_store)
|
||||
.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;
|
||||
|
||||
let _ = axum::serve(
|
||||
listener,
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.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() {
|
||||
async fn index_handler_returns_oauth_page() {
|
||||
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();
|
||||
assert_eq!(response.0, INDEX_HTML);
|
||||
assert!(response.0.contains("OAuth"));
|
||||
}
|
||||
|
||||
#[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());
|
||||
fn production_requires_strong_approval_secret() {
|
||||
assert!(approval_secret_for_mode(false, None).is_err());
|
||||
assert!(approval_secret_for_mode(false, Some("too-short".to_string())).is_err());
|
||||
let configured = "resource-owner-secret-with-32-characters".to_string();
|
||||
assert_eq!(
|
||||
no_auth_token_result.unwrap_err(),
|
||||
"No third-party token available for session"
|
||||
approval_secret_for_mode(false, Some(configured.clone())).unwrap(),
|
||||
configured
|
||||
);
|
||||
|
||||
// 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(),
|
||||
#[test]
|
||||
fn production_requires_explicit_https_public_base_url() {
|
||||
assert!(public_base_url_for_mode(false, None, "0.0.0.0:4000").is_err());
|
||||
assert!(
|
||||
public_base_url_for_mode(
|
||||
false,
|
||||
Some("http://mcp.example".to_string()),
|
||||
"0.0.0.0:4000"
|
||||
)
|
||||
.await;
|
||||
|
||||
let session2_id = store
|
||||
.create_auth_session(
|
||||
"mcp-client".to_string(),
|
||||
Some("profile".to_string()),
|
||||
None,
|
||||
"security-test-2".to_string(),
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
public_base_url_for_mode(
|
||||
false,
|
||||
Some("https://mcp.example".to_string()),
|
||||
"0.0.0.0:4000"
|
||||
)
|
||||
.await;
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
"https://mcp.example/"
|
||||
);
|
||||
}
|
||||
|
||||
// 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());
|
||||
#[test]
|
||||
fn development_uses_safe_loopback_metadata_address() {
|
||||
assert_eq!(
|
||||
public_base_url_for_mode(true, None, "0.0.0.0:4000")
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
"http://127.0.0.1:4000/"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +72,25 @@
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
.approval-secret {
|
||||
margin-bottom: 1.5rem;
|
||||
}
|
||||
|
||||
.approval-secret label {
|
||||
display: block;
|
||||
font-weight: 600;
|
||||
margin-bottom: 0.5rem;
|
||||
}
|
||||
|
||||
.approval-secret input {
|
||||
box-sizing: border-box;
|
||||
width: 100%;
|
||||
padding: 0.75rem;
|
||||
border: 1px solid var(--border-color);
|
||||
border-radius: 6px;
|
||||
font-size: 1rem;
|
||||
}
|
||||
|
||||
.btn {
|
||||
padding: 0.75rem 1.5rem;
|
||||
border-radius: 6px;
|
||||
@@ -110,11 +129,15 @@
|
||||
</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 }}">
|
||||
|
||||
<input type="hidden" name="transaction_id" value="{{ transaction_id }}">
|
||||
<input type="hidden" name="consent_nonce" value="{{ consent_nonce }}">
|
||||
|
||||
<div class="approval-secret">
|
||||
<label for="approval_secret">Approval secret</label>
|
||||
<input id="approval_secret" name="approval_secret" type="password"
|
||||
minlength="32" autocomplete="current-password" required>
|
||||
</div>
|
||||
|
||||
<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>
|
||||
@@ -122,4 +145,4 @@
|
||||
</form>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
</html>
|
||||
|
||||
Reference in New Issue
Block a user