feat!: remove SSE transport support (#562)
SSE transport has been removed from the MCP specification in favor of streamable HTTP. This removes all SSE-specific transport code: - Remove `transport-sse-client` and `transport-sse-server` features - Remove `SseClientTransport` and `SseServer` types - Remove SSE-specific examples (`counter_sse`, `counter_sse_directly`) - Migrate auth examples from SSE to streamable HTTP - Update tests to remove SSE transport usage - Update documentation BREAKING CHANGE: The following have been removed: - `transport-sse-client` feature - `transport-sse-client-reqwest` feature - `transport-sse-server` feature - `SseClientTransport` type - `SseServer` type - `sse_client` and `sse_server` modules Users should migrate to streamable HTTP transport which provides equivalent functionality. See `StreamableHttpClientTransport` and `StreamableHttpService` for the replacement APIs. Ref: https://github.com/modelcontextprotocol/rust-sdk/pull/561#issuecomment-3576551699
This commit is contained in:
adam jones
GitHub
parent
3dee024325
commit
eb5a7f7408
6 changed files with 164 additions and 1688 deletions
705
examples/servers/src/complex_auth_streamhttp.rs
Normal file
705
examples/servers/src/complex_auth_streamhttp.rs
Normal file
|
|
@ -0,0 +1,705 @@
|
|||
use std::{collections::HashMap, net::SocketAddr, sync::Arc};
|
||||
|
||||
use anyhow::Result;
|
||||
use askama::Template;
|
||||
use axum::{
|
||||
Json, Router,
|
||||
body::Body,
|
||||
extract::{Form, Query, State},
|
||||
http::{Request, StatusCode},
|
||||
middleware::{self, Next},
|
||||
response::{Html, IntoResponse, Redirect, Response},
|
||||
routing::{get, post},
|
||||
};
|
||||
use rand::{Rng, distr::Alphanumeric};
|
||||
use rmcp::transport::{
|
||||
StreamableHttpServerConfig,
|
||||
auth::{AuthorizationMetadata, ClientRegistrationResponse, OAuthClientConfig},
|
||||
streamable_http_server::{session::local::LocalSessionManager, tower::StreamableHttpService},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::RwLock;
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tracing::{debug, error, info, warn};
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||
use uuid::Uuid;
|
||||
// Import Counter tool for MCP service
|
||||
mod common;
|
||||
use common::counter::Counter;
|
||||
|
||||
const BIND_ADDRESS: &str = "127.0.0.1:3000";
|
||||
const INDEX_HTML: &str = include_str!("html/mcp_oauth_index.html");
|
||||
|
||||
// Local registration request - only uses fields needed for this demo
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct LocalClientRegistrationRequest {
|
||||
client_name: String,
|
||||
redirect_uris: Vec<String>,
|
||||
}
|
||||
|
||||
// A easy way to manage MCP OAuth Store for managing tokens and sessions
|
||||
#[derive(Clone, Debug)]
|
||||
struct McpOAuthStore {
|
||||
clients: Arc<RwLock<HashMap<String, OAuthClientConfig>>>,
|
||||
auth_sessions: Arc<RwLock<HashMap<String, AuthSession>>>,
|
||||
access_tokens: Arc<RwLock<HashMap<String, McpAccessToken>>>,
|
||||
}
|
||||
|
||||
impl McpOAuthStore {
|
||||
fn new() -> Self {
|
||||
let mut clients = HashMap::new();
|
||||
clients.insert(
|
||||
"mcp-client".to_string(),
|
||||
OAuthClientConfig {
|
||||
client_id: "mcp-client".to_string(),
|
||||
client_secret: Some("mcp-client-secret".to_string()),
|
||||
scopes: vec!["profile".to_string(), "email".to_string()],
|
||||
redirect_uri: "http://localhost:8080/callback".to_string(),
|
||||
},
|
||||
);
|
||||
|
||||
Self {
|
||||
clients: Arc::new(RwLock::new(clients)),
|
||||
auth_sessions: Arc::new(RwLock::new(HashMap::new())),
|
||||
access_tokens: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn validate_client(
|
||||
&self,
|
||||
client_id: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Option<OAuthClientConfig> {
|
||||
let clients = self.clients.read().await;
|
||||
if let Some(client) = clients.get(client_id) {
|
||||
if client.redirect_uri.contains(&redirect_uri.to_string()) {
|
||||
return Some(client.clone());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
async fn create_auth_session(
|
||||
&self,
|
||||
client_id: String,
|
||||
scope: Option<String>,
|
||||
state: Option<String>,
|
||||
session_id: String,
|
||||
) -> String {
|
||||
let session = AuthSession {
|
||||
client_id,
|
||||
scope,
|
||||
_state: state,
|
||||
_created_at: chrono::Utc::now(),
|
||||
auth_token: None,
|
||||
};
|
||||
|
||||
self.auth_sessions
|
||||
.write()
|
||||
.await
|
||||
.insert(session_id.clone(), session);
|
||||
session_id
|
||||
}
|
||||
|
||||
async fn update_auth_session_token(
|
||||
&self,
|
||||
session_id: &str,
|
||||
token: AuthToken,
|
||||
) -> Result<(), String> {
|
||||
let mut sessions = self.auth_sessions.write().await;
|
||||
if let Some(session) = sessions.get_mut(session_id) {
|
||||
session.auth_token = Some(token);
|
||||
Ok(())
|
||||
} else {
|
||||
Err("Session not found".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_mcp_token(&self, session_id: &str) -> Result<McpAccessToken, String> {
|
||||
let sessions = self.auth_sessions.read().await;
|
||||
if let Some(session) = sessions.get(session_id) {
|
||||
if let Some(auth_token) = &session.auth_token {
|
||||
let access_token = format!("mcp-token-{}", Uuid::new_v4());
|
||||
let token = McpAccessToken {
|
||||
access_token: access_token.clone(),
|
||||
token_type: "Bearer".to_string(),
|
||||
expires_in: 3600,
|
||||
refresh_token: format!("mcp-refresh-{}", Uuid::new_v4()),
|
||||
scope: session.scope.clone(),
|
||||
auth_token: auth_token.clone(),
|
||||
client_id: session.client_id.clone(),
|
||||
};
|
||||
|
||||
self.access_tokens
|
||||
.write()
|
||||
.await
|
||||
.insert(access_token.clone(), token.clone());
|
||||
Ok(token)
|
||||
} else {
|
||||
Err("No third-party token available for session".to_string())
|
||||
}
|
||||
} else {
|
||||
Err("Session not found".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
async fn validate_token(&self, token: &str) -> Option<McpAccessToken> {
|
||||
self.access_tokens.read().await.get(token).cloned()
|
||||
}
|
||||
}
|
||||
|
||||
// a simple session record for auth session
|
||||
#[derive(Clone, Debug)]
|
||||
struct AuthSession {
|
||||
client_id: String,
|
||||
scope: Option<String>,
|
||||
_state: Option<String>,
|
||||
_created_at: chrono::DateTime<chrono::Utc>,
|
||||
auth_token: Option<AuthToken>,
|
||||
}
|
||||
|
||||
// a simple token record for auth token
|
||||
// not used oauth2 token for avoid include oauth2 crate in this example
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
struct AuthToken {
|
||||
access_token: String,
|
||||
token_type: String,
|
||||
expires_in: u64,
|
||||
refresh_token: String,
|
||||
scope: Option<String>,
|
||||
}
|
||||
|
||||
// a simple token record for mcp token ,
|
||||
// not used oauth2 token for avoid include oauth2 crate in this example
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
struct McpAccessToken {
|
||||
access_token: String,
|
||||
token_type: String,
|
||||
expires_in: u64,
|
||||
refresh_token: String,
|
||||
scope: Option<String>,
|
||||
auth_token: AuthToken,
|
||||
client_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AuthorizeQuery {
|
||||
#[allow(dead_code)]
|
||||
response_type: String,
|
||||
client_id: String,
|
||||
redirect_uri: String,
|
||||
scope: Option<String>,
|
||||
state: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct TokenRequest {
|
||||
grant_type: String,
|
||||
#[serde(default)]
|
||||
code: String,
|
||||
#[serde(default)]
|
||||
client_id: String,
|
||||
#[serde(default)]
|
||||
client_secret: String,
|
||||
#[serde(default)]
|
||||
redirect_uri: String,
|
||||
#[serde(default)]
|
||||
code_verifier: Option<String>,
|
||||
#[serde(default)]
|
||||
refresh_token: String,
|
||||
}
|
||||
|
||||
fn generate_random_string(length: usize) -> String {
|
||||
rand::rng()
|
||||
.sample_iter(&Alphanumeric)
|
||||
.take(length)
|
||||
.map(char::from)
|
||||
.collect()
|
||||
}
|
||||
|
||||
// Root path handler
|
||||
async fn index() -> Html<&'static str> {
|
||||
Html(INDEX_HTML)
|
||||
}
|
||||
|
||||
#[derive(Template)]
|
||||
#[template(path = "mcp_oauth_authorize.html")]
|
||||
struct OAuthAuthorizeTemplate {
|
||||
client_id: String,
|
||||
redirect_uri: String,
|
||||
scope: String,
|
||||
state: String,
|
||||
scopes: String,
|
||||
}
|
||||
|
||||
// Initial OAuth authorize endpoint
|
||||
async fn oauth_authorize(
|
||||
Query(params): Query<AuthorizeQuery>,
|
||||
State(state): State<Arc<McpOAuthStore>>,
|
||||
) -> impl IntoResponse {
|
||||
debug!("doing oauth_authorize");
|
||||
if let Some(_client) = state
|
||||
.validate_client(¶ms.client_id, ¶ms.redirect_uri)
|
||||
.await
|
||||
{
|
||||
let template = OAuthAuthorizeTemplate {
|
||||
client_id: params.client_id,
|
||||
redirect_uri: params.redirect_uri,
|
||||
scope: params.scope.clone().unwrap_or_default(),
|
||||
state: params.state.clone().unwrap_or_default(),
|
||||
scopes: params
|
||||
.scope
|
||||
.clone()
|
||||
.unwrap_or_else(|| "Basic scope".to_string()),
|
||||
};
|
||||
|
||||
Html(template.render().unwrap()).into_response()
|
||||
} else {
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "invalid_request",
|
||||
"error_description": "invalid client id or redirect uri"
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
// handle approval of authorization
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ApprovalForm {
|
||||
client_id: String,
|
||||
redirect_uri: String,
|
||||
scope: String,
|
||||
state: String,
|
||||
approved: String,
|
||||
}
|
||||
|
||||
async fn oauth_approve(
|
||||
State(state): State<Arc<McpOAuthStore>>,
|
||||
Form(form): Form<ApprovalForm>,
|
||||
) -> impl IntoResponse {
|
||||
if form.approved != "true" {
|
||||
// user rejected the authorization request
|
||||
let redirect_url = format!(
|
||||
"{}?error=access_denied&error_description={}{}",
|
||||
form.redirect_uri,
|
||||
"user rejected the authorization request",
|
||||
if form.state.is_empty() {
|
||||
"".to_string()
|
||||
} else {
|
||||
format!("&state={}", form.state)
|
||||
}
|
||||
);
|
||||
return Redirect::to(&redirect_url).into_response();
|
||||
}
|
||||
|
||||
// user approved the authorization request, generate authorization code
|
||||
let session_id = Uuid::new_v4().to_string();
|
||||
let auth_code = format!("mcp-code-{}", session_id);
|
||||
|
||||
// create new session record authorization information
|
||||
let session_id = state
|
||||
.create_auth_session(
|
||||
form.client_id.clone(),
|
||||
Some(form.scope.clone()),
|
||||
Some(form.state.clone()),
|
||||
session_id.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
// create token
|
||||
let created_token = AuthToken {
|
||||
access_token: format!("tp-token-{}", Uuid::new_v4()),
|
||||
token_type: "Bearer".to_string(),
|
||||
expires_in: 3600,
|
||||
refresh_token: format!("tp-refresh-{}", Uuid::new_v4()),
|
||||
scope: Some(form.scope),
|
||||
};
|
||||
|
||||
// update session token
|
||||
if let Err(e) = state
|
||||
.update_auth_session_token(&session_id, created_token)
|
||||
.await
|
||||
{
|
||||
error!("Failed to update session token: {}", e);
|
||||
}
|
||||
|
||||
// redirect back to client, with authorization code
|
||||
let redirect_url = format!(
|
||||
"{}?code={}{}",
|
||||
form.redirect_uri,
|
||||
auth_code,
|
||||
if form.state.is_empty() {
|
||||
"".to_string()
|
||||
} else {
|
||||
format!("&state={}", form.state)
|
||||
}
|
||||
);
|
||||
|
||||
info!("authorization approved, redirecting to: {}", redirect_url);
|
||||
Redirect::to(&redirect_url).into_response()
|
||||
}
|
||||
|
||||
// Handle token request from the MCP client
|
||||
async fn oauth_token(
|
||||
State(state): State<Arc<McpOAuthStore>>,
|
||||
request: axum::http::Request<Body>,
|
||||
) -> impl IntoResponse {
|
||||
info!("Received token request");
|
||||
|
||||
let bytes = match axum::body::to_bytes(request.into_body(), usize::MAX).await {
|
||||
Ok(bytes) => bytes,
|
||||
Err(e) => {
|
||||
error!("can't read request body: {}", e);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "invalid_request",
|
||||
"error_description": "can't read request body"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
let body_str = String::from_utf8_lossy(&bytes);
|
||||
info!("request body: {}", body_str);
|
||||
|
||||
let token_req = match serde_urlencoded::from_bytes::<TokenRequest>(&bytes) {
|
||||
Ok(form) => {
|
||||
info!("successfully parsed form data: {:?}", form);
|
||||
form
|
||||
}
|
||||
Err(e) => {
|
||||
error!("can't parse form data: {}", e);
|
||||
return (
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
Json(serde_json::json!({
|
||||
"error": "invalid_request",
|
||||
"error_description": format!("can't parse form data: {}", e)
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
if token_req.grant_type == "refresh_token" {
|
||||
warn!("this easy server only support authorization_code now");
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "unsupported_grant_type",
|
||||
"error_description": "only authorization_code is supported"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
if token_req.grant_type != "authorization_code" {
|
||||
info!("unsupported grant type: {}", token_req.grant_type);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "unsupported_grant_type",
|
||||
"error_description": "only authorization_code is supported"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
// get session_id from code
|
||||
if !token_req.code.starts_with("mcp-code-") {
|
||||
info!("invalid authorization code: {}", token_req.code);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "invalid_grant",
|
||||
"error_description": "invalid authorization code"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
// handle empty client_id
|
||||
let client_id = if token_req.client_id.is_empty() {
|
||||
"mcp-client".to_string()
|
||||
} else {
|
||||
token_req.client_id.clone()
|
||||
};
|
||||
|
||||
// validate client
|
||||
match state
|
||||
.validate_client(&client_id, &token_req.redirect_uri)
|
||||
.await
|
||||
{
|
||||
Some(_) => {
|
||||
let session_id = token_req.code.replace("mcp-code-", "");
|
||||
info!("got session id: {}", session_id);
|
||||
|
||||
// create mcp access token
|
||||
match state.create_mcp_token(&session_id).await {
|
||||
Ok(token) => {
|
||||
info!("successfully created access token");
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({
|
||||
"access_token": token.access_token,
|
||||
"token_type": token.token_type,
|
||||
"expires_in": token.expires_in,
|
||||
"refresh_token": token.refresh_token,
|
||||
"scope": token.scope,
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
Err(e) => {
|
||||
error!("failed to create access token: {}", e);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({
|
||||
"error": "server_error",
|
||||
"error_description": format!("failed to create access token: {}", e)
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
None => {
|
||||
info!(
|
||||
"invalid client id or redirect uri: {} / {}",
|
||||
client_id, token_req.redirect_uri
|
||||
);
|
||||
(
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "invalid_client",
|
||||
"error_description": "invalid client id or redirect uri"
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Auth middleware for MCP connections
|
||||
async fn validate_token_middleware(
|
||||
State(token_store): State<Arc<McpOAuthStore>>,
|
||||
request: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
debug!("validate_token_middleware");
|
||||
// Extract the access token from the Authorization header
|
||||
let auth_header = request.headers().get("Authorization");
|
||||
let token = match auth_header {
|
||||
Some(header) => {
|
||||
let header_str = header.to_str().unwrap_or("");
|
||||
if let Some(stripped) = header_str.strip_prefix("Bearer ") {
|
||||
stripped.to_string()
|
||||
} else {
|
||||
return StatusCode::UNAUTHORIZED.into_response();
|
||||
}
|
||||
}
|
||||
None => {
|
||||
return StatusCode::UNAUTHORIZED.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
// Validate the token
|
||||
match token_store.validate_token(&token).await {
|
||||
Some(_) => next.run(request).await,
|
||||
None => StatusCode::UNAUTHORIZED.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
// handle oauth server metadata request
|
||||
async fn oauth_authorization_server() -> impl IntoResponse {
|
||||
let mut additional_fields = HashMap::new();
|
||||
additional_fields.insert(
|
||||
"response_types_supported".into(),
|
||||
Value::Array(vec![Value::String("code".into())]),
|
||||
);
|
||||
additional_fields.insert(
|
||||
"code_challenge_methods_supported".into(),
|
||||
Value::Array(vec![Value::String("S256".into())]),
|
||||
);
|
||||
let metadata = AuthorizationMetadata {
|
||||
authorization_endpoint: format!("http://{}/oauth/authorize", BIND_ADDRESS),
|
||||
token_endpoint: format!("http://{}/oauth/token", BIND_ADDRESS),
|
||||
scopes_supported: Some(vec!["profile".to_string(), "email".to_string()]),
|
||||
registration_endpoint: Some(format!("http://{}/oauth/register", BIND_ADDRESS)),
|
||||
issuer: Some(BIND_ADDRESS.to_string()),
|
||||
jwks_uri: Some(format!("http://{}/oauth/jwks", BIND_ADDRESS)),
|
||||
additional_fields,
|
||||
};
|
||||
debug!("metadata: {:?}", metadata);
|
||||
(StatusCode::OK, Json(metadata))
|
||||
}
|
||||
|
||||
// handle client registration request
|
||||
async fn oauth_register(
|
||||
State(state): State<Arc<McpOAuthStore>>,
|
||||
Json(req): Json<LocalClientRegistrationRequest>,
|
||||
) -> impl IntoResponse {
|
||||
debug!("register request: {:?}", req);
|
||||
if req.redirect_uris.is_empty() {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": "invalid_request",
|
||||
"error_description": "at least one redirect uri is required"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
// generate client id and secret
|
||||
let client_id = format!("client-{}", Uuid::new_v4());
|
||||
let client_secret = generate_random_string(32);
|
||||
|
||||
let client = OAuthClientConfig {
|
||||
client_id: client_id.clone(),
|
||||
client_secret: Some(client_secret.clone()),
|
||||
redirect_uri: req.redirect_uris[0].clone(),
|
||||
scopes: vec![],
|
||||
};
|
||||
|
||||
state
|
||||
.clients
|
||||
.write()
|
||||
.await
|
||||
.insert(client_id.clone(), client);
|
||||
|
||||
// return client information
|
||||
let response = ClientRegistrationResponse {
|
||||
client_id,
|
||||
client_secret: Some(client_secret),
|
||||
client_name: Some(req.client_name),
|
||||
redirect_uris: req.redirect_uris,
|
||||
additional_fields: HashMap::new(),
|
||||
};
|
||||
|
||||
(StatusCode::CREATED, Json(response)).into_response()
|
||||
}
|
||||
|
||||
// Log all HTTP requests
|
||||
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!(
|
||||
"{} {} {:?}{}\nContent-Type: {}",
|
||||
method, uri, version, header_log, 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
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
// Initialize logging
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "debug".to_string().into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
// Create the OAuth store
|
||||
let oauth_store = Arc::new(McpOAuthStore::new());
|
||||
|
||||
// Set up port
|
||||
let addr = BIND_ADDRESS.parse::<SocketAddr>()?;
|
||||
|
||||
// Create streamable HTTP service for MCP
|
||||
let mcp_service: StreamableHttpService<Counter, LocalSessionManager> =
|
||||
StreamableHttpService::new(
|
||||
|| Ok(Counter::new()),
|
||||
LocalSessionManager::default().into(),
|
||||
StreamableHttpServerConfig::default(),
|
||||
);
|
||||
|
||||
// Create protected MCP routes (require authorization)
|
||||
let protected_mcp_router =
|
||||
Router::new()
|
||||
.nest_service("/mcp", mcp_service)
|
||||
.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).options(oauth_authorization_server),
|
||||
)
|
||||
.route("/oauth/token", post(oauth_token).options(oauth_token))
|
||||
.route(
|
||||
"/oauth/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("/oauth/authorize", get(oauth_authorize))
|
||||
.route("/oauth/approve", post(oauth_approve))
|
||||
.merge(oauth_server_router) // Merge the CORS-enabled oauth server router
|
||||
.merge(protected_mcp_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?;
|
||||
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
tokio::signal::ctrl_c().await.ok();
|
||||
info!("Shutting down...");
|
||||
})
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -1,54 +0,0 @@
|
|||
use rmcp::transport::sse_server::{SseServer, SseServerConfig};
|
||||
use tracing_subscriber::{
|
||||
layer::SubscriberExt,
|
||||
util::SubscriberInitExt,
|
||||
{self},
|
||||
};
|
||||
mod common;
|
||||
use common::counter::Counter;
|
||||
|
||||
const BIND_ADDRESS: &str = "127.0.0.1:8000";
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "debug".to_string().into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
let config = SseServerConfig {
|
||||
bind: BIND_ADDRESS.parse()?,
|
||||
sse_path: "/sse".to_string(),
|
||||
post_path: "/message".to_string(),
|
||||
ct: tokio_util::sync::CancellationToken::new(),
|
||||
sse_keep_alive: None,
|
||||
};
|
||||
|
||||
let (sse_server, router) = SseServer::new(config);
|
||||
|
||||
// Do something with the router, e.g., add routes or middleware
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
|
||||
|
||||
let ct = sse_server.config.ct.child_token();
|
||||
|
||||
let server = axum::serve(listener, router).with_graceful_shutdown(async move {
|
||||
ct.cancelled().await;
|
||||
tracing::info!("sse server cancelled");
|
||||
});
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = server.await {
|
||||
tracing::error!(error = %e, "sse server shutdown with error");
|
||||
}
|
||||
});
|
||||
|
||||
let ct = sse_server.with_service(Counter::new);
|
||||
|
||||
tokio::signal::ctrl_c().await?;
|
||||
ct.cancel();
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -1,29 +0,0 @@
|
|||
use rmcp::transport::sse_server::SseServer;
|
||||
use tracing_subscriber::{
|
||||
layer::SubscriberExt,
|
||||
util::SubscriberInitExt,
|
||||
{self},
|
||||
};
|
||||
mod common;
|
||||
use common::counter::Counter;
|
||||
|
||||
const BIND_ADDRESS: &str = "127.0.0.1:8000";
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "debug".to_string().into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
let ct = SseServer::serve(BIND_ADDRESS.parse()?)
|
||||
.await?
|
||||
.with_service_directly(Counter::new);
|
||||
|
||||
tokio::signal::ctrl_c().await?;
|
||||
ct.cancel();
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -1,33 +0,0 @@
|
|||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>RMCP Authorized SSE Server</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; margin: 40px; line-height: 1.6; }
|
||||
h1 { color: #333; }
|
||||
.endpoint { background: #f4f4f4; padding: 10px; border-radius: 5px; }
|
||||
pre { background: #eee; padding: 10px; border-radius: 5px; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>RMCP Authorized SSE Server</h1>
|
||||
<p>This is a Server-Sent Events server example that requires OAuth authorization.</p>
|
||||
|
||||
<h2>Available Endpoints:</h2>
|
||||
<ul>
|
||||
<li class="endpoint"><code>/api/health</code> - Health check</li>
|
||||
<li class="endpoint"><code>/api/token/{token_id}</code> - Get test token (available: demo, test)</li>
|
||||
<li class="endpoint"><code>/sse</code> - SSE connection endpoint (requires authorization)</li>
|
||||
<li class="endpoint"><code>/message</code> - Message sending endpoint (requires authorization)</li>
|
||||
</ul>
|
||||
|
||||
<h2>Usage:</h2>
|
||||
<pre>
|
||||
# Get a token
|
||||
curl http://127.0.0.1:8000/api/token/demo
|
||||
|
||||
# Connect to SSE using the token
|
||||
curl -H "Authorization: Bearer demo-token" http://127.0.0.1:8000/sse
|
||||
</pre>
|
||||
</body>
|
||||
</html>
|
||||
|
|
@ -3,7 +3,6 @@ use std::env;
|
|||
use rmcp::{
|
||||
ServiceExt,
|
||||
transport::{
|
||||
sse_server::{SseServer, SseServerConfig},
|
||||
stdio,
|
||||
streamable_http_server::{StreamableHttpService, session::local::LocalSessionManager},
|
||||
},
|
||||
|
|
@ -12,7 +11,6 @@ use rmcp::{
|
|||
mod common;
|
||||
use common::progress_demo::ProgressDemo;
|
||||
|
||||
const SSE_BIND_ADDRESS: &str = "127.0.0.1:8000";
|
||||
const HTTP_BIND_ADDRESS: &str = "127.0.0.1:8001";
|
||||
|
||||
#[tokio::main]
|
||||
|
|
@ -24,14 +22,10 @@ async fn main() -> anyhow::Result<()> {
|
|||
|
||||
match transport_mode.as_str() {
|
||||
"stdio" => run_stdio().await,
|
||||
"sse" => run_sse().await,
|
||||
"http" | "streamhttp" => run_streamable_http().await,
|
||||
"all" => run_all_transports().await,
|
||||
_ => {
|
||||
eprintln!(
|
||||
"Usage: {} [stdio|sse|http|all]",
|
||||
env::args().next().unwrap()
|
||||
);
|
||||
eprintln!("Usage: {} [stdio|http|all]", env::args().next().unwrap());
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
|
|
@ -47,47 +41,6 @@ async fn run_stdio() -> anyhow::Result<()> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_sse() -> anyhow::Result<()> {
|
||||
println!("Running SSE server");
|
||||
let config = SseServerConfig {
|
||||
bind: SSE_BIND_ADDRESS.parse()?,
|
||||
sse_path: "/sse".to_string(),
|
||||
post_path: "/message".to_string(),
|
||||
ct: tokio_util::sync::CancellationToken::new(),
|
||||
sse_keep_alive: None,
|
||||
};
|
||||
|
||||
let (sse_server, router) = SseServer::new(config);
|
||||
|
||||
// Start the HTTP server for SSE
|
||||
let listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
|
||||
let ct = sse_server.config.ct.child_token();
|
||||
|
||||
let server = axum::serve(listener, router).with_graceful_shutdown(async move {
|
||||
ct.cancelled().await;
|
||||
tracing::info!("SSE server cancelled");
|
||||
});
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = server.await {
|
||||
tracing::error!(error = %e, "SSE server shutdown with error");
|
||||
}
|
||||
});
|
||||
|
||||
// Start the MCP service with SSE transport
|
||||
let ct = sse_server.with_service(ProgressDemo::new);
|
||||
|
||||
tracing::info!(
|
||||
"Progress Demo SSE server started at http://{}/sse",
|
||||
SSE_BIND_ADDRESS
|
||||
);
|
||||
tracing::info!("Press Ctrl+C to shutdown");
|
||||
|
||||
tokio::signal::ctrl_c().await?;
|
||||
ct.cancel();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_streamable_http() -> anyhow::Result<()> {
|
||||
println!("Running Streamable HTTP server");
|
||||
let service = StreamableHttpService::new(
|
||||
|
|
@ -114,18 +67,6 @@ async fn run_streamable_http() -> anyhow::Result<()> {
|
|||
|
||||
async fn run_all_transports() -> anyhow::Result<()> {
|
||||
println!("Running all transports");
|
||||
// Start SSE server
|
||||
let sse_config = SseServerConfig {
|
||||
bind: SSE_BIND_ADDRESS.parse()?,
|
||||
sse_path: "/sse".to_string(),
|
||||
post_path: "/message".to_string(),
|
||||
ct: tokio_util::sync::CancellationToken::new(),
|
||||
sse_keep_alive: None,
|
||||
};
|
||||
|
||||
let (sse_server, sse_router) = SseServer::new(sse_config);
|
||||
let sse_listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
|
||||
let sse_ct = sse_server.config.ct.child_token();
|
||||
|
||||
// Start Streamable HTTP server
|
||||
let http_service = StreamableHttpService::new(
|
||||
|
|
@ -136,19 +77,6 @@ async fn run_all_transports() -> anyhow::Result<()> {
|
|||
let http_router = axum::Router::new().nest_service("/mcp", http_service);
|
||||
let http_listener = tokio::net::TcpListener::bind(HTTP_BIND_ADDRESS).await?;
|
||||
|
||||
// Start SSE HTTP server
|
||||
let sse_http_server =
|
||||
axum::serve(sse_listener, sse_router).with_graceful_shutdown(async move {
|
||||
sse_ct.cancelled().await;
|
||||
tracing::info!("SSE server cancelled");
|
||||
});
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = sse_http_server.await {
|
||||
tracing::error!(error = %e, "SSE server shutdown with error");
|
||||
}
|
||||
});
|
||||
|
||||
// Start Streamable HTTP server
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(http_listener, http_router)
|
||||
|
|
@ -156,11 +84,13 @@ async fn run_all_transports() -> anyhow::Result<()> {
|
|||
.await;
|
||||
});
|
||||
|
||||
// Start MCP service with SSE
|
||||
let mcp_sse_ct = sse_server.with_service(ProgressDemo::new);
|
||||
tracing::info!(
|
||||
"Progress Demo HTTP server started at http://{}/mcp",
|
||||
HTTP_BIND_ADDRESS
|
||||
);
|
||||
tracing::info!("Press Ctrl+C to shutdown");
|
||||
|
||||
tokio::signal::ctrl_c().await?;
|
||||
mcp_sse_ct.cancel();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
173
examples/servers/src/simple_auth_streamhttp.rs
Normal file
173
examples/servers/src/simple_auth_streamhttp.rs
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
/// This example shows how to use the RMCP streamable HTTP server with simple token authorization.
|
||||
/// Use the inspector to view this server https://github.com/modelcontextprotocol/inspector
|
||||
/// The default index page is available at http://127.0.0.1:8000/
|
||||
/// # Get a token
|
||||
/// curl http://127.0.0.1:8000/api/token/demo
|
||||
/// # Connect using the token
|
||||
/// curl -X POST -H "Authorization: Bearer demo-token" -H "Content-Type: application/json" \
|
||||
/// -d '{"jsonrpc":"2.0","method":"initialize","params":{},"id":1}' \
|
||||
/// http://127.0.0.1:8000/mcp
|
||||
use std::{net::SocketAddr, sync::Arc};
|
||||
|
||||
use anyhow::Result;
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, State},
|
||||
http::{HeaderMap, Request, StatusCode},
|
||||
middleware::{self, Next},
|
||||
response::{Html, Response},
|
||||
routing::get,
|
||||
};
|
||||
use rmcp::transport::{
|
||||
StreamableHttpServerConfig,
|
||||
streamable_http_server::{session::local::LocalSessionManager, tower::StreamableHttpService},
|
||||
};
|
||||
mod common;
|
||||
use common::counter::Counter;
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||
|
||||
const BIND_ADDRESS: &str = "127.0.0.1:8000";
|
||||
const INDEX_HTML: &str = r#"<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>MCP Streamable HTTP Auth Server</title>
|
||||
</head>
|
||||
<body>
|
||||
<h1>MCP Streamable HTTP Server with Auth</h1>
|
||||
<p>Get a token: <code>curl http://127.0.0.1:8000/api/token/demo</code></p>
|
||||
<p>Connect: <code>curl -X POST -H "Authorization: Bearer demo-token" -H "Content-Type: application/json" -d '{"jsonrpc":"2.0","method":"initialize","params":{},"id":1}' http://127.0.0.1:8000/mcp</code></p>
|
||||
</body>
|
||||
</html>"#;
|
||||
|
||||
// A simple token store
|
||||
struct TokenStore {
|
||||
valid_tokens: Vec<String>,
|
||||
}
|
||||
|
||||
impl TokenStore {
|
||||
fn new() -> Self {
|
||||
// For demonstration purposes, use more secure token management in production
|
||||
Self {
|
||||
valid_tokens: vec!["demo-token".to_string(), "test-token".to_string()],
|
||||
}
|
||||
}
|
||||
|
||||
fn is_valid(&self, token: &str) -> bool {
|
||||
self.valid_tokens.contains(&token.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
// Extract authorization token
|
||||
fn extract_token(headers: &HeaderMap) -> Option<String> {
|
||||
headers
|
||||
.get("Authorization")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|auth_header| {
|
||||
auth_header
|
||||
.strip_prefix("Bearer ")
|
||||
.map(|stripped| stripped.to_string())
|
||||
})
|
||||
}
|
||||
|
||||
// Authorization middleware
|
||||
async fn auth_middleware(
|
||||
State(token_store): State<Arc<TokenStore>>,
|
||||
headers: HeaderMap,
|
||||
request: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> Result<Response, StatusCode> {
|
||||
match extract_token(&headers) {
|
||||
Some(token) if token_store.is_valid(&token) => {
|
||||
// Token is valid, proceed with the request
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
_ => {
|
||||
// Token is invalid, return 401 error
|
||||
Err(StatusCode::UNAUTHORIZED)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Root path handler
|
||||
async fn index() -> Html<&'static str> {
|
||||
Html(INDEX_HTML)
|
||||
}
|
||||
|
||||
// Health check endpoint
|
||||
async fn health_check() -> &'static str {
|
||||
"OK"
|
||||
}
|
||||
|
||||
// Token generation endpoint (simplified example)
|
||||
async fn get_token(Path(token_id): Path<String>) -> Result<Json<serde_json::Value>, StatusCode> {
|
||||
// In a real application, you should authenticate the user and generate a real token
|
||||
if token_id == "demo" || token_id == "test" {
|
||||
let token = format!("{}-token", token_id);
|
||||
Ok(Json(serde_json::json!({
|
||||
"access_token": token,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600
|
||||
})))
|
||||
} else {
|
||||
Err(StatusCode::UNAUTHORIZED)
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
// Initialize logging
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "debug".to_string().into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
// Create token store
|
||||
let token_store = Arc::new(TokenStore::new());
|
||||
|
||||
// Set up port
|
||||
let addr = BIND_ADDRESS.parse::<SocketAddr>()?;
|
||||
|
||||
// Create streamable HTTP service
|
||||
let mcp_service: StreamableHttpService<Counter, LocalSessionManager> =
|
||||
StreamableHttpService::new(
|
||||
|| Ok(Counter::new()),
|
||||
LocalSessionManager::default().into(),
|
||||
StreamableHttpServerConfig::default(),
|
||||
);
|
||||
|
||||
// Create API routes
|
||||
let api_routes = Router::new()
|
||||
.route("/health", get(health_check))
|
||||
.route("/token/{token_id}", get(get_token));
|
||||
|
||||
// Create protected MCP routes (require authorization)
|
||||
let protected_mcp_router =
|
||||
Router::new()
|
||||
.nest_service("/mcp", mcp_service)
|
||||
.layer(middleware::from_fn_with_state(
|
||||
token_store.clone(),
|
||||
auth_middleware,
|
||||
));
|
||||
|
||||
// Create main router, public endpoints don't require authorization
|
||||
let app = Router::new()
|
||||
.route("/", get(index))
|
||||
.nest("/api", api_routes)
|
||||
.merge(protected_mcp_router);
|
||||
|
||||
// Start HTTP server
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
tracing::info!("Server started on {}", addr);
|
||||
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
tokio::signal::ctrl_c().await.ok();
|
||||
println!("Shutting down...");
|
||||
})
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Loading…
Reference in a new issue