rmcp-server-kit 1.3.2

Reusable MCP server framework with auth, RBAC, and Streamable HTTP transport (built on the rmcp SDK)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294

//! Admin diagnostic endpoints.
//!
//! When enabled, the server exposes a small `/admin/*` surface that returns
//! read-only diagnostic JSON: uptime, active auth configuration (no
//! secrets), auth counters, and an RBAC policy summary.
//!
//! The admin router is always wrapped in the existing auth + RBAC stack
//! and additionally requires the caller's role to match the `role` field
//! on [`crate::admin::AdminConfig`]. Configuration validation refuses to
//! enable admin without auth.

use std::{
    sync::Arc,
    time::{Instant, SystemTime, UNIX_EPOCH},
};

use arc_swap::ArcSwap;
use axum::{
    Json, Router,
    body::Body,
    extract::{Request, State},
    http::StatusCode,
    middleware::Next,
    response::{IntoResponse, Response},
    routing::get,
};
use serde::Serialize;

use crate::{auth::AuthState, rbac::RbacPolicy};

/// Admin endpoint configuration.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct AdminConfig {
    /// RBAC role required to access the admin endpoints.
    pub role: String,
}

impl Default for AdminConfig {
    fn default() -> Self {
        Self {
            role: "admin".to_owned(),
        }
    }
}

/// Shared state used by admin endpoint handlers.
#[allow(
    missing_debug_implementations,
    reason = "contains Arc<AuthState> and ArcSwap<RbacPolicy> without Debug impls"
)]
#[derive(Clone)]
#[non_exhaustive]
pub(crate) struct AdminState {
    /// Server start instant, used for uptime.
    pub started_at: Instant,
    /// Server name for /admin/status.
    pub name: String,
    /// Server version for /admin/status.
    pub version: String,
    /// Shared auth state (optional for test constructions).
    pub auth: Option<Arc<AuthState>>,
    /// Shared RBAC policy for diagnostics.
    pub rbac: Arc<ArcSwap<RbacPolicy>>,
}

/// `/admin/status` response body.
#[derive(Debug, Clone, Serialize)]
#[non_exhaustive]
pub struct AdminStatus {
    /// Server name.
    pub name: String,
    /// Server version string.
    pub version: String,
    /// Seconds since the server process started.
    pub uptime_seconds: u64,
    /// Wall-clock UNIX epoch at startup.
    pub started_at_epoch: u64,
}

fn admin_status(state: &AdminState) -> AdminStatus {
    let started_epoch = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or_default()
        .saturating_sub(state.started_at.elapsed().as_secs());
    AdminStatus {
        name: state.name.clone(),
        version: state.version.clone(),
        uptime_seconds: state.started_at.elapsed().as_secs(),
        started_at_epoch: started_epoch,
    }
}

async fn status_handler(State(state): State<AdminState>) -> Json<AdminStatus> {
    Json(admin_status(&state))
}

async fn auth_keys_handler(State(state): State<AdminState>) -> Response {
    state.auth.as_ref().map_or_else(
        || not_available("auth is not configured"),
        |auth| Json(auth.api_key_summaries()).into_response(),
    )
}

async fn auth_counters_handler(State(state): State<AdminState>) -> Response {
    state.auth.as_ref().map_or_else(
        || not_available("auth is not configured"),
        |auth| Json(auth.counters_snapshot()).into_response(),
    )
}

async fn rbac_handler(State(state): State<AdminState>) -> Response {
    Json(state.rbac.load().summary()).into_response()
}

fn not_available(reason: &str) -> Response {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(serde_json::json!({
            "error": "unavailable",
            "error_description": reason,
        })),
    )
        .into_response()
}

/// Role-check middleware for admin routes.
///
/// Reads the caller's role from the `AuthIdentity` request extension
/// (populated by the outer auth middleware) and rejects requests whose
/// role does not match `expected_role`.
pub async fn require_admin_role(
    expected_role: Arc<str>,
    req: Request<Body>,
    next: Next,
) -> Response {
    let role = req
        .extensions()
        .get::<crate::auth::AuthIdentity>()
        .map_or("", |id| id.role.as_str());
    if role != expected_role.as_ref() {
        return (
            StatusCode::FORBIDDEN,
            Json(serde_json::json!({
                "error": "forbidden",
                "error_description": "admin role required",
            })),
        )
            .into_response();
    }
    next.run(req).await
}

/// Build the `/admin` router layered with the admin role check.
///
/// The caller is expected to merge this router on top of their top-level
/// router *after* the auth + RBAC middleware has been installed, so that
/// by the time a request reaches this router the task-local role is set.
pub(crate) fn admin_router(state: AdminState, config: &AdminConfig) -> Router {
    let role: Arc<str> = Arc::from(config.role.as_str());
    Router::new()
        .route("/admin/status", get(status_handler))
        .route("/admin/auth/keys", get(auth_keys_handler))
        .route("/admin/auth/counters", get(auth_counters_handler))
        .route("/admin/rbac", get(rbac_handler))
        .with_state(state)
        .layer(axum::middleware::from_fn(move |req, next| {
            let r = Arc::clone(&role);
            require_admin_role(r, req, next)
        }))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::expect_used)]
    use std::sync::Mutex;

    use axum::http::Request;
    use tower::ServiceExt as _;

    use super::*;
    use crate::{
        auth::{ApiKeyEntry, AuthCounters, AuthIdentity, AuthMethod, AuthState},
        rbac::{RbacConfig, RbacPolicy, RoleConfig},
    };

    fn make_auth_state() -> Arc<AuthState> {
        Arc::new(AuthState {
            api_keys: ArcSwap::from_pointee(vec![ApiKeyEntry::new(
                "test-key",
                "argon2id-hash",
                "admin",
            )]),
            rate_limiter: None,
            pre_auth_limiter: None,
            #[cfg(feature = "oauth")]
            jwks_cache: None,
            seen_identities: Mutex::new(std::collections::HashSet::default()),
            counters: AuthCounters::default(),
        })
    }

    fn make_state() -> AdminState {
        AdminState {
            started_at: Instant::now(),
            name: "test".into(),
            version: "0.0.0".into(),
            auth: Some(make_auth_state()),
            rbac: Arc::new(ArcSwap::from_pointee(RbacPolicy::new(
                &RbacConfig::with_roles(vec![RoleConfig::new(
                    "admin",
                    vec!["*".into()],
                    vec!["*".into()],
                )]),
            ))),
        }
    }

    fn admin_req(uri: &str, role: Option<&str>) -> Request<Body> {
        let mut req = Request::builder().uri(uri).body(Body::empty()).unwrap();
        if let Some(r) = role {
            req.extensions_mut().insert(AuthIdentity {
                name: "tester".into(),
                role: r.to_owned(),
                method: AuthMethod::BearerToken,
                raw_token: None,
                sub: None,
            });
        }
        req
    }

    #[tokio::test]
    async fn keys_endpoint_omits_hash() {
        let app = admin_router(make_state(), &AdminConfig::default());
        let resp = app
            .oneshot(admin_req("/admin/auth/keys", Some("admin")))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = axum::body::to_bytes(resp.into_body(), 64 * 1024)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        let arr = json.as_array().unwrap();
        assert_eq!(arr.len(), 1);
        assert_eq!(arr[0]["name"], "test-key");
        assert!(arr[0].get("hash").is_none());
    }

    #[tokio::test]
    async fn wrong_role_gets_403() {
        let app = admin_router(make_state(), &AdminConfig::default());
        let resp = app
            .oneshot(admin_req("/admin/status", Some("viewer")))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn no_identity_gets_403() {
        let app = admin_router(make_state(), &AdminConfig::default());
        let resp = app.oneshot(admin_req("/admin/status", None)).await.unwrap();
        assert_eq!(resp.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn status_returns_uptime() {
        let app = admin_router(make_state(), &AdminConfig::default());
        let resp = app
            .oneshot(admin_req("/admin/status", Some("admin")))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn rbac_summary_includes_role_list() {
        let app = admin_router(make_state(), &AdminConfig::default());
        let resp = app
            .oneshot(admin_req("/admin/rbac", Some("admin")))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = axum::body::to_bytes(resp.into_body(), 64 * 1024)
            .await
            .unwrap();
        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(json["enabled"], true);
        assert_eq!(json["roles"][0]["name"], "admin");
    }
}