1use std::{net::IpAddr, num::NonZeroU32, sync::Arc, time::Duration};
12
13use axum::{
14 body::Body,
15 http::{Method, Request, StatusCode},
16 middleware::Next,
17 response::{IntoResponse, Response},
18};
19use hmac::{Hmac, KeyInit, Mac};
20use http_body_util::BodyExt;
21use secrecy::{ExposeSecret, SecretString};
22use serde::Deserialize;
23use sha2::Sha256;
24
25use crate::{auth::AuthIdentity, bounded_limiter::BoundedKeyedLimiter, error::McpxError};
26
27pub(crate) type ToolRateLimiter = BoundedKeyedLimiter<IpAddr>;
30
31const DEFAULT_TOOL_RATE: NonZeroU32 = NonZeroU32::new(120).unwrap();
34
35const DEFAULT_TOOL_MAX_TRACKED_KEYS: usize = 10_000;
38
39const DEFAULT_TOOL_IDLE_EVICTION: Duration = Duration::from_mins(15);
41
42#[must_use]
48pub(crate) fn build_tool_rate_limiter(
49 max_per_minute: u32,
50 burst: Option<u32>,
51) -> Arc<ToolRateLimiter> {
52 build_tool_rate_limiter_with_bounds(
53 max_per_minute,
54 burst,
55 DEFAULT_TOOL_MAX_TRACKED_KEYS,
56 DEFAULT_TOOL_IDLE_EVICTION,
57 )
58}
59
60#[must_use]
66pub(crate) fn build_tool_rate_limiter_with_bounds(
67 max_per_minute: u32,
68 burst: Option<u32>,
69 max_tracked_keys: usize,
70 idle_eviction: Duration,
71) -> Arc<ToolRateLimiter> {
72 let mut quota =
73 governor::Quota::per_minute(NonZeroU32::new(max_per_minute).unwrap_or(DEFAULT_TOOL_RATE));
74 if let Some(b) = burst.and_then(NonZeroU32::new) {
75 quota = quota.allow_burst(b);
76 }
77 Arc::new(BoundedKeyedLimiter::new(
78 quota,
79 max_tracked_keys,
80 idle_eviction,
81 ))
82}
83
84tokio::task_local! {
91 static CURRENT_ROLE: String;
92 static CURRENT_IDENTITY: String;
93 static CURRENT_TOKEN: SecretString;
94 static CURRENT_SUB: String;
95}
96
97#[must_use]
100pub fn current_role() -> Option<String> {
101 CURRENT_ROLE.try_with(Clone::clone).ok()
102}
103
104#[must_use]
107pub fn current_identity() -> Option<String> {
108 CURRENT_IDENTITY.try_with(Clone::clone).ok()
109}
110
111#[must_use]
125pub fn current_token() -> Option<SecretString> {
126 CURRENT_TOKEN
127 .try_with(|t| {
128 if t.expose_secret().is_empty() {
129 None
130 } else {
131 Some(t.clone())
132 }
133 })
134 .ok()
135 .flatten()
136}
137
138#[must_use]
142pub fn current_sub() -> Option<String> {
143 CURRENT_SUB
144 .try_with(Clone::clone)
145 .ok()
146 .filter(|s| !s.is_empty())
147}
148
149pub async fn with_token_scope<F: Future>(token: SecretString, f: F) -> F::Output {
156 CURRENT_TOKEN.scope(token, f).await
157}
158
159pub async fn with_rbac_scope<F: Future>(
166 role: String,
167 identity: String,
168 token: SecretString,
169 sub: String,
170 f: F,
171) -> F::Output {
172 CURRENT_ROLE
173 .scope(
174 role,
175 CURRENT_IDENTITY.scope(
176 identity,
177 CURRENT_TOKEN.scope(token, CURRENT_SUB.scope(sub, f)),
178 ),
179 )
180 .await
181}
182
183#[derive(Debug, Clone, Deserialize)]
185#[non_exhaustive]
186pub struct RoleConfig {
187 pub name: String,
189 #[serde(default)]
191 pub description: Option<String>,
192 #[serde(default)]
194 pub allow: Vec<String>,
195 #[serde(default)]
197 pub deny: Vec<String>,
198 #[serde(default = "default_hosts")]
200 pub hosts: Vec<String>,
201 #[serde(default)]
205 pub argument_allowlists: Vec<ArgumentAllowlist>,
206}
207
208impl RoleConfig {
209 #[must_use]
211 pub fn new(name: impl Into<String>, allow: Vec<String>, hosts: Vec<String>) -> Self {
212 Self {
213 name: name.into(),
214 description: None,
215 allow,
216 deny: vec![],
217 hosts,
218 argument_allowlists: vec![],
219 }
220 }
221
222 #[must_use]
224 pub fn with_argument_allowlists(mut self, allowlists: Vec<ArgumentAllowlist>) -> Self {
225 self.argument_allowlists = allowlists;
226 self
227 }
228}
229
230#[derive(Debug, Clone, Deserialize)]
253#[non_exhaustive]
254pub struct ArgumentAllowlist {
255 pub tool: String,
257 pub argument: String,
259 #[serde(default)]
261 pub allowed: Vec<String>,
262}
263
264impl ArgumentAllowlist {
265 #[must_use]
267 pub fn new(tool: impl Into<String>, argument: impl Into<String>, allowed: Vec<String>) -> Self {
268 Self {
269 tool: tool.into(),
270 argument: argument.into(),
271 allowed,
272 }
273 }
274}
275
276fn default_hosts() -> Vec<String> {
277 vec!["*".into()]
278}
279
280#[derive(Debug, Clone, Default, Deserialize)]
282#[non_exhaustive]
283pub struct RbacConfig {
284 #[serde(default)]
286 pub enabled: bool,
287 #[serde(default)]
289 pub roles: Vec<RoleConfig>,
290 #[serde(default)]
299 pub redaction_salt: Option<SecretString>,
300}
301
302impl RbacConfig {
303 #[must_use]
305 pub fn with_roles(roles: Vec<RoleConfig>) -> Self {
306 Self {
307 enabled: true,
308 roles,
309 redaction_salt: None,
310 }
311 }
312}
313
314#[derive(Debug, Clone, Copy, PartialEq, Eq)]
316#[non_exhaustive]
317pub enum RbacDecision {
318 Allow,
320 Deny,
322}
323
324#[derive(Debug, Clone, serde::Serialize)]
326#[non_exhaustive]
327pub struct RbacRoleSummary {
328 pub name: String,
330 pub allow: usize,
332 pub deny: usize,
334 pub hosts: usize,
336 pub argument_allowlists: usize,
338}
339
340#[derive(Debug, Clone, serde::Serialize)]
342#[non_exhaustive]
343pub struct RbacPolicySummary {
344 pub enabled: bool,
346 pub roles: Vec<RbacRoleSummary>,
348}
349
350#[derive(Debug, Clone)]
356#[non_exhaustive]
357pub struct RbacPolicy {
358 roles: Vec<RoleConfig>,
359 enabled: bool,
360 redaction_salt: Arc<SecretString>,
363}
364
365impl RbacPolicy {
366 #[must_use]
369 pub fn new(config: &RbacConfig) -> Self {
370 let salt = config
371 .redaction_salt
372 .clone()
373 .unwrap_or_else(|| process_redaction_salt().clone());
374 Self {
375 roles: config.roles.clone(),
376 enabled: config.enabled,
377 redaction_salt: Arc::new(salt),
378 }
379 }
380
381 #[must_use]
383 pub fn disabled() -> Self {
384 Self {
385 roles: Vec::new(),
386 enabled: false,
387 redaction_salt: Arc::new(process_redaction_salt().clone()),
388 }
389 }
390
391 #[must_use]
393 pub fn is_enabled(&self) -> bool {
394 self.enabled
395 }
396
397 #[must_use]
402 pub fn summary(&self) -> RbacPolicySummary {
403 let roles = self
404 .roles
405 .iter()
406 .map(|r| RbacRoleSummary {
407 name: r.name.clone(),
408 allow: r.allow.len(),
409 deny: r.deny.len(),
410 hosts: r.hosts.len(),
411 argument_allowlists: r.argument_allowlists.len(),
412 })
413 .collect();
414 RbacPolicySummary {
415 enabled: self.enabled,
416 roles,
417 }
418 }
419
420 #[must_use]
425 pub fn check_operation(&self, role: &str, operation: &str) -> RbacDecision {
426 if !self.enabled {
427 return RbacDecision::Allow;
428 }
429 let Some(role_cfg) = self.find_role(role) else {
430 return RbacDecision::Deny;
431 };
432 if role_cfg.deny.iter().any(|d| d == operation) {
433 return RbacDecision::Deny;
434 }
435 if role_cfg.allow.iter().any(|a| a == "*" || a == operation) {
436 return RbacDecision::Allow;
437 }
438 RbacDecision::Deny
439 }
440
441 #[must_use]
448 pub fn check(&self, role: &str, operation: &str, host: &str) -> RbacDecision {
449 if !self.enabled {
450 return RbacDecision::Allow;
451 }
452 let Some(role_cfg) = self.find_role(role) else {
453 return RbacDecision::Deny;
454 };
455 if role_cfg.deny.iter().any(|d| d == operation) {
456 return RbacDecision::Deny;
457 }
458 if !role_cfg.allow.iter().any(|a| a == "*" || a == operation) {
459 return RbacDecision::Deny;
460 }
461 if !Self::host_matches(&role_cfg.hosts, host) {
462 return RbacDecision::Deny;
463 }
464 RbacDecision::Allow
465 }
466
467 #[must_use]
469 pub fn host_visible(&self, role: &str, host: &str) -> bool {
470 if !self.enabled {
471 return true;
472 }
473 let Some(role_cfg) = self.find_role(role) else {
474 return false;
475 };
476 Self::host_matches(&role_cfg.hosts, host)
477 }
478
479 #[must_use]
481 pub fn host_patterns(&self, role: &str) -> Option<&[String]> {
482 self.find_role(role).map(|r| r.hosts.as_slice())
483 }
484
485 #[must_use]
514 pub fn argument_allowed(&self, role: &str, tool: &str, argument: &str, value: &str) -> bool {
515 if !self.enabled {
516 return true;
517 }
518 let Some(role_cfg) = self.find_role(role) else {
519 return false;
520 };
521 for al in &role_cfg.argument_allowlists {
522 if al.tool != tool && !glob_match(&al.tool, tool) {
523 continue;
524 }
525 if al.argument != argument {
526 continue;
527 }
528 if al.allowed.is_empty() {
529 continue;
530 }
531 let Some(tokens) = shlex::split(value) else {
536 return false;
537 };
538 let Some(first_token) = tokens.first() else {
539 return false;
540 };
541 if first_token.is_empty() {
545 return false;
546 }
547 let basename = first_token
551 .rsplit('/')
552 .next()
553 .unwrap_or(first_token.as_str());
554 if !al.allowed.iter().any(|a| a == first_token || a == basename) {
555 return false;
556 }
557 }
558 true
559 }
560
561 #[must_use]
571 pub fn has_argument_allowlist(&self, role: &str, tool: &str, argument: &str) -> bool {
572 if !self.enabled {
573 return false;
574 }
575 let Some(role_cfg) = self.find_role(role) else {
576 return false;
577 };
578 role_cfg.argument_allowlists.iter().any(|al| {
579 (al.tool == tool || glob_match(&al.tool, tool))
580 && al.argument == argument
581 && !al.allowed.is_empty()
582 })
583 }
584
585 fn find_role(&self, name: &str) -> Option<&RoleConfig> {
587 self.roles.iter().find(|r| r.name == name)
588 }
589
590 fn host_matches(patterns: &[String], host: &str) -> bool {
592 patterns.iter().any(|p| glob_match(p, host))
593 }
594
595 #[must_use]
604 pub fn redact_arg(&self, value: &str) -> String {
605 redact_with_salt(self.redaction_salt.expose_secret().as_bytes(), value)
606 }
607}
608
609fn process_redaction_salt() -> &'static SecretString {
612 use base64::{Engine as _, engine::general_purpose::STANDARD_NO_PAD};
613 static PROCESS_SALT: std::sync::OnceLock<SecretString> = std::sync::OnceLock::new();
614 PROCESS_SALT.get_or_init(|| {
615 let mut bytes = [0u8; 32];
616 rand::fill(&mut bytes);
617 SecretString::from(STANDARD_NO_PAD.encode(bytes))
620 })
621}
622
623fn redact_with_salt(salt: &[u8], value: &str) -> String {
628 use std::fmt::Write as _;
629
630 use sha2::Digest as _;
631
632 type HmacSha256 = Hmac<Sha256>;
633 let mut mac = if let Ok(m) = HmacSha256::new_from_slice(salt) {
639 m
640 } else {
641 let digest = Sha256::digest(salt);
642 #[allow(
643 clippy::expect_used,
644 reason = "32-byte SHA-256 digest is unconditionally valid as an HMAC-SHA256 key (RFC 2104 allows any key length); see surrounding comment"
645 )]
646 HmacSha256::new_from_slice(&digest).expect("32-byte SHA256 digest is valid HMAC key")
647 };
648 mac.update(value.as_bytes());
649 let bytes = mac.finalize().into_bytes();
650 let prefix = bytes.get(..4).unwrap_or(&[0; 4]);
652 let mut out = String::with_capacity(8);
653 for b in prefix {
654 let _ = write!(out, "{b:02x}");
655 }
656 out
657}
658
659#[allow(
680 clippy::too_many_lines,
681 reason = "linear request lifecycle (body collect → JSON-RPC parse → policy dispatch) kept inline for security review visibility; helpers already extracted"
682)]
683pub(crate) async fn rbac_middleware(
684 policy: Arc<RbacPolicy>,
685 tool_limiter: Option<Arc<ToolRateLimiter>>,
686 req: Request<Body>,
687 next: Next,
688) -> Response {
689 if req.method() != Method::POST {
691 return next.run(req).await;
692 }
693
694 let peer_ip: Option<IpAddr> = crate::transport::limiter_client_ip(req.extensions());
697
698 let identity = req.extensions().get::<AuthIdentity>();
700 let identity_name = identity.map(|id| id.name.clone()).unwrap_or_default();
701 let role = identity.map(|id| id.role.clone()).unwrap_or_default();
702 let raw_token: SecretString = identity
705 .and_then(|id| id.raw_token.clone())
706 .unwrap_or_else(|| SecretString::from(String::new()));
707 let sub = identity.and_then(|id| id.sub.clone()).unwrap_or_default();
708
709 if policy.is_enabled() && identity.is_none() {
711 return McpxError::Rbac("no authenticated identity".into()).into_response();
712 }
713
714 let (parts, body) = req.into_parts();
716 let bytes = match body.collect().await {
717 Ok(collected) => collected.to_bytes(),
718 Err(e) => {
719 tracing::error!(error = %e, "failed to read request body");
720 return (
721 StatusCode::INTERNAL_SERVER_ERROR,
722 "failed to read request body",
723 )
724 .into_response();
725 }
726 };
727
728 if let Ok(json) = serde_json::from_slice::<serde_json::Value>(&bytes) {
730 let tool_calls = extract_tool_calls(&json);
731 if !tool_calls.is_empty() {
732 for params in tool_calls {
733 if let Some(resp) = enforce_rate_limit(tool_limiter.as_deref(), peer_ip) {
734 #[cfg(feature = "metrics")]
735 crate::metrics::record_rate_limit_deny(&parts.extensions, "tool");
736 return resp;
737 }
738 if policy.is_enabled()
739 && let Some(resp) = enforce_tool_policy(&policy, &identity_name, &role, params)
740 {
741 return resp;
742 }
743 }
744 }
745 }
746 let req = Request::from_parts(parts, Body::from(bytes));
750
751 if role.is_empty() {
753 next.run(req).await
754 } else {
755 CURRENT_ROLE
756 .scope(
757 role,
758 CURRENT_IDENTITY.scope(
759 identity_name,
760 CURRENT_TOKEN.scope(raw_token, CURRENT_SUB.scope(sub, next.run(req))),
761 ),
762 )
763 .await
764 }
765}
766
767fn extract_tool_calls(value: &serde_json::Value) -> Vec<&serde_json::Value> {
773 match value {
774 serde_json::Value::Object(map) => map
775 .get("method")
776 .and_then(serde_json::Value::as_str)
777 .filter(|method| *method == "tools/call")
778 .and_then(|_| map.get("params"))
779 .into_iter()
780 .collect(),
781 serde_json::Value::Array(items) => items
782 .iter()
783 .filter_map(|item| match item {
784 serde_json::Value::Object(map) => map
785 .get("method")
786 .and_then(serde_json::Value::as_str)
787 .filter(|method| *method == "tools/call")
788 .and_then(|_| map.get("params")),
789 serde_json::Value::Null
790 | serde_json::Value::Bool(_)
791 | serde_json::Value::Number(_)
792 | serde_json::Value::String(_)
793 | serde_json::Value::Array(_) => None,
794 })
795 .collect(),
796 serde_json::Value::Null
797 | serde_json::Value::Bool(_)
798 | serde_json::Value::Number(_)
799 | serde_json::Value::String(_) => Vec::new(),
800 }
801}
802
803fn enforce_rate_limit(
806 tool_limiter: Option<&ToolRateLimiter>,
807 peer_ip: Option<IpAddr>,
808) -> Option<Response> {
809 let limiter = tool_limiter?;
810 let ip = peer_ip?;
811 if let Err(wait) = limiter.check_key_wait(&ip) {
812 tracing::warn!(%ip, "tool invocation rate limited");
813 return Some(
814 McpxError::RateLimitedFor {
815 message: "too many tool invocations".into(),
816 retry_after: wait,
817 }
818 .into_response(),
819 );
820 }
821 None
822}
823
824fn enforce_tool_policy(
833 policy: &RbacPolicy,
834 identity_name: &str,
835 role: &str,
836 params: &serde_json::Value,
837) -> Option<Response> {
838 let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
839 let host = params
840 .get("arguments")
841 .and_then(|a| a.get("host"))
842 .and_then(|h| h.as_str());
843
844 let decision = if let Some(host) = host {
845 policy.check(role, tool_name, host)
846 } else {
847 policy.check_operation(role, tool_name)
848 };
849 if decision == RbacDecision::Deny {
850 tracing::warn!(
851 user = %identity_name,
852 role = %role,
853 tool = tool_name,
854 host = host.unwrap_or("-"),
855 "RBAC denied"
856 );
857 return Some(
858 McpxError::Rbac(format!("{tool_name} denied for role '{role}'")).into_response(),
859 );
860 }
861
862 let args = params.get("arguments").and_then(|a| a.as_object())?;
863 for (arg_key, arg_val) in args {
864 if let Some(resp) = check_argument(policy, identity_name, role, tool_name, arg_key, arg_val)
865 {
866 return Some(resp);
867 }
868 }
869 None
870}
871
872fn check_argument(
873 policy: &RbacPolicy,
874 identity_name: &str,
875 role: &str,
876 tool_name: &str,
877 arg_key: &str,
878 arg_val: &serde_json::Value,
879) -> Option<Response> {
880 if !policy.has_argument_allowlist(role, tool_name, arg_key) {
881 return None;
882 }
883 let Some(val_str) = arg_val.as_str() else {
884 tracing::warn!(
890 user = %identity_name,
891 role = %role,
892 tool = tool_name,
893 argument = arg_key,
894 value_type = json_value_type(arg_val),
895 "non-string argument rejected by allowlist"
896 );
897 return Some(
898 McpxError::Rbac(format!(
899 "argument '{arg_key}' must be a string for tool '{tool_name}'"
900 ))
901 .into_response(),
902 );
903 };
904 if policy.argument_allowed(role, tool_name, arg_key, val_str) {
905 return None;
906 }
907 tracing::warn!(
912 user = %identity_name,
913 role = %role,
914 tool = tool_name,
915 argument = arg_key,
916 arg_hmac = %policy.redact_arg(val_str),
917 "argument not in allowlist"
918 );
919 Some(
920 McpxError::Rbac(format!(
921 "argument '{arg_key}' value not in allowlist for tool '{tool_name}'"
922 ))
923 .into_response(),
924 )
925}
926
927fn json_value_type(v: &serde_json::Value) -> &'static str {
928 match v {
929 serde_json::Value::Null => "null",
930 serde_json::Value::Bool(_) => "bool",
931 serde_json::Value::Number(_) => "number",
932 serde_json::Value::String(_) => "string",
933 serde_json::Value::Array(_) => "array",
934 serde_json::Value::Object(_) => "object",
935 }
936}
937
938fn glob_match(pattern: &str, text: &str) -> bool {
948 let parts: Vec<&str> = pattern.split('*').collect();
949 if parts.len() == 1 {
950 return pattern == text;
952 }
953
954 let pos = if let Some(&first) = parts.first()
956 && !first.is_empty()
957 {
958 if !text.starts_with(first) {
959 return false;
960 }
961 first.len()
962 } else {
963 0
964 };
965
966 if let Some(&last) = parts.last()
968 && !last.is_empty()
969 {
970 if !text.get(pos..).unwrap_or_default().ends_with(last) {
971 return false;
972 }
973 let end = text.len() - last.len();
975 if pos > end {
976 return false;
977 }
978 let middle = text.get(pos..end).unwrap_or_default();
980 let middle_parts = parts.get(1..parts.len() - 1).unwrap_or_default();
981 return match_middle(middle, middle_parts);
982 }
983
984 let middle = text.get(pos..).unwrap_or_default();
986 let middle_parts = parts.get(1..parts.len() - 1).unwrap_or_default();
987 match_middle(middle, middle_parts)
988}
989
990fn match_middle(mut text: &str, parts: &[&str]) -> bool {
992 for part in parts {
993 if part.is_empty() {
994 continue;
995 }
996 if let Some(idx) = text.find(part) {
997 text = text.get(idx + part.len()..).unwrap_or_default();
998 } else {
999 return false;
1000 }
1001 }
1002 true
1003}
1004
1005#[cfg(test)]
1006mod tests {
1007 use super::*;
1008
1009 #[test]
1014 fn tool_limiter_burst_allows_initial_spike() {
1015 let limiter = build_tool_rate_limiter(2, Some(4));
1016 let ip: IpAddr = "10.9.9.9".parse().unwrap();
1017 for i in 0..4 {
1018 assert!(
1019 limiter.check_key(&ip).is_ok(),
1020 "burst request {i} should pass"
1021 );
1022 }
1023 assert!(
1024 limiter.check_key(&ip).is_err(),
1025 "request 5 must exceed the burst bucket"
1026 );
1027 }
1028
1029 #[test]
1031 fn tool_limiter_deny_sets_retry_after() {
1032 let limiter = build_tool_rate_limiter(1, None);
1033 let ip: IpAddr = "10.8.8.8".parse().unwrap();
1034 assert!(enforce_rate_limit(Some(&limiter), Some(ip)).is_none());
1035 let resp = enforce_rate_limit(Some(&limiter), Some(ip))
1036 .expect("second call within the window must deny");
1037 assert_eq!(resp.status(), axum::http::StatusCode::TOO_MANY_REQUESTS);
1038 let retry_after = resp
1039 .headers()
1040 .get(axum::http::header::RETRY_AFTER)
1041 .expect("Retry-After present")
1042 .to_str()
1043 .unwrap()
1044 .parse::<u64>()
1045 .unwrap();
1046 assert!(retry_after >= 1, "delta-seconds must be >= 1");
1047 }
1048
1049 fn test_policy() -> RbacPolicy {
1050 RbacPolicy::new(&RbacConfig {
1051 enabled: true,
1052 roles: vec![
1053 RoleConfig {
1054 name: "viewer".into(),
1055 description: Some("Read-only".into()),
1056 allow: vec![
1057 "list_hosts".into(),
1058 "resource_list".into(),
1059 "resource_inspect".into(),
1060 "resource_logs".into(),
1061 "system_info".into(),
1062 ],
1063 deny: vec![],
1064 hosts: vec!["*".into()],
1065 argument_allowlists: vec![],
1066 },
1067 RoleConfig {
1068 name: "deploy".into(),
1069 description: Some("Lifecycle management".into()),
1070 allow: vec![
1071 "list_hosts".into(),
1072 "resource_list".into(),
1073 "resource_run".into(),
1074 "resource_start".into(),
1075 "resource_stop".into(),
1076 "resource_restart".into(),
1077 "resource_logs".into(),
1078 "image_pull".into(),
1079 ],
1080 deny: vec!["resource_delete".into(), "resource_exec".into()],
1081 hosts: vec!["web-*".into(), "api-*".into()],
1082 argument_allowlists: vec![],
1083 },
1084 RoleConfig {
1085 name: "ops".into(),
1086 description: Some("Full access".into()),
1087 allow: vec!["*".into()],
1088 deny: vec![],
1089 hosts: vec!["*".into()],
1090 argument_allowlists: vec![],
1091 },
1092 RoleConfig {
1093 name: "restricted-exec".into(),
1094 description: Some("Exec with argument allowlist".into()),
1095 allow: vec!["resource_exec".into()],
1096 deny: vec![],
1097 hosts: vec!["dev-*".into()],
1098 argument_allowlists: vec![ArgumentAllowlist {
1099 tool: "resource_exec".into(),
1100 argument: "cmd".into(),
1101 allowed: vec![
1102 "sh".into(),
1103 "bash".into(),
1104 "cat".into(),
1105 "ls".into(),
1106 "ps".into(),
1107 ],
1108 }],
1109 },
1110 ],
1111 redaction_salt: None,
1112 })
1113 }
1114
1115 #[test]
1118 fn glob_exact_match() {
1119 assert!(glob_match("web-prod-1", "web-prod-1"));
1120 assert!(!glob_match("web-prod-1", "web-prod-2"));
1121 }
1122
1123 #[test]
1124 fn glob_star_suffix() {
1125 assert!(glob_match("web-*", "web-prod-1"));
1126 assert!(glob_match("web-*", "web-staging"));
1127 assert!(!glob_match("web-*", "api-prod"));
1128 }
1129
1130 #[test]
1131 fn glob_star_prefix() {
1132 assert!(glob_match("*-prod", "web-prod"));
1133 assert!(glob_match("*-prod", "api-prod"));
1134 assert!(!glob_match("*-prod", "web-staging"));
1135 }
1136
1137 #[test]
1138 fn glob_star_middle() {
1139 assert!(glob_match("web-*-prod", "web-us-prod"));
1140 assert!(glob_match("web-*-prod", "web-eu-east-prod"));
1141 assert!(!glob_match("web-*-prod", "web-staging"));
1142 }
1143
1144 #[test]
1145 fn glob_star_only() {
1146 assert!(glob_match("*", "anything"));
1147 assert!(glob_match("*", ""));
1148 }
1149
1150 #[test]
1151 fn glob_multiple_stars() {
1152 assert!(glob_match("*web*prod*", "my-web-us-prod-1"));
1153 assert!(!glob_match("*web*prod*", "my-api-us-staging"));
1154 }
1155
1156 #[test]
1161 fn glob_match_multibyte_utf8() {
1162 assert!(glob_match("hé*llo", "héllo"));
1163 assert!(glob_match("*ö*", "wörld"));
1164 assert!(glob_match("über*", "übermensch"));
1165 assert!(glob_match("*界", "世界"));
1166 assert!(!glob_match("hé*llo", "hello"));
1167 assert!(!glob_match("界*", "世界"));
1168 assert!(glob_match("世*界", "世界"));
1169 }
1170
1171 #[test]
1183 fn glob_prefix_and_suffix_meet_exactly() {
1184 assert!(glob_match("ab*cd", "abcd"));
1187 }
1188
1189 #[test]
1194 fn glob_middle_segment_required_with_suffix() {
1195 assert!(!glob_match("a*b*c", "axyc"));
1200 }
1201
1202 #[test]
1208 fn glob_match_middle_advances_past_matched_part() {
1209 assert!(!glob_match("*ab*ab*", "xxab_yz"));
1214 }
1215
1216 #[test]
1221 fn glob_match_middle_uses_addition_not_multiplication() {
1222 assert!(glob_match("*abcde*X*", "yyyyyyyyabcde_X"));
1226 }
1227
1228 #[test]
1237 fn argument_allowed_glob_pattern_with_literal_mismatch_still_enforced() {
1238 let role = RoleConfig::new("viewer", vec!["run-foo".into()], vec!["*".into()])
1246 .with_argument_allowlists(vec![ArgumentAllowlist::new(
1247 "run-*",
1248 "cmd",
1249 vec!["ls".into()],
1250 )]);
1251 let mut config = RbacConfig::with_roles(vec![role]);
1252 config.enabled = true;
1253 let policy = RbacPolicy::new(&config);
1254 assert!(!policy.argument_allowed("viewer", "run-foo", "cmd", "rm"));
1255 }
1256
1257 #[test]
1260 fn disabled_policy_allows_everything() {
1261 let policy = RbacPolicy::new(&RbacConfig {
1262 enabled: false,
1263 roles: vec![],
1264 redaction_salt: None,
1265 });
1266 assert_eq!(
1267 policy.check("nonexistent", "resource_delete", "any-host"),
1268 RbacDecision::Allow
1269 );
1270 }
1271
1272 #[test]
1273 fn unknown_role_denied() {
1274 let policy = test_policy();
1275 assert_eq!(
1276 policy.check("unknown", "resource_list", "web-prod-1"),
1277 RbacDecision::Deny
1278 );
1279 }
1280
1281 #[test]
1282 fn viewer_allowed_read_ops() {
1283 let policy = test_policy();
1284 assert_eq!(
1285 policy.check("viewer", "resource_list", "web-prod-1"),
1286 RbacDecision::Allow
1287 );
1288 assert_eq!(
1289 policy.check("viewer", "system_info", "db-host"),
1290 RbacDecision::Allow
1291 );
1292 }
1293
1294 #[test]
1295 fn viewer_denied_write_ops() {
1296 let policy = test_policy();
1297 assert_eq!(
1298 policy.check("viewer", "resource_run", "web-prod-1"),
1299 RbacDecision::Deny
1300 );
1301 assert_eq!(
1302 policy.check("viewer", "resource_delete", "web-prod-1"),
1303 RbacDecision::Deny
1304 );
1305 }
1306
1307 #[test]
1308 fn deploy_allowed_on_matching_hosts() {
1309 let policy = test_policy();
1310 assert_eq!(
1311 policy.check("deploy", "resource_run", "web-prod-1"),
1312 RbacDecision::Allow
1313 );
1314 assert_eq!(
1315 policy.check("deploy", "resource_start", "api-staging"),
1316 RbacDecision::Allow
1317 );
1318 }
1319
1320 #[test]
1321 fn deploy_denied_on_non_matching_host() {
1322 let policy = test_policy();
1323 assert_eq!(
1324 policy.check("deploy", "resource_run", "db-prod-1"),
1325 RbacDecision::Deny
1326 );
1327 }
1328
1329 #[test]
1330 fn deny_overrides_allow() {
1331 let policy = test_policy();
1332 assert_eq!(
1333 policy.check("deploy", "resource_delete", "web-prod-1"),
1334 RbacDecision::Deny
1335 );
1336 assert_eq!(
1337 policy.check("deploy", "resource_exec", "web-prod-1"),
1338 RbacDecision::Deny
1339 );
1340 }
1341
1342 #[test]
1343 fn ops_wildcard_allows_everything() {
1344 let policy = test_policy();
1345 assert_eq!(
1346 policy.check("ops", "resource_delete", "any-host"),
1347 RbacDecision::Allow
1348 );
1349 assert_eq!(
1350 policy.check("ops", "secret_create", "db-host"),
1351 RbacDecision::Allow
1352 );
1353 }
1354
1355 #[test]
1358 fn host_visible_respects_globs() {
1359 let policy = test_policy();
1360 assert!(policy.host_visible("deploy", "web-prod-1"));
1361 assert!(policy.host_visible("deploy", "api-staging"));
1362 assert!(!policy.host_visible("deploy", "db-prod-1"));
1363 assert!(policy.host_visible("ops", "anything"));
1364 assert!(policy.host_visible("viewer", "anything"));
1365 }
1366
1367 #[test]
1368 fn host_visible_unknown_role() {
1369 let policy = test_policy();
1370 assert!(!policy.host_visible("unknown", "web-prod-1"));
1371 }
1372
1373 #[test]
1376 fn argument_allowed_no_allowlist() {
1377 let policy = test_policy();
1378 assert!(policy.argument_allowed("ops", "resource_exec", "cmd", "rm -rf /"));
1380 assert!(policy.argument_allowed("ops", "resource_exec", "cmd", "bash"));
1381 }
1382
1383 #[test]
1384 fn argument_allowed_with_allowlist() {
1385 let policy = test_policy();
1386 assert!(policy.argument_allowed("restricted-exec", "resource_exec", "cmd", "sh"));
1387 assert!(policy.argument_allowed(
1388 "restricted-exec",
1389 "resource_exec",
1390 "cmd",
1391 "bash -c 'echo hi'"
1392 ));
1393 assert!(policy.argument_allowed(
1394 "restricted-exec",
1395 "resource_exec",
1396 "cmd",
1397 "cat /etc/hosts"
1398 ));
1399 assert!(policy.argument_allowed(
1400 "restricted-exec",
1401 "resource_exec",
1402 "cmd",
1403 "/usr/bin/ls -la"
1404 ));
1405 }
1406
1407 #[test]
1408 fn argument_denied_not_in_allowlist() {
1409 let policy = test_policy();
1410 assert!(!policy.argument_allowed("restricted-exec", "resource_exec", "cmd", "rm -rf /"));
1411 assert!(!policy.argument_allowed(
1412 "restricted-exec",
1413 "resource_exec",
1414 "cmd",
1415 "python3 exploit.py"
1416 ));
1417 assert!(!policy.argument_allowed(
1418 "restricted-exec",
1419 "resource_exec",
1420 "cmd",
1421 "/usr/bin/curl evil.com"
1422 ));
1423 }
1424
1425 #[test]
1426 fn argument_denied_unknown_role() {
1427 let policy = test_policy();
1428 assert!(!policy.argument_allowed("unknown", "resource_exec", "cmd", "sh"));
1429 }
1430
1431 fn shlex_policy(allowed: Vec<String>) -> RbacPolicy {
1440 let role = RoleConfig::new("viewer", vec!["run".into()], vec!["*".into()])
1441 .with_argument_allowlists(vec![ArgumentAllowlist::new("run", "cmd", allowed)]);
1442 let mut config = RbacConfig::with_roles(vec![role]);
1443 config.enabled = true;
1444 RbacPolicy::new(&config)
1445 }
1446
1447 #[test]
1448 fn argument_allowed_matches_quoted_path_with_spaces() {
1449 let policy = shlex_policy(vec!["/usr/bin/my tool".into()]);
1450 assert!(policy.argument_allowed("viewer", "run", "cmd", r#""/usr/bin/my tool" --flag"#));
1451 }
1452
1453 #[test]
1454 fn argument_allowed_matches_basename_of_quoted_path() {
1455 let policy = shlex_policy(vec!["my tool".into()]);
1456 assert!(policy.argument_allowed("viewer", "run", "cmd", r#""/usr/bin/my tool" --flag"#));
1457 }
1458
1459 #[test]
1460 fn argument_allowed_fails_closed_on_unbalanced_quote() {
1461 let policy = shlex_policy(vec!["unbalanced".into()]);
1462 assert!(!policy.argument_allowed("viewer", "run", "cmd", r"unbalanced 'quote"));
1463 }
1464
1465 #[test]
1466 fn argument_allowed_fails_closed_on_empty_string() {
1467 let policy = shlex_policy(vec![String::new()]);
1468 assert!(!policy.argument_allowed("viewer", "run", "cmd", ""));
1469 }
1470
1471 #[test]
1472 fn argument_allowed_handles_single_quoted_executable() {
1473 let policy = shlex_policy(vec!["/bin/sh".into()]);
1474 assert!(policy.argument_allowed("viewer", "run", "cmd", r"'/bin/sh' -c 'echo hi'"));
1475 }
1476
1477 #[test]
1478 fn argument_allowed_handles_tab_separator() {
1479 let policy = shlex_policy(vec!["ls".into()]);
1480 assert!(policy.argument_allowed("viewer", "run", "cmd", "ls\t/etc/passwd"));
1481 }
1482
1483 #[test]
1484 fn argument_allowed_plain_token_unchanged() {
1485 let policy = shlex_policy(vec!["ls".into()]);
1486 assert!(policy.argument_allowed("viewer", "run", "cmd", "ls"));
1487 }
1488
1489 #[test]
1495 fn argument_allowed_fails_closed_on_quoted_empty_first_token() {
1496 let policy = shlex_policy(vec![String::new()]);
1500 assert!(!policy.argument_allowed("viewer", "run", "cmd", r#""""#));
1501 }
1502
1503 #[test]
1504 fn argument_allowed_quoted_literal_token_no_longer_matches() {
1505 let policy = shlex_policy(vec!["'bash'".into()]);
1511 assert!(!policy.argument_allowed("viewer", "run", "cmd", "'bash' -c true"));
1512 }
1513
1514 #[test]
1515 fn argument_allowed_backslash_literal_token_no_longer_matches() {
1516 let policy = shlex_policy(vec![r"foo\bar".into()]);
1521 assert!(!policy.argument_allowed("viewer", "run", "cmd", r"foo\bar --x"));
1522 }
1523
1524 #[test]
1525 fn argument_allowed_windows_path_no_longer_matches() {
1526 let policy = shlex_policy(vec![r"C:\Windows\System32\cmd.exe".into()]);
1531 assert!(!policy.argument_allowed(
1532 "viewer",
1533 "run",
1534 "cmd",
1535 r"C:\Windows\System32\cmd.exe /c dir"
1536 ));
1537 }
1538
1539 #[test]
1542 fn host_patterns_returns_globs() {
1543 let policy = test_policy();
1544 assert_eq!(
1545 policy.host_patterns("deploy"),
1546 Some(vec!["web-*".to_owned(), "api-*".to_owned()].as_slice())
1547 );
1548 assert_eq!(
1549 policy.host_patterns("ops"),
1550 Some(vec!["*".to_owned()].as_slice())
1551 );
1552 assert!(policy.host_patterns("nonexistent").is_none());
1553 }
1554
1555 #[test]
1558 fn check_operation_allows_without_host() {
1559 let policy = test_policy();
1560 assert_eq!(
1561 policy.check_operation("deploy", "resource_run"),
1562 RbacDecision::Allow
1563 );
1564 assert_eq!(
1566 policy.check("deploy", "resource_run", "db-prod-1"),
1567 RbacDecision::Deny
1568 );
1569 }
1570
1571 #[test]
1572 fn check_operation_deny_overrides() {
1573 let policy = test_policy();
1574 assert_eq!(
1575 policy.check_operation("deploy", "resource_delete"),
1576 RbacDecision::Deny
1577 );
1578 }
1579
1580 #[test]
1581 fn check_operation_unknown_role() {
1582 let policy = test_policy();
1583 assert_eq!(
1584 policy.check_operation("unknown", "resource_list"),
1585 RbacDecision::Deny
1586 );
1587 }
1588
1589 #[test]
1590 fn check_operation_disabled() {
1591 let policy = RbacPolicy::new(&RbacConfig {
1592 enabled: false,
1593 roles: vec![],
1594 redaction_salt: None,
1595 });
1596 assert_eq!(
1597 policy.check_operation("nonexistent", "anything"),
1598 RbacDecision::Allow
1599 );
1600 }
1601
1602 #[test]
1605 fn current_role_returns_none_outside_scope() {
1606 assert!(current_role().is_none());
1607 }
1608
1609 #[test]
1610 fn current_identity_returns_none_outside_scope() {
1611 assert!(current_identity().is_none());
1612 }
1613
1614 use axum::{
1617 body::Body,
1618 http::{Method, Request, StatusCode},
1619 };
1620 use tower::ServiceExt as _;
1621
1622 fn tool_call_body(tool: &str, args: &serde_json::Value) -> String {
1623 serde_json::json!({
1624 "jsonrpc": "2.0",
1625 "id": 1,
1626 "method": "tools/call",
1627 "params": {
1628 "name": tool,
1629 "arguments": args
1630 }
1631 })
1632 .to_string()
1633 }
1634
1635 fn rbac_router(policy: Arc<RbacPolicy>) -> axum::Router {
1636 axum::Router::new()
1637 .route("/mcp", axum::routing::post(|| async { "ok" }))
1638 .layer(axum::middleware::from_fn(move |req, next| {
1639 let p = Arc::clone(&policy);
1640 rbac_middleware(p, None, req, next)
1641 }))
1642 }
1643
1644 fn rbac_router_with_identity(policy: Arc<RbacPolicy>, identity: AuthIdentity) -> axum::Router {
1645 axum::Router::new()
1646 .route("/mcp", axum::routing::post(|| async { "ok" }))
1647 .layer(axum::middleware::from_fn(
1648 move |mut req: Request<Body>, next: Next| {
1649 let p = Arc::clone(&policy);
1650 let id = identity.clone();
1651 async move {
1652 req.extensions_mut().insert(id);
1653 rbac_middleware(p, None, req, next).await
1654 }
1655 },
1656 ))
1657 }
1658
1659 #[cfg(feature = "metrics")]
1663 #[tokio::test]
1664 async fn tool_limiter_deny_increments_counter() {
1665 use axum::extract::ConnectInfo;
1666
1667 let policy = Arc::new(test_policy());
1668 let limiter = build_tool_rate_limiter(1, None);
1669 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
1670 let identity = AuthIdentity {
1671 method: crate::auth::AuthMethod::BearerToken,
1672 name: "alice".into(),
1673 role: "viewer".into(),
1674 raw_token: None,
1675 sub: None,
1676 };
1677 let app = {
1678 let metrics = Arc::clone(&metrics);
1679 axum::Router::new()
1680 .route("/mcp", axum::routing::post(|| async { "ok" }))
1681 .layer(axum::middleware::from_fn(
1682 move |mut req: Request<Body>, next: Next| {
1683 let p = Arc::clone(&policy);
1684 let l = Arc::clone(&limiter);
1685 let id = identity.clone();
1686 let m = Arc::clone(&metrics);
1687 async move {
1688 req.extensions_mut().insert(id);
1689 req.extensions_mut().insert(m);
1690 let peer: std::net::SocketAddr =
1691 "10.9.9.1:40000".parse().expect("static socket addr parses");
1692 req.extensions_mut().insert(ConnectInfo(peer));
1693 rbac_middleware(p, Some(l), req, next).await
1694 }
1695 },
1696 ))
1697 };
1698 let mk = || {
1699 Request::builder()
1700 .method(Method::POST)
1701 .uri("/mcp")
1702 .header("content-type", "application/json")
1703 .body(Body::from(tool_call_body(
1704 "resource_list",
1705 &serde_json::json!({}),
1706 )))
1707 .unwrap()
1708 };
1709 let counter = || {
1710 metrics
1711 .rate_limited_total
1712 .with_label_values(&["tool"])
1713 .get()
1714 };
1715
1716 let first = app.clone().oneshot(mk()).await.unwrap();
1717 assert_eq!(first.status(), StatusCode::OK);
1718 assert_eq!(counter(), 0, "successful call must not count");
1719
1720 let denied = app.clone().oneshot(mk()).await.unwrap();
1721 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
1722 assert_eq!(counter(), 1, "deny must increment the tool label");
1723 }
1724
1725 #[tokio::test]
1726 async fn middleware_passes_non_post() {
1727 let policy = Arc::new(test_policy());
1728 let app = rbac_router(policy);
1729 let req = Request::builder()
1731 .method(Method::GET)
1732 .uri("/mcp")
1733 .body(Body::empty())
1734 .unwrap();
1735 let resp = app.oneshot(req).await.unwrap();
1738 assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
1739 }
1740
1741 #[tokio::test]
1742 async fn middleware_denies_without_identity() {
1743 let policy = Arc::new(test_policy());
1744 let app = rbac_router(policy);
1745 let body = tool_call_body("resource_list", &serde_json::json!({}));
1746 let req = Request::builder()
1747 .method(Method::POST)
1748 .uri("/mcp")
1749 .header("content-type", "application/json")
1750 .body(Body::from(body))
1751 .unwrap();
1752 let resp = app.oneshot(req).await.unwrap();
1753 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1754 }
1755
1756 #[tokio::test]
1757 async fn middleware_allows_permitted_tool() {
1758 let policy = Arc::new(test_policy());
1759 let id = AuthIdentity {
1760 method: crate::auth::AuthMethod::BearerToken,
1761 name: "alice".into(),
1762 role: "viewer".into(),
1763 raw_token: None,
1764 sub: None,
1765 };
1766 let app = rbac_router_with_identity(policy, id);
1767 let body = tool_call_body("resource_list", &serde_json::json!({}));
1768 let req = Request::builder()
1769 .method(Method::POST)
1770 .uri("/mcp")
1771 .header("content-type", "application/json")
1772 .body(Body::from(body))
1773 .unwrap();
1774 let resp = app.oneshot(req).await.unwrap();
1775 assert_eq!(resp.status(), StatusCode::OK);
1776 }
1777
1778 #[tokio::test]
1779 async fn middleware_denies_unpermitted_tool() {
1780 let policy = Arc::new(test_policy());
1781 let id = AuthIdentity {
1782 method: crate::auth::AuthMethod::BearerToken,
1783 name: "alice".into(),
1784 role: "viewer".into(),
1785 raw_token: None,
1786 sub: None,
1787 };
1788 let app = rbac_router_with_identity(policy, id);
1789 let body = tool_call_body("resource_delete", &serde_json::json!({}));
1790 let req = Request::builder()
1791 .method(Method::POST)
1792 .uri("/mcp")
1793 .header("content-type", "application/json")
1794 .body(Body::from(body))
1795 .unwrap();
1796 let resp = app.oneshot(req).await.unwrap();
1797 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1798 }
1799
1800 #[tokio::test]
1801 async fn middleware_passes_non_tool_call_post() {
1802 let policy = Arc::new(test_policy());
1803 let id = AuthIdentity {
1804 method: crate::auth::AuthMethod::BearerToken,
1805 name: "alice".into(),
1806 role: "viewer".into(),
1807 raw_token: None,
1808 sub: None,
1809 };
1810 let app = rbac_router_with_identity(policy, id);
1811 let body = serde_json::json!({
1813 "jsonrpc": "2.0",
1814 "id": 1,
1815 "method": "resources/list"
1816 })
1817 .to_string();
1818 let req = Request::builder()
1819 .method(Method::POST)
1820 .uri("/mcp")
1821 .header("content-type", "application/json")
1822 .body(Body::from(body))
1823 .unwrap();
1824 let resp = app.oneshot(req).await.unwrap();
1825 assert_eq!(resp.status(), StatusCode::OK);
1826 }
1827
1828 #[tokio::test]
1829 async fn middleware_enforces_argument_allowlist() {
1830 let policy = Arc::new(test_policy());
1831 let id = AuthIdentity {
1832 method: crate::auth::AuthMethod::BearerToken,
1833 name: "dev".into(),
1834 role: "restricted-exec".into(),
1835 raw_token: None,
1836 sub: None,
1837 };
1838 let app = rbac_router_with_identity(Arc::clone(&policy), id.clone());
1840 let body = tool_call_body(
1841 "resource_exec",
1842 &serde_json::json!({"cmd": "ls -la", "host": "dev-1"}),
1843 );
1844 let req = Request::builder()
1845 .method(Method::POST)
1846 .uri("/mcp")
1847 .body(Body::from(body))
1848 .unwrap();
1849 let resp = app.oneshot(req).await.unwrap();
1850 assert_eq!(resp.status(), StatusCode::OK);
1851
1852 let app = rbac_router_with_identity(policy, id);
1854 let body = tool_call_body(
1855 "resource_exec",
1856 &serde_json::json!({"cmd": "rm -rf /", "host": "dev-1"}),
1857 );
1858 let req = Request::builder()
1859 .method(Method::POST)
1860 .uri("/mcp")
1861 .body(Body::from(body))
1862 .unwrap();
1863 let resp = app.oneshot(req).await.unwrap();
1864 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1865 }
1866
1867 #[tokio::test]
1868 async fn middleware_disabled_policy_passes_everything() {
1869 let policy = Arc::new(RbacPolicy::disabled());
1870 let app = rbac_router(policy);
1871 let body = tool_call_body("anything", &serde_json::json!({}));
1873 let req = Request::builder()
1874 .method(Method::POST)
1875 .uri("/mcp")
1876 .body(Body::from(body))
1877 .unwrap();
1878 let resp = app.oneshot(req).await.unwrap();
1879 assert_eq!(resp.status(), StatusCode::OK);
1880 }
1881
1882 #[tokio::test]
1883 async fn middleware_batch_all_allowed_passes() {
1884 let policy = Arc::new(test_policy());
1885 let id = AuthIdentity {
1886 method: crate::auth::AuthMethod::BearerToken,
1887 name: "alice".into(),
1888 role: "viewer".into(),
1889 raw_token: None,
1890 sub: None,
1891 };
1892 let app = rbac_router_with_identity(policy, id);
1893 let body = serde_json::json!([
1894 {
1895 "jsonrpc": "2.0",
1896 "id": 1,
1897 "method": "tools/call",
1898 "params": { "name": "resource_list", "arguments": {} }
1899 },
1900 {
1901 "jsonrpc": "2.0",
1902 "id": 2,
1903 "method": "tools/call",
1904 "params": { "name": "system_info", "arguments": {} }
1905 }
1906 ])
1907 .to_string();
1908 let req = Request::builder()
1909 .method(Method::POST)
1910 .uri("/mcp")
1911 .header("content-type", "application/json")
1912 .body(Body::from(body))
1913 .unwrap();
1914 let resp = app.oneshot(req).await.unwrap();
1915 assert_eq!(resp.status(), StatusCode::OK);
1916 }
1917
1918 #[tokio::test]
1919 async fn middleware_batch_with_denied_call_rejects_entire_batch() {
1920 let policy = Arc::new(test_policy());
1921 let id = AuthIdentity {
1922 method: crate::auth::AuthMethod::BearerToken,
1923 name: "alice".into(),
1924 role: "viewer".into(),
1925 raw_token: None,
1926 sub: None,
1927 };
1928 let app = rbac_router_with_identity(policy, id);
1929 let body = serde_json::json!([
1930 {
1931 "jsonrpc": "2.0",
1932 "id": 1,
1933 "method": "tools/call",
1934 "params": { "name": "resource_list", "arguments": {} }
1935 },
1936 {
1937 "jsonrpc": "2.0",
1938 "id": 2,
1939 "method": "tools/call",
1940 "params": { "name": "resource_delete", "arguments": {} }
1941 }
1942 ])
1943 .to_string();
1944 let req = Request::builder()
1945 .method(Method::POST)
1946 .uri("/mcp")
1947 .header("content-type", "application/json")
1948 .body(Body::from(body))
1949 .unwrap();
1950 let resp = app.oneshot(req).await.unwrap();
1951 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1952 }
1953
1954 #[tokio::test]
1955 async fn middleware_batch_mixed_allowed_and_denied_rejects() {
1956 let policy = Arc::new(test_policy());
1957 let id = AuthIdentity {
1958 method: crate::auth::AuthMethod::BearerToken,
1959 name: "dev".into(),
1960 role: "restricted-exec".into(),
1961 raw_token: None,
1962 sub: None,
1963 };
1964 let app = rbac_router_with_identity(policy, id);
1965 let body = serde_json::json!([
1966 {
1967 "jsonrpc": "2.0",
1968 "id": 1,
1969 "method": "tools/call",
1970 "params": {
1971 "name": "resource_exec",
1972 "arguments": { "cmd": "ls -la", "host": "dev-1" }
1973 }
1974 },
1975 {
1976 "jsonrpc": "2.0",
1977 "id": 2,
1978 "method": "tools/call",
1979 "params": {
1980 "name": "resource_exec",
1981 "arguments": { "cmd": "rm -rf /", "host": "dev-1" }
1982 }
1983 }
1984 ])
1985 .to_string();
1986 let req = Request::builder()
1987 .method(Method::POST)
1988 .uri("/mcp")
1989 .header("content-type", "application/json")
1990 .body(Body::from(body))
1991 .unwrap();
1992 let resp = app.oneshot(req).await.unwrap();
1993 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1994 }
1995
1996 #[test]
1999 fn redact_with_salt_is_deterministic_per_salt() {
2000 let salt = b"unit-test-salt";
2001 let a = redact_with_salt(salt, "rm -rf /");
2002 let b = redact_with_salt(salt, "rm -rf /");
2003 assert_eq!(a, b, "same input + salt must yield identical hash");
2004 assert_eq!(a.len(), 8, "redacted hash is 8 hex chars (4 bytes)");
2005 assert!(
2006 a.chars().all(|c| c.is_ascii_hexdigit()),
2007 "redacted hash must be lowercase hex: {a}"
2008 );
2009 }
2010
2011 #[test]
2012 fn redact_with_salt_differs_across_salts() {
2013 let v = "the-same-value";
2014 let h1 = redact_with_salt(b"salt-one", v);
2015 let h2 = redact_with_salt(b"salt-two", v);
2016 assert_ne!(
2017 h1, h2,
2018 "different salts must produce different hashes for the same value"
2019 );
2020 }
2021
2022 #[test]
2023 fn redact_with_salt_distinguishes_values() {
2024 let salt = b"k";
2025 let h1 = redact_with_salt(salt, "alpha");
2026 let h2 = redact_with_salt(salt, "beta");
2027 assert_ne!(h1, h2, "different values must produce different hashes");
2029 }
2030
2031 #[test]
2032 fn policy_with_configured_salt_redacts_consistently() {
2033 let cfg = RbacConfig {
2034 enabled: true,
2035 roles: vec![],
2036 redaction_salt: Some(SecretString::from("my-stable-salt")),
2037 };
2038 let p1 = RbacPolicy::new(&cfg);
2039 let p2 = RbacPolicy::new(&cfg);
2040 assert_eq!(
2041 p1.redact_arg("payload"),
2042 p2.redact_arg("payload"),
2043 "policies built from the same configured salt must agree"
2044 );
2045 }
2046
2047 #[test]
2048 fn policy_without_configured_salt_uses_process_salt() {
2049 let cfg = RbacConfig {
2050 enabled: true,
2051 roles: vec![],
2052 redaction_salt: None,
2053 };
2054 let p1 = RbacPolicy::new(&cfg);
2055 let p2 = RbacPolicy::new(&cfg);
2056 assert_eq!(
2058 p1.redact_arg("payload"),
2059 p2.redact_arg("payload"),
2060 "process-wide salt must be consistent within one process"
2061 );
2062 }
2063
2064 #[test]
2065 fn redact_arg_is_fast_enough() {
2066 let salt = b"perf-sanity-salt-32-bytes-padded";
2070 let value = "x".repeat(256);
2071 let start = std::time::Instant::now();
2072 let _ = redact_with_salt(salt, &value);
2073 let elapsed = start.elapsed();
2074 assert!(
2075 elapsed < Duration::from_millis(5),
2076 "single redact_with_salt took {elapsed:?}, expected <5 ms even in debug"
2077 );
2078 }
2079
2080 #[tokio::test]
2092 async fn deny_path_uses_explicit_identity_not_task_local() {
2093 let policy = Arc::new(test_policy());
2094 let id = AuthIdentity {
2095 method: crate::auth::AuthMethod::BearerToken,
2096 name: "alice-the-auditor".into(),
2097 role: "viewer".into(),
2098 raw_token: None,
2099 sub: None,
2100 };
2101 let app = rbac_router_with_identity(policy, id);
2102 let body = tool_call_body("resource_delete", &serde_json::json!({}));
2104 let req = Request::builder()
2105 .method(Method::POST)
2106 .uri("/mcp")
2107 .header("content-type", "application/json")
2108 .body(Body::from(body))
2109 .unwrap();
2110 let resp = app.oneshot(req).await.unwrap();
2111 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2112 }
2113
2114 fn restricted_exec_identity() -> AuthIdentity {
2117 AuthIdentity {
2118 method: crate::auth::AuthMethod::BearerToken,
2119 name: "carol".into(),
2120 role: "restricted-exec".into(),
2121 raw_token: None,
2122 sub: None,
2123 }
2124 }
2125
2126 #[test]
2127 fn has_argument_allowlist_matches_configured_tool_argument() {
2128 let policy = test_policy();
2129 assert!(policy.has_argument_allowlist("restricted-exec", "resource_exec", "cmd"));
2130 assert!(!policy.has_argument_allowlist("restricted-exec", "resource_exec", "host"));
2131 assert!(!policy.has_argument_allowlist("restricted-exec", "other_tool", "cmd"));
2132 assert!(!policy.has_argument_allowlist("ops", "resource_exec", "cmd"));
2133 }
2134
2135 #[tokio::test]
2136 async fn array_arg_with_matching_allowlist_is_denied() {
2137 let policy = Arc::new(test_policy());
2138 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2139 let body = tool_call_body(
2140 "resource_exec",
2141 &serde_json::json!({ "host": "dev-1", "cmd": ["bash", "-c", "evil"] }),
2142 );
2143 let req = Request::builder()
2144 .method(Method::POST)
2145 .uri("/mcp")
2146 .header("content-type", "application/json")
2147 .body(Body::from(body))
2148 .unwrap();
2149 let resp = app.oneshot(req).await.unwrap();
2150 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2151 }
2152
2153 #[tokio::test]
2154 async fn object_arg_with_matching_allowlist_is_denied() {
2155 let policy = Arc::new(test_policy());
2156 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2157 let body = tool_call_body(
2158 "resource_exec",
2159 &serde_json::json!({ "host": "dev-1", "cmd": { "raw": "sh" } }),
2160 );
2161 let req = Request::builder()
2162 .method(Method::POST)
2163 .uri("/mcp")
2164 .header("content-type", "application/json")
2165 .body(Body::from(body))
2166 .unwrap();
2167 let resp = app.oneshot(req).await.unwrap();
2168 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2169 }
2170
2171 #[tokio::test]
2172 async fn number_arg_with_matching_allowlist_is_denied() {
2173 let policy = Arc::new(test_policy());
2174 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2175 let body = tool_call_body(
2176 "resource_exec",
2177 &serde_json::json!({ "host": "dev-1", "cmd": 42 }),
2178 );
2179 let req = Request::builder()
2180 .method(Method::POST)
2181 .uri("/mcp")
2182 .header("content-type", "application/json")
2183 .body(Body::from(body))
2184 .unwrap();
2185 let resp = app.oneshot(req).await.unwrap();
2186 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2187 }
2188
2189 #[tokio::test]
2190 async fn bool_arg_with_matching_allowlist_is_denied() {
2191 let policy = Arc::new(test_policy());
2192 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2193 let body = tool_call_body(
2194 "resource_exec",
2195 &serde_json::json!({ "host": "dev-1", "cmd": true }),
2196 );
2197 let req = Request::builder()
2198 .method(Method::POST)
2199 .uri("/mcp")
2200 .header("content-type", "application/json")
2201 .body(Body::from(body))
2202 .unwrap();
2203 let resp = app.oneshot(req).await.unwrap();
2204 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2205 }
2206
2207 #[tokio::test]
2208 async fn null_arg_with_matching_allowlist_is_denied() {
2209 let policy = Arc::new(test_policy());
2210 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2211 let body = tool_call_body(
2212 "resource_exec",
2213 &serde_json::json!({ "host": "dev-1", "cmd": null }),
2214 );
2215 let req = Request::builder()
2216 .method(Method::POST)
2217 .uri("/mcp")
2218 .header("content-type", "application/json")
2219 .body(Body::from(body))
2220 .unwrap();
2221 let resp = app.oneshot(req).await.unwrap();
2222 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2223 }
2224
2225 #[tokio::test]
2226 async fn non_string_arg_without_allowlist_is_passthrough() {
2227 let policy = Arc::new(test_policy());
2231 let id = AuthIdentity {
2232 method: crate::auth::AuthMethod::BearerToken,
2233 name: "olivia".into(),
2234 role: "ops".into(),
2235 raw_token: None,
2236 sub: None,
2237 };
2238 let app = rbac_router_with_identity(policy, id);
2239 let body = tool_call_body(
2240 "resource_exec",
2241 &serde_json::json!({ "host": "dev-1", "cmd": ["bash"] }),
2242 );
2243 let req = Request::builder()
2244 .method(Method::POST)
2245 .uri("/mcp")
2246 .header("content-type", "application/json")
2247 .body(Body::from(body))
2248 .unwrap();
2249 let resp = app.oneshot(req).await.unwrap();
2250 assert_ne!(resp.status(), StatusCode::FORBIDDEN);
2251 }
2252
2253 #[tokio::test]
2254 async fn string_arg_in_allowlist_still_passes() {
2255 let policy = Arc::new(test_policy());
2256 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2257 let body = tool_call_body(
2258 "resource_exec",
2259 &serde_json::json!({ "host": "dev-1", "cmd": "bash" }),
2260 );
2261 let req = Request::builder()
2262 .method(Method::POST)
2263 .uri("/mcp")
2264 .header("content-type", "application/json")
2265 .body(Body::from(body))
2266 .unwrap();
2267 let resp = app.oneshot(req).await.unwrap();
2268 assert_ne!(resp.status(), StatusCode::FORBIDDEN);
2269 }
2270}