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)]
261#[non_exhaustive]
262pub struct ArgumentAllowlist {
263 pub tool: String,
265 pub argument: String,
267 #[serde(default)]
269 pub allowed: Vec<String>,
270 #[serde(default)]
283 pub required: bool,
284}
285
286impl ArgumentAllowlist {
287 #[must_use]
292 pub fn new(tool: impl Into<String>, argument: impl Into<String>, allowed: Vec<String>) -> Self {
293 Self {
294 tool: tool.into(),
295 argument: argument.into(),
296 allowed,
297 required: false,
298 }
299 }
300
301 #[must_use]
303 pub const fn with_required(mut self, required: bool) -> Self {
304 self.required = required;
305 self
306 }
307}
308
309fn default_hosts() -> Vec<String> {
310 vec!["*".into()]
311}
312
313#[derive(Debug, Clone, Default, Deserialize)]
315#[non_exhaustive]
316pub struct RbacConfig {
317 #[serde(default)]
319 pub enabled: bool,
320 #[serde(default)]
322 pub roles: Vec<RoleConfig>,
323 #[serde(default)]
332 pub redaction_salt: Option<SecretString>,
333}
334
335impl RbacConfig {
336 #[must_use]
338 pub fn with_roles(roles: Vec<RoleConfig>) -> Self {
339 Self {
340 enabled: true,
341 roles,
342 redaction_salt: None,
343 }
344 }
345}
346
347#[derive(Debug, Clone, Copy, PartialEq, Eq)]
349#[non_exhaustive]
350pub enum RbacDecision {
351 Allow,
353 Deny,
355}
356
357#[derive(Debug, Clone, serde::Serialize)]
359#[non_exhaustive]
360pub struct RbacRoleSummary {
361 pub name: String,
363 pub allow: usize,
365 pub deny: usize,
367 pub hosts: usize,
369 pub argument_allowlists: usize,
371}
372
373#[derive(Debug, Clone, serde::Serialize)]
375#[non_exhaustive]
376pub struct RbacPolicySummary {
377 pub enabled: bool,
379 pub roles: Vec<RbacRoleSummary>,
381}
382
383#[derive(Debug, Clone)]
389#[non_exhaustive]
390pub struct RbacPolicy {
391 roles: Vec<RoleConfig>,
392 enabled: bool,
393 redaction_salt: Arc<SecretString>,
396}
397
398impl RbacPolicy {
399 #[must_use]
402 pub fn new(config: &RbacConfig) -> Self {
403 let salt = config
404 .redaction_salt
405 .clone()
406 .unwrap_or_else(|| process_redaction_salt().clone());
407 Self {
408 roles: config.roles.clone(),
409 enabled: config.enabled,
410 redaction_salt: Arc::new(salt),
411 }
412 }
413
414 #[must_use]
416 pub fn disabled() -> Self {
417 Self {
418 roles: Vec::new(),
419 enabled: false,
420 redaction_salt: Arc::new(process_redaction_salt().clone()),
421 }
422 }
423
424 #[must_use]
426 pub fn is_enabled(&self) -> bool {
427 self.enabled
428 }
429
430 #[must_use]
435 pub fn summary(&self) -> RbacPolicySummary {
436 let roles = self
437 .roles
438 .iter()
439 .map(|r| RbacRoleSummary {
440 name: r.name.clone(),
441 allow: r.allow.len(),
442 deny: r.deny.len(),
443 hosts: r.hosts.len(),
444 argument_allowlists: r.argument_allowlists.len(),
445 })
446 .collect();
447 RbacPolicySummary {
448 enabled: self.enabled,
449 roles,
450 }
451 }
452
453 #[must_use]
458 pub fn check_operation(&self, role: &str, operation: &str) -> RbacDecision {
459 if !self.enabled {
460 return RbacDecision::Allow;
461 }
462 let Some(role_cfg) = self.find_role(role) else {
463 return RbacDecision::Deny;
464 };
465 if role_cfg.deny.iter().any(|d| d == operation) {
466 return RbacDecision::Deny;
467 }
468 if role_cfg.allow.iter().any(|a| a == "*" || a == operation) {
469 return RbacDecision::Allow;
470 }
471 RbacDecision::Deny
472 }
473
474 #[must_use]
481 pub fn check(&self, role: &str, operation: &str, host: &str) -> RbacDecision {
482 if !self.enabled {
483 return RbacDecision::Allow;
484 }
485 let Some(role_cfg) = self.find_role(role) else {
486 return RbacDecision::Deny;
487 };
488 if role_cfg.deny.iter().any(|d| d == operation) {
489 return RbacDecision::Deny;
490 }
491 if !role_cfg.allow.iter().any(|a| a == "*" || a == operation) {
492 return RbacDecision::Deny;
493 }
494 if !Self::host_matches(&role_cfg.hosts, host) {
495 return RbacDecision::Deny;
496 }
497 RbacDecision::Allow
498 }
499
500 #[must_use]
502 pub fn host_visible(&self, role: &str, host: &str) -> bool {
503 if !self.enabled {
504 return true;
505 }
506 let Some(role_cfg) = self.find_role(role) else {
507 return false;
508 };
509 Self::host_matches(&role_cfg.hosts, host)
510 }
511
512 #[must_use]
514 pub fn host_patterns(&self, role: &str) -> Option<&[String]> {
515 self.find_role(role).map(|r| r.hosts.as_slice())
516 }
517
518 #[must_use]
547 pub fn argument_allowed(&self, role: &str, tool: &str, argument: &str, value: &str) -> bool {
548 if !self.enabled {
549 return true;
550 }
551 let Some(role_cfg) = self.find_role(role) else {
552 return false;
553 };
554 for al in &role_cfg.argument_allowlists {
555 if al.tool != tool && !glob_match(&al.tool, tool) {
556 continue;
557 }
558 if al.argument != argument {
559 continue;
560 }
561 if al.allowed.is_empty() {
562 continue;
563 }
564 let Some(tokens) = shlex::split(value) else {
569 return false;
570 };
571 let Some(first_token) = tokens.first() else {
572 return false;
573 };
574 if first_token.is_empty() {
578 return false;
579 }
580 let basename = first_token
584 .rsplit('/')
585 .next()
586 .unwrap_or(first_token.as_str());
587 if !al.allowed.iter().any(|a| a == first_token || a == basename) {
588 return false;
589 }
590 }
591 true
592 }
593
594 #[must_use]
604 pub fn has_argument_allowlist(&self, role: &str, tool: &str, argument: &str) -> bool {
605 if !self.enabled {
606 return false;
607 }
608 let Some(role_cfg) = self.find_role(role) else {
609 return false;
610 };
611 role_cfg.argument_allowlists.iter().any(|al| {
612 (al.tool == tool || glob_match(&al.tool, tool))
613 && al.argument == argument
614 && !al.allowed.is_empty()
615 })
616 }
617
618 fn find_role(&self, name: &str) -> Option<&RoleConfig> {
620 self.roles.iter().find(|r| r.name == name)
621 }
622
623 fn missing_required_argument(
634 &self,
635 role: &str,
636 tool: &str,
637 args: Option<&serde_json::Map<String, serde_json::Value>>,
638 ) -> Option<&str> {
639 if !self.enabled {
640 return None;
641 }
642 let role_cfg = self.find_role(role)?;
643 role_cfg
644 .argument_allowlists
645 .iter()
646 .filter(|al| al.required)
647 .filter(|al| al.tool == tool || glob_match(&al.tool, tool))
651 .find(|al| {
652 !args.is_some_and(|a| {
653 a.get(&al.argument)
654 .is_some_and(serde_json::Value::is_string)
655 })
656 })
657 .map(|al| al.argument.as_str())
658 }
659
660 fn host_matches(patterns: &[String], host: &str) -> bool {
662 patterns.iter().any(|p| glob_match(p, host))
663 }
664
665 #[must_use]
674 pub fn redact_arg(&self, value: &str) -> String {
675 redact_with_salt(self.redaction_salt.expose_secret().as_bytes(), value)
676 }
677}
678
679fn process_redaction_salt() -> &'static SecretString {
682 use base64::{Engine as _, engine::general_purpose::STANDARD_NO_PAD};
683 static PROCESS_SALT: std::sync::OnceLock<SecretString> = std::sync::OnceLock::new();
684 PROCESS_SALT.get_or_init(|| {
685 let mut bytes = [0u8; 32];
686 rand::fill(&mut bytes);
687 SecretString::from(STANDARD_NO_PAD.encode(bytes))
690 })
691}
692
693fn redact_with_salt(salt: &[u8], value: &str) -> String {
698 use std::fmt::Write as _;
699
700 use sha2::Digest as _;
701
702 type HmacSha256 = Hmac<Sha256>;
703 let mut mac = if let Ok(m) = HmacSha256::new_from_slice(salt) {
709 m
710 } else {
711 let digest = Sha256::digest(salt);
712 #[allow(
713 clippy::expect_used,
714 reason = "32-byte SHA-256 digest is unconditionally valid as an HMAC-SHA256 key (RFC 2104 allows any key length); see surrounding comment"
715 )]
716 HmacSha256::new_from_slice(&digest).expect("32-byte SHA256 digest is valid HMAC key")
717 };
718 mac.update(value.as_bytes());
719 let bytes = mac.finalize().into_bytes();
720 let prefix = bytes.get(..4).unwrap_or(&[0; 4]);
722 let mut out = String::with_capacity(8);
723 for b in prefix {
724 let _ = write!(out, "{b:02x}");
725 }
726 out
727}
728
729#[allow(
750 clippy::too_many_lines,
751 reason = "linear request lifecycle (body collect → JSON-RPC parse → policy dispatch) kept inline for security review visibility; helpers already extracted"
752)]
753pub(crate) async fn rbac_middleware(
754 policy: Arc<RbacPolicy>,
755 tool_limiter: Option<Arc<ToolRateLimiter>>,
756 req: Request<Body>,
757 next: Next,
758) -> Response {
759 if req.method() != Method::POST {
761 return next.run(req).await;
762 }
763
764 let peer_ip: Option<IpAddr> = crate::transport::limiter_client_ip(req.extensions());
767
768 let identity = req.extensions().get::<AuthIdentity>();
770 let identity_name = identity.map(|id| id.name.clone()).unwrap_or_default();
771 let role = identity.map(|id| id.role.clone()).unwrap_or_default();
772 let raw_token: SecretString = identity
775 .and_then(|id| id.raw_token.clone())
776 .unwrap_or_else(|| SecretString::from(String::new()));
777 let sub = identity.and_then(|id| id.sub.clone()).unwrap_or_default();
778
779 if policy.is_enabled() && identity.is_none() {
781 return McpxError::Rbac("no authenticated identity".into()).into_response();
782 }
783
784 let (parts, body) = req.into_parts();
786 let bytes = match body.collect().await {
787 Ok(collected) => collected.to_bytes(),
788 Err(e) => {
789 tracing::error!(error = %e, "failed to read request body");
790 return (
791 StatusCode::INTERNAL_SERVER_ERROR,
792 "failed to read request body",
793 )
794 .into_response();
795 }
796 };
797
798 if let Ok(json) = serde_json::from_slice::<serde_json::Value>(&bytes) {
800 let tool_calls = extract_tool_calls(&json);
801 if !tool_calls.is_empty() {
802 for params in tool_calls {
803 if let Some(resp) = enforce_rate_limit(tool_limiter.as_deref(), peer_ip) {
804 #[cfg(feature = "metrics")]
805 crate::metrics::record_rate_limit_deny(&parts.extensions, "tool");
806 return resp;
807 }
808 if policy.is_enabled()
809 && let Some(resp) = enforce_tool_policy(&policy, &identity_name, &role, params)
810 {
811 return resp;
812 }
813 }
814 }
815 }
816 let req = Request::from_parts(parts, Body::from(bytes));
820
821 if role.is_empty() {
823 next.run(req).await
824 } else {
825 CURRENT_ROLE
826 .scope(
827 role,
828 CURRENT_IDENTITY.scope(
829 identity_name,
830 CURRENT_TOKEN.scope(raw_token, CURRENT_SUB.scope(sub, next.run(req))),
831 ),
832 )
833 .await
834 }
835}
836
837fn extract_tool_calls(value: &serde_json::Value) -> Vec<&serde_json::Value> {
843 match value {
844 serde_json::Value::Object(map) => map
845 .get("method")
846 .and_then(serde_json::Value::as_str)
847 .filter(|method| *method == "tools/call")
848 .and_then(|_| map.get("params"))
849 .into_iter()
850 .collect(),
851 serde_json::Value::Array(items) => items
852 .iter()
853 .filter_map(|item| match item {
854 serde_json::Value::Object(map) => map
855 .get("method")
856 .and_then(serde_json::Value::as_str)
857 .filter(|method| *method == "tools/call")
858 .and_then(|_| map.get("params")),
859 serde_json::Value::Null
860 | serde_json::Value::Bool(_)
861 | serde_json::Value::Number(_)
862 | serde_json::Value::String(_)
863 | serde_json::Value::Array(_) => None,
864 })
865 .collect(),
866 serde_json::Value::Null
867 | serde_json::Value::Bool(_)
868 | serde_json::Value::Number(_)
869 | serde_json::Value::String(_) => Vec::new(),
870 }
871}
872
873fn enforce_rate_limit(
876 tool_limiter: Option<&ToolRateLimiter>,
877 peer_ip: Option<IpAddr>,
878) -> Option<Response> {
879 let limiter = tool_limiter?;
880 let ip = peer_ip?;
881 if let Err(wait) = limiter.check_key_wait(&ip) {
882 tracing::warn!(%ip, "tool invocation rate limited");
883 return Some(
884 McpxError::RateLimitedFor {
885 message: "too many tool invocations".into(),
886 retry_after: wait,
887 }
888 .into_response(),
889 );
890 }
891 None
892}
893
894fn enforce_tool_policy(
903 policy: &RbacPolicy,
904 identity_name: &str,
905 role: &str,
906 params: &serde_json::Value,
907) -> Option<Response> {
908 let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
909 let host_value = params.get("arguments").and_then(|a| a.get("host"));
910
911 if let Some(value) = host_value
919 && !value.is_string()
920 {
921 tracing::warn!(
922 user = %identity_name,
923 role = %role,
924 tool = tool_name,
925 value_type = json_value_type(value),
926 "non-string host argument rejected"
927 );
928 return Some(
929 McpxError::Rbac(format!(
930 "argument 'host' must be a string for tool '{tool_name}'"
931 ))
932 .into_response(),
933 );
934 }
935 let host = host_value.and_then(|h| h.as_str());
938
939 let decision = if let Some(host) = host {
940 policy.check(role, tool_name, host)
941 } else {
942 policy.check_operation(role, tool_name)
943 };
944 if decision == RbacDecision::Deny {
945 tracing::warn!(
946 user = %identity_name,
947 role = %role,
948 tool = tool_name,
949 host = host.unwrap_or("-"),
950 "RBAC denied"
951 );
952 return Some(
953 McpxError::Rbac(format!("{tool_name} denied for role '{role}'")).into_response(),
954 );
955 }
956
957 let args = params.get("arguments").and_then(|a| a.as_object());
958 if let Some(args) = args {
959 for (arg_key, arg_val) in args {
960 if let Some(resp) =
961 check_argument(policy, identity_name, role, tool_name, arg_key, arg_val)
962 {
963 return Some(resp);
964 }
965 }
966 }
967 check_required_arguments(policy, identity_name, role, tool_name, args)
968}
969
970fn check_required_arguments(
978 policy: &RbacPolicy,
979 identity_name: &str,
980 role: &str,
981 tool_name: &str,
982 args: Option<&serde_json::Map<String, serde_json::Value>>,
983) -> Option<Response> {
984 let missing = policy.missing_required_argument(role, tool_name, args)?;
985 tracing::warn!(
986 user = %identity_name,
987 role = %role,
988 tool = tool_name,
989 argument = missing,
990 "required argument missing"
991 );
992 Some(
993 McpxError::Rbac(format!(
994 "argument '{missing}' is required for tool '{tool_name}'"
995 ))
996 .into_response(),
997 )
998}
999
1000fn check_argument(
1001 policy: &RbacPolicy,
1002 identity_name: &str,
1003 role: &str,
1004 tool_name: &str,
1005 arg_key: &str,
1006 arg_val: &serde_json::Value,
1007) -> Option<Response> {
1008 if !policy.has_argument_allowlist(role, tool_name, arg_key) {
1009 return None;
1010 }
1011 let Some(val_str) = arg_val.as_str() else {
1012 tracing::warn!(
1018 user = %identity_name,
1019 role = %role,
1020 tool = tool_name,
1021 argument = arg_key,
1022 value_type = json_value_type(arg_val),
1023 "non-string argument rejected by allowlist"
1024 );
1025 return Some(
1026 McpxError::Rbac(format!(
1027 "argument '{arg_key}' must be a string for tool '{tool_name}'"
1028 ))
1029 .into_response(),
1030 );
1031 };
1032 if policy.argument_allowed(role, tool_name, arg_key, val_str) {
1033 return None;
1034 }
1035 tracing::warn!(
1040 user = %identity_name,
1041 role = %role,
1042 tool = tool_name,
1043 argument = arg_key,
1044 arg_hmac = %policy.redact_arg(val_str),
1045 "argument not in allowlist"
1046 );
1047 Some(
1048 McpxError::Rbac(format!(
1049 "argument '{arg_key}' value not in allowlist for tool '{tool_name}'"
1050 ))
1051 .into_response(),
1052 )
1053}
1054
1055fn json_value_type(v: &serde_json::Value) -> &'static str {
1056 match v {
1057 serde_json::Value::Null => "null",
1058 serde_json::Value::Bool(_) => "bool",
1059 serde_json::Value::Number(_) => "number",
1060 serde_json::Value::String(_) => "string",
1061 serde_json::Value::Array(_) => "array",
1062 serde_json::Value::Object(_) => "object",
1063 }
1064}
1065
1066fn glob_match(pattern: &str, text: &str) -> bool {
1076 let parts: Vec<&str> = pattern.split('*').collect();
1077 if parts.len() == 1 {
1078 return pattern == text;
1080 }
1081
1082 let pos = if let Some(&first) = parts.first()
1084 && !first.is_empty()
1085 {
1086 if !text.starts_with(first) {
1087 return false;
1088 }
1089 first.len()
1090 } else {
1091 0
1092 };
1093
1094 if let Some(&last) = parts.last()
1096 && !last.is_empty()
1097 {
1098 if !text.get(pos..).unwrap_or_default().ends_with(last) {
1099 return false;
1100 }
1101 let end = text.len() - last.len();
1103 if pos > end {
1104 return false;
1105 }
1106 let middle = text.get(pos..end).unwrap_or_default();
1108 let middle_parts = parts.get(1..parts.len() - 1).unwrap_or_default();
1109 return match_middle(middle, middle_parts);
1110 }
1111
1112 let middle = text.get(pos..).unwrap_or_default();
1114 let middle_parts = parts.get(1..parts.len() - 1).unwrap_or_default();
1115 match_middle(middle, middle_parts)
1116}
1117
1118fn match_middle(mut text: &str, parts: &[&str]) -> bool {
1120 for part in parts {
1121 if part.is_empty() {
1122 continue;
1123 }
1124 if let Some(idx) = text.find(part) {
1125 text = text.get(idx + part.len()..).unwrap_or_default();
1126 } else {
1127 return false;
1128 }
1129 }
1130 true
1131}
1132
1133#[cfg(test)]
1134mod tests {
1135 use super::*;
1136
1137 #[test]
1142 fn tool_limiter_burst_allows_initial_spike() {
1143 let limiter = build_tool_rate_limiter(2, Some(4));
1144 let ip: IpAddr = "10.9.9.9".parse().unwrap();
1145 for i in 0..4 {
1146 assert!(
1147 limiter.check_key(&ip).is_ok(),
1148 "burst request {i} should pass"
1149 );
1150 }
1151 assert!(
1152 limiter.check_key(&ip).is_err(),
1153 "request 5 must exceed the burst bucket"
1154 );
1155 }
1156
1157 #[test]
1159 fn tool_limiter_deny_sets_retry_after() {
1160 let limiter = build_tool_rate_limiter(1, None);
1161 let ip: IpAddr = "10.8.8.8".parse().unwrap();
1162 assert!(enforce_rate_limit(Some(&limiter), Some(ip)).is_none());
1163 let resp = enforce_rate_limit(Some(&limiter), Some(ip))
1164 .expect("second call within the window must deny");
1165 assert_eq!(resp.status(), axum::http::StatusCode::TOO_MANY_REQUESTS);
1166 let retry_after = resp
1167 .headers()
1168 .get(axum::http::header::RETRY_AFTER)
1169 .expect("Retry-After present")
1170 .to_str()
1171 .unwrap()
1172 .parse::<u64>()
1173 .unwrap();
1174 assert!(retry_after >= 1, "delta-seconds must be >= 1");
1175 }
1176
1177 fn test_policy() -> RbacPolicy {
1178 RbacPolicy::new(&RbacConfig {
1179 enabled: true,
1180 roles: vec![
1181 RoleConfig {
1182 name: "viewer".into(),
1183 description: Some("Read-only".into()),
1184 allow: vec![
1185 "list_hosts".into(),
1186 "resource_list".into(),
1187 "resource_inspect".into(),
1188 "resource_logs".into(),
1189 "system_info".into(),
1190 ],
1191 deny: vec![],
1192 hosts: vec!["*".into()],
1193 argument_allowlists: vec![],
1194 },
1195 RoleConfig {
1196 name: "deploy".into(),
1197 description: Some("Lifecycle management".into()),
1198 allow: vec![
1199 "list_hosts".into(),
1200 "resource_list".into(),
1201 "resource_run".into(),
1202 "resource_start".into(),
1203 "resource_stop".into(),
1204 "resource_restart".into(),
1205 "resource_logs".into(),
1206 "image_pull".into(),
1207 ],
1208 deny: vec!["resource_delete".into(), "resource_exec".into()],
1209 hosts: vec!["web-*".into(), "api-*".into()],
1210 argument_allowlists: vec![],
1211 },
1212 RoleConfig {
1213 name: "ops".into(),
1214 description: Some("Full access".into()),
1215 allow: vec!["*".into()],
1216 deny: vec![],
1217 hosts: vec!["*".into()],
1218 argument_allowlists: vec![],
1219 },
1220 RoleConfig {
1221 name: "restricted-exec".into(),
1222 description: Some("Exec with argument allowlist".into()),
1223 allow: vec!["resource_exec".into()],
1224 deny: vec![],
1225 hosts: vec!["dev-*".into()],
1226 argument_allowlists: vec![ArgumentAllowlist {
1227 tool: "resource_exec".into(),
1228 argument: "cmd".into(),
1229 allowed: vec![
1230 "sh".into(),
1231 "bash".into(),
1232 "cat".into(),
1233 "ls".into(),
1234 "ps".into(),
1235 ],
1236 required: false,
1237 }],
1238 },
1239 ],
1240 redaction_salt: None,
1241 })
1242 }
1243
1244 #[test]
1247 fn glob_exact_match() {
1248 assert!(glob_match("web-prod-1", "web-prod-1"));
1249 assert!(!glob_match("web-prod-1", "web-prod-2"));
1250 }
1251
1252 #[test]
1253 fn glob_star_suffix() {
1254 assert!(glob_match("web-*", "web-prod-1"));
1255 assert!(glob_match("web-*", "web-staging"));
1256 assert!(!glob_match("web-*", "api-prod"));
1257 }
1258
1259 #[test]
1260 fn glob_star_prefix() {
1261 assert!(glob_match("*-prod", "web-prod"));
1262 assert!(glob_match("*-prod", "api-prod"));
1263 assert!(!glob_match("*-prod", "web-staging"));
1264 }
1265
1266 #[test]
1267 fn glob_star_middle() {
1268 assert!(glob_match("web-*-prod", "web-us-prod"));
1269 assert!(glob_match("web-*-prod", "web-eu-east-prod"));
1270 assert!(!glob_match("web-*-prod", "web-staging"));
1271 }
1272
1273 #[test]
1274 fn glob_star_only() {
1275 assert!(glob_match("*", "anything"));
1276 assert!(glob_match("*", ""));
1277 }
1278
1279 #[test]
1280 fn glob_multiple_stars() {
1281 assert!(glob_match("*web*prod*", "my-web-us-prod-1"));
1282 assert!(!glob_match("*web*prod*", "my-api-us-staging"));
1283 }
1284
1285 #[test]
1290 fn glob_match_multibyte_utf8() {
1291 assert!(glob_match("hé*llo", "héllo"));
1292 assert!(glob_match("*ö*", "wörld"));
1293 assert!(glob_match("über*", "übermensch"));
1294 assert!(glob_match("*界", "世界"));
1295 assert!(!glob_match("hé*llo", "hello"));
1296 assert!(!glob_match("界*", "世界"));
1297 assert!(glob_match("世*界", "世界"));
1298 }
1299
1300 #[test]
1312 fn glob_prefix_and_suffix_meet_exactly() {
1313 assert!(glob_match("ab*cd", "abcd"));
1316 }
1317
1318 #[test]
1323 fn glob_middle_segment_required_with_suffix() {
1324 assert!(!glob_match("a*b*c", "axyc"));
1329 }
1330
1331 #[test]
1337 fn glob_match_middle_advances_past_matched_part() {
1338 assert!(!glob_match("*ab*ab*", "xxab_yz"));
1343 }
1344
1345 #[test]
1350 fn glob_match_middle_uses_addition_not_multiplication() {
1351 assert!(glob_match("*abcde*X*", "yyyyyyyyabcde_X"));
1355 }
1356
1357 #[test]
1366 fn argument_allowed_glob_pattern_with_literal_mismatch_still_enforced() {
1367 let role = RoleConfig::new("viewer", vec!["run-foo".into()], vec!["*".into()])
1375 .with_argument_allowlists(vec![ArgumentAllowlist::new(
1376 "run-*",
1377 "cmd",
1378 vec!["ls".into()],
1379 )]);
1380 let mut config = RbacConfig::with_roles(vec![role]);
1381 config.enabled = true;
1382 let policy = RbacPolicy::new(&config);
1383 assert!(!policy.argument_allowed("viewer", "run-foo", "cmd", "rm"));
1384 }
1385
1386 #[test]
1389 fn disabled_policy_allows_everything() {
1390 let policy = RbacPolicy::new(&RbacConfig {
1391 enabled: false,
1392 roles: vec![],
1393 redaction_salt: None,
1394 });
1395 assert_eq!(
1396 policy.check("nonexistent", "resource_delete", "any-host"),
1397 RbacDecision::Allow
1398 );
1399 }
1400
1401 #[test]
1402 fn unknown_role_denied() {
1403 let policy = test_policy();
1404 assert_eq!(
1405 policy.check("unknown", "resource_list", "web-prod-1"),
1406 RbacDecision::Deny
1407 );
1408 }
1409
1410 #[test]
1411 fn viewer_allowed_read_ops() {
1412 let policy = test_policy();
1413 assert_eq!(
1414 policy.check("viewer", "resource_list", "web-prod-1"),
1415 RbacDecision::Allow
1416 );
1417 assert_eq!(
1418 policy.check("viewer", "system_info", "db-host"),
1419 RbacDecision::Allow
1420 );
1421 }
1422
1423 #[test]
1424 fn viewer_denied_write_ops() {
1425 let policy = test_policy();
1426 assert_eq!(
1427 policy.check("viewer", "resource_run", "web-prod-1"),
1428 RbacDecision::Deny
1429 );
1430 assert_eq!(
1431 policy.check("viewer", "resource_delete", "web-prod-1"),
1432 RbacDecision::Deny
1433 );
1434 }
1435
1436 #[test]
1437 fn deploy_allowed_on_matching_hosts() {
1438 let policy = test_policy();
1439 assert_eq!(
1440 policy.check("deploy", "resource_run", "web-prod-1"),
1441 RbacDecision::Allow
1442 );
1443 assert_eq!(
1444 policy.check("deploy", "resource_start", "api-staging"),
1445 RbacDecision::Allow
1446 );
1447 }
1448
1449 #[test]
1450 fn deploy_denied_on_non_matching_host() {
1451 let policy = test_policy();
1452 assert_eq!(
1453 policy.check("deploy", "resource_run", "db-prod-1"),
1454 RbacDecision::Deny
1455 );
1456 }
1457
1458 #[test]
1459 fn deny_overrides_allow() {
1460 let policy = test_policy();
1461 assert_eq!(
1462 policy.check("deploy", "resource_delete", "web-prod-1"),
1463 RbacDecision::Deny
1464 );
1465 assert_eq!(
1466 policy.check("deploy", "resource_exec", "web-prod-1"),
1467 RbacDecision::Deny
1468 );
1469 }
1470
1471 #[test]
1472 fn ops_wildcard_allows_everything() {
1473 let policy = test_policy();
1474 assert_eq!(
1475 policy.check("ops", "resource_delete", "any-host"),
1476 RbacDecision::Allow
1477 );
1478 assert_eq!(
1479 policy.check("ops", "secret_create", "db-host"),
1480 RbacDecision::Allow
1481 );
1482 }
1483
1484 #[test]
1487 fn host_visible_respects_globs() {
1488 let policy = test_policy();
1489 assert!(policy.host_visible("deploy", "web-prod-1"));
1490 assert!(policy.host_visible("deploy", "api-staging"));
1491 assert!(!policy.host_visible("deploy", "db-prod-1"));
1492 assert!(policy.host_visible("ops", "anything"));
1493 assert!(policy.host_visible("viewer", "anything"));
1494 }
1495
1496 #[test]
1497 fn host_visible_unknown_role() {
1498 let policy = test_policy();
1499 assert!(!policy.host_visible("unknown", "web-prod-1"));
1500 }
1501
1502 #[test]
1505 fn argument_allowed_no_allowlist() {
1506 let policy = test_policy();
1507 assert!(policy.argument_allowed("ops", "resource_exec", "cmd", "rm -rf /"));
1509 assert!(policy.argument_allowed("ops", "resource_exec", "cmd", "bash"));
1510 }
1511
1512 #[test]
1513 fn argument_allowed_with_allowlist() {
1514 let policy = test_policy();
1515 assert!(policy.argument_allowed("restricted-exec", "resource_exec", "cmd", "sh"));
1516 assert!(policy.argument_allowed(
1517 "restricted-exec",
1518 "resource_exec",
1519 "cmd",
1520 "bash -c 'echo hi'"
1521 ));
1522 assert!(policy.argument_allowed(
1523 "restricted-exec",
1524 "resource_exec",
1525 "cmd",
1526 "cat /etc/hosts"
1527 ));
1528 assert!(policy.argument_allowed(
1529 "restricted-exec",
1530 "resource_exec",
1531 "cmd",
1532 "/usr/bin/ls -la"
1533 ));
1534 }
1535
1536 #[test]
1537 fn argument_denied_not_in_allowlist() {
1538 let policy = test_policy();
1539 assert!(!policy.argument_allowed("restricted-exec", "resource_exec", "cmd", "rm -rf /"));
1540 assert!(!policy.argument_allowed(
1541 "restricted-exec",
1542 "resource_exec",
1543 "cmd",
1544 "python3 exploit.py"
1545 ));
1546 assert!(!policy.argument_allowed(
1547 "restricted-exec",
1548 "resource_exec",
1549 "cmd",
1550 "/usr/bin/curl evil.com"
1551 ));
1552 }
1553
1554 #[test]
1555 fn argument_denied_unknown_role() {
1556 let policy = test_policy();
1557 assert!(!policy.argument_allowed("unknown", "resource_exec", "cmd", "sh"));
1558 }
1559
1560 fn shlex_policy(allowed: Vec<String>) -> RbacPolicy {
1569 let role = RoleConfig::new("viewer", vec!["run".into()], vec!["*".into()])
1570 .with_argument_allowlists(vec![ArgumentAllowlist::new("run", "cmd", allowed)]);
1571 let mut config = RbacConfig::with_roles(vec![role]);
1572 config.enabled = true;
1573 RbacPolicy::new(&config)
1574 }
1575
1576 #[test]
1577 fn argument_allowed_matches_quoted_path_with_spaces() {
1578 let policy = shlex_policy(vec!["/usr/bin/my tool".into()]);
1579 assert!(policy.argument_allowed("viewer", "run", "cmd", r#""/usr/bin/my tool" --flag"#));
1580 }
1581
1582 #[test]
1583 fn argument_allowed_matches_basename_of_quoted_path() {
1584 let policy = shlex_policy(vec!["my tool".into()]);
1585 assert!(policy.argument_allowed("viewer", "run", "cmd", r#""/usr/bin/my tool" --flag"#));
1586 }
1587
1588 #[test]
1589 fn argument_allowed_fails_closed_on_unbalanced_quote() {
1590 let policy = shlex_policy(vec!["unbalanced".into()]);
1591 assert!(!policy.argument_allowed("viewer", "run", "cmd", r"unbalanced 'quote"));
1592 }
1593
1594 #[test]
1595 fn argument_allowed_fails_closed_on_empty_string() {
1596 let policy = shlex_policy(vec![String::new()]);
1597 assert!(!policy.argument_allowed("viewer", "run", "cmd", ""));
1598 }
1599
1600 #[test]
1601 fn argument_allowed_handles_single_quoted_executable() {
1602 let policy = shlex_policy(vec!["/bin/sh".into()]);
1603 assert!(policy.argument_allowed("viewer", "run", "cmd", r"'/bin/sh' -c 'echo hi'"));
1604 }
1605
1606 #[test]
1607 fn argument_allowed_handles_tab_separator() {
1608 let policy = shlex_policy(vec!["ls".into()]);
1609 assert!(policy.argument_allowed("viewer", "run", "cmd", "ls\t/etc/passwd"));
1610 }
1611
1612 #[test]
1613 fn argument_allowed_plain_token_unchanged() {
1614 let policy = shlex_policy(vec!["ls".into()]);
1615 assert!(policy.argument_allowed("viewer", "run", "cmd", "ls"));
1616 }
1617
1618 #[test]
1624 fn argument_allowed_fails_closed_on_quoted_empty_first_token() {
1625 let policy = shlex_policy(vec![String::new()]);
1629 assert!(!policy.argument_allowed("viewer", "run", "cmd", r#""""#));
1630 }
1631
1632 #[test]
1633 fn argument_allowed_quoted_literal_token_no_longer_matches() {
1634 let policy = shlex_policy(vec!["'bash'".into()]);
1640 assert!(!policy.argument_allowed("viewer", "run", "cmd", "'bash' -c true"));
1641 }
1642
1643 #[test]
1644 fn argument_allowed_backslash_literal_token_no_longer_matches() {
1645 let policy = shlex_policy(vec![r"foo\bar".into()]);
1650 assert!(!policy.argument_allowed("viewer", "run", "cmd", r"foo\bar --x"));
1651 }
1652
1653 #[test]
1654 fn argument_allowed_windows_path_no_longer_matches() {
1655 let policy = shlex_policy(vec![r"C:\Windows\System32\cmd.exe".into()]);
1660 assert!(!policy.argument_allowed(
1661 "viewer",
1662 "run",
1663 "cmd",
1664 r"C:\Windows\System32\cmd.exe /c dir"
1665 ));
1666 }
1667
1668 #[test]
1671 fn host_patterns_returns_globs() {
1672 let policy = test_policy();
1673 assert_eq!(
1674 policy.host_patterns("deploy"),
1675 Some(vec!["web-*".to_owned(), "api-*".to_owned()].as_slice())
1676 );
1677 assert_eq!(
1678 policy.host_patterns("ops"),
1679 Some(vec!["*".to_owned()].as_slice())
1680 );
1681 assert!(policy.host_patterns("nonexistent").is_none());
1682 }
1683
1684 #[test]
1687 fn check_operation_allows_without_host() {
1688 let policy = test_policy();
1689 assert_eq!(
1690 policy.check_operation("deploy", "resource_run"),
1691 RbacDecision::Allow
1692 );
1693 assert_eq!(
1695 policy.check("deploy", "resource_run", "db-prod-1"),
1696 RbacDecision::Deny
1697 );
1698 }
1699
1700 #[test]
1701 fn check_operation_deny_overrides() {
1702 let policy = test_policy();
1703 assert_eq!(
1704 policy.check_operation("deploy", "resource_delete"),
1705 RbacDecision::Deny
1706 );
1707 }
1708
1709 #[test]
1710 fn check_operation_unknown_role() {
1711 let policy = test_policy();
1712 assert_eq!(
1713 policy.check_operation("unknown", "resource_list"),
1714 RbacDecision::Deny
1715 );
1716 }
1717
1718 #[test]
1719 fn check_operation_disabled() {
1720 let policy = RbacPolicy::new(&RbacConfig {
1721 enabled: false,
1722 roles: vec![],
1723 redaction_salt: None,
1724 });
1725 assert_eq!(
1726 policy.check_operation("nonexistent", "anything"),
1727 RbacDecision::Allow
1728 );
1729 }
1730
1731 #[test]
1734 fn current_role_returns_none_outside_scope() {
1735 assert!(current_role().is_none());
1736 }
1737
1738 #[test]
1739 fn current_identity_returns_none_outside_scope() {
1740 assert!(current_identity().is_none());
1741 }
1742
1743 use axum::{
1746 body::Body,
1747 http::{Method, Request, StatusCode},
1748 };
1749 use tower::ServiceExt as _;
1750
1751 fn tool_call_body(tool: &str, args: &serde_json::Value) -> String {
1752 serde_json::json!({
1753 "jsonrpc": "2.0",
1754 "id": 1,
1755 "method": "tools/call",
1756 "params": {
1757 "name": tool,
1758 "arguments": args
1759 }
1760 })
1761 .to_string()
1762 }
1763
1764 fn rbac_router(policy: Arc<RbacPolicy>) -> axum::Router {
1765 axum::Router::new()
1766 .route("/mcp", axum::routing::post(|| async { "ok" }))
1767 .layer(axum::middleware::from_fn(move |req, next| {
1768 let p = Arc::clone(&policy);
1769 rbac_middleware(p, None, req, next)
1770 }))
1771 }
1772
1773 fn rbac_router_with_identity(policy: Arc<RbacPolicy>, identity: AuthIdentity) -> axum::Router {
1774 axum::Router::new()
1775 .route("/mcp", axum::routing::post(|| async { "ok" }))
1776 .layer(axum::middleware::from_fn(
1777 move |mut req: Request<Body>, next: Next| {
1778 let p = Arc::clone(&policy);
1779 let id = identity.clone();
1780 async move {
1781 req.extensions_mut().insert(id);
1782 rbac_middleware(p, None, req, next).await
1783 }
1784 },
1785 ))
1786 }
1787
1788 #[cfg(feature = "metrics")]
1792 #[tokio::test]
1793 async fn tool_limiter_deny_increments_counter() {
1794 use axum::extract::ConnectInfo;
1795
1796 let policy = Arc::new(test_policy());
1797 let limiter = build_tool_rate_limiter(1, None);
1798 let metrics = Arc::new(crate::metrics::McpMetrics::new().unwrap());
1799 let identity = AuthIdentity {
1800 method: crate::auth::AuthMethod::BearerToken,
1801 name: "alice".into(),
1802 role: "viewer".into(),
1803 raw_token: None,
1804 sub: None,
1805 };
1806 let app = {
1807 let metrics = Arc::clone(&metrics);
1808 axum::Router::new()
1809 .route("/mcp", axum::routing::post(|| async { "ok" }))
1810 .layer(axum::middleware::from_fn(
1811 move |mut req: Request<Body>, next: Next| {
1812 let p = Arc::clone(&policy);
1813 let l = Arc::clone(&limiter);
1814 let id = identity.clone();
1815 let m = Arc::clone(&metrics);
1816 async move {
1817 req.extensions_mut().insert(id);
1818 req.extensions_mut().insert(m);
1819 let peer: std::net::SocketAddr =
1820 "10.9.9.1:40000".parse().expect("static socket addr parses");
1821 req.extensions_mut().insert(ConnectInfo(peer));
1822 rbac_middleware(p, Some(l), req, next).await
1823 }
1824 },
1825 ))
1826 };
1827 let mk = || {
1828 Request::builder()
1829 .method(Method::POST)
1830 .uri("/mcp")
1831 .header("content-type", "application/json")
1832 .body(Body::from(tool_call_body(
1833 "resource_list",
1834 &serde_json::json!({}),
1835 )))
1836 .unwrap()
1837 };
1838 let counter = || {
1839 metrics
1840 .rate_limited_total
1841 .with_label_values(&["tool"])
1842 .get()
1843 };
1844
1845 let first = app.clone().oneshot(mk()).await.unwrap();
1846 assert_eq!(first.status(), StatusCode::OK);
1847 assert_eq!(counter(), 0, "successful call must not count");
1848
1849 let denied = app.clone().oneshot(mk()).await.unwrap();
1850 assert_eq!(denied.status(), StatusCode::TOO_MANY_REQUESTS);
1851 assert_eq!(counter(), 1, "deny must increment the tool label");
1852 }
1853
1854 #[tokio::test]
1855 async fn middleware_passes_non_post() {
1856 let policy = Arc::new(test_policy());
1857 let app = rbac_router(policy);
1858 let req = Request::builder()
1860 .method(Method::GET)
1861 .uri("/mcp")
1862 .body(Body::empty())
1863 .unwrap();
1864 let resp = app.oneshot(req).await.unwrap();
1867 assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
1868 }
1869
1870 #[tokio::test]
1871 async fn middleware_denies_without_identity() {
1872 let policy = Arc::new(test_policy());
1873 let app = rbac_router(policy);
1874 let body = tool_call_body("resource_list", &serde_json::json!({}));
1875 let req = Request::builder()
1876 .method(Method::POST)
1877 .uri("/mcp")
1878 .header("content-type", "application/json")
1879 .body(Body::from(body))
1880 .unwrap();
1881 let resp = app.oneshot(req).await.unwrap();
1882 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1883 }
1884
1885 #[tokio::test]
1886 async fn middleware_allows_permitted_tool() {
1887 let policy = Arc::new(test_policy());
1888 let id = AuthIdentity {
1889 method: crate::auth::AuthMethod::BearerToken,
1890 name: "alice".into(),
1891 role: "viewer".into(),
1892 raw_token: None,
1893 sub: None,
1894 };
1895 let app = rbac_router_with_identity(policy, id);
1896 let body = tool_call_body("resource_list", &serde_json::json!({}));
1897 let req = Request::builder()
1898 .method(Method::POST)
1899 .uri("/mcp")
1900 .header("content-type", "application/json")
1901 .body(Body::from(body))
1902 .unwrap();
1903 let resp = app.oneshot(req).await.unwrap();
1904 assert_eq!(resp.status(), StatusCode::OK);
1905 }
1906
1907 #[tokio::test]
1908 async fn middleware_denies_unpermitted_tool() {
1909 let policy = Arc::new(test_policy());
1910 let id = AuthIdentity {
1911 method: crate::auth::AuthMethod::BearerToken,
1912 name: "alice".into(),
1913 role: "viewer".into(),
1914 raw_token: None,
1915 sub: None,
1916 };
1917 let app = rbac_router_with_identity(policy, id);
1918 let body = tool_call_body("resource_delete", &serde_json::json!({}));
1919 let req = Request::builder()
1920 .method(Method::POST)
1921 .uri("/mcp")
1922 .header("content-type", "application/json")
1923 .body(Body::from(body))
1924 .unwrap();
1925 let resp = app.oneshot(req).await.unwrap();
1926 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
1927 }
1928
1929 #[tokio::test]
1930 async fn middleware_passes_non_tool_call_post() {
1931 let policy = Arc::new(test_policy());
1932 let id = AuthIdentity {
1933 method: crate::auth::AuthMethod::BearerToken,
1934 name: "alice".into(),
1935 role: "viewer".into(),
1936 raw_token: None,
1937 sub: None,
1938 };
1939 let app = rbac_router_with_identity(policy, id);
1940 let body = serde_json::json!({
1942 "jsonrpc": "2.0",
1943 "id": 1,
1944 "method": "resources/list"
1945 })
1946 .to_string();
1947 let req = Request::builder()
1948 .method(Method::POST)
1949 .uri("/mcp")
1950 .header("content-type", "application/json")
1951 .body(Body::from(body))
1952 .unwrap();
1953 let resp = app.oneshot(req).await.unwrap();
1954 assert_eq!(resp.status(), StatusCode::OK);
1955 }
1956
1957 #[tokio::test]
1958 async fn middleware_enforces_argument_allowlist() {
1959 let policy = Arc::new(test_policy());
1960 let id = AuthIdentity {
1961 method: crate::auth::AuthMethod::BearerToken,
1962 name: "dev".into(),
1963 role: "restricted-exec".into(),
1964 raw_token: None,
1965 sub: None,
1966 };
1967 let app = rbac_router_with_identity(Arc::clone(&policy), id.clone());
1969 let body = tool_call_body(
1970 "resource_exec",
1971 &serde_json::json!({"cmd": "ls -la", "host": "dev-1"}),
1972 );
1973 let req = Request::builder()
1974 .method(Method::POST)
1975 .uri("/mcp")
1976 .body(Body::from(body))
1977 .unwrap();
1978 let resp = app.oneshot(req).await.unwrap();
1979 assert_eq!(resp.status(), StatusCode::OK);
1980
1981 let app = rbac_router_with_identity(policy, id);
1983 let body = tool_call_body(
1984 "resource_exec",
1985 &serde_json::json!({"cmd": "rm -rf /", "host": "dev-1"}),
1986 );
1987 let req = Request::builder()
1988 .method(Method::POST)
1989 .uri("/mcp")
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 #[tokio::test]
1997 async fn middleware_disabled_policy_passes_everything() {
1998 let policy = Arc::new(RbacPolicy::disabled());
1999 let app = rbac_router(policy);
2000 let body = tool_call_body("anything", &serde_json::json!({}));
2002 let req = Request::builder()
2003 .method(Method::POST)
2004 .uri("/mcp")
2005 .body(Body::from(body))
2006 .unwrap();
2007 let resp = app.oneshot(req).await.unwrap();
2008 assert_eq!(resp.status(), StatusCode::OK);
2009 }
2010
2011 #[tokio::test]
2012 async fn middleware_batch_all_allowed_passes() {
2013 let policy = Arc::new(test_policy());
2014 let id = AuthIdentity {
2015 method: crate::auth::AuthMethod::BearerToken,
2016 name: "alice".into(),
2017 role: "viewer".into(),
2018 raw_token: None,
2019 sub: None,
2020 };
2021 let app = rbac_router_with_identity(policy, id);
2022 let body = serde_json::json!([
2023 {
2024 "jsonrpc": "2.0",
2025 "id": 1,
2026 "method": "tools/call",
2027 "params": { "name": "resource_list", "arguments": {} }
2028 },
2029 {
2030 "jsonrpc": "2.0",
2031 "id": 2,
2032 "method": "tools/call",
2033 "params": { "name": "system_info", "arguments": {} }
2034 }
2035 ])
2036 .to_string();
2037 let req = Request::builder()
2038 .method(Method::POST)
2039 .uri("/mcp")
2040 .header("content-type", "application/json")
2041 .body(Body::from(body))
2042 .unwrap();
2043 let resp = app.oneshot(req).await.unwrap();
2044 assert_eq!(resp.status(), StatusCode::OK);
2045 }
2046
2047 #[tokio::test]
2048 async fn middleware_batch_with_denied_call_rejects_entire_batch() {
2049 let policy = Arc::new(test_policy());
2050 let id = AuthIdentity {
2051 method: crate::auth::AuthMethod::BearerToken,
2052 name: "alice".into(),
2053 role: "viewer".into(),
2054 raw_token: None,
2055 sub: None,
2056 };
2057 let app = rbac_router_with_identity(policy, id);
2058 let body = serde_json::json!([
2059 {
2060 "jsonrpc": "2.0",
2061 "id": 1,
2062 "method": "tools/call",
2063 "params": { "name": "resource_list", "arguments": {} }
2064 },
2065 {
2066 "jsonrpc": "2.0",
2067 "id": 2,
2068 "method": "tools/call",
2069 "params": { "name": "resource_delete", "arguments": {} }
2070 }
2071 ])
2072 .to_string();
2073 let req = Request::builder()
2074 .method(Method::POST)
2075 .uri("/mcp")
2076 .header("content-type", "application/json")
2077 .body(Body::from(body))
2078 .unwrap();
2079 let resp = app.oneshot(req).await.unwrap();
2080 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2081 }
2082
2083 #[tokio::test]
2084 async fn middleware_batch_mixed_allowed_and_denied_rejects() {
2085 let policy = Arc::new(test_policy());
2086 let id = AuthIdentity {
2087 method: crate::auth::AuthMethod::BearerToken,
2088 name: "dev".into(),
2089 role: "restricted-exec".into(),
2090 raw_token: None,
2091 sub: None,
2092 };
2093 let app = rbac_router_with_identity(policy, id);
2094 let body = serde_json::json!([
2095 {
2096 "jsonrpc": "2.0",
2097 "id": 1,
2098 "method": "tools/call",
2099 "params": {
2100 "name": "resource_exec",
2101 "arguments": { "cmd": "ls -la", "host": "dev-1" }
2102 }
2103 },
2104 {
2105 "jsonrpc": "2.0",
2106 "id": 2,
2107 "method": "tools/call",
2108 "params": {
2109 "name": "resource_exec",
2110 "arguments": { "cmd": "rm -rf /", "host": "dev-1" }
2111 }
2112 }
2113 ])
2114 .to_string();
2115 let req = Request::builder()
2116 .method(Method::POST)
2117 .uri("/mcp")
2118 .header("content-type", "application/json")
2119 .body(Body::from(body))
2120 .unwrap();
2121 let resp = app.oneshot(req).await.unwrap();
2122 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2123 }
2124
2125 #[test]
2128 fn redact_with_salt_is_deterministic_per_salt() {
2129 let salt = b"unit-test-salt";
2130 let a = redact_with_salt(salt, "rm -rf /");
2131 let b = redact_with_salt(salt, "rm -rf /");
2132 assert_eq!(a, b, "same input + salt must yield identical hash");
2133 assert_eq!(a.len(), 8, "redacted hash is 8 hex chars (4 bytes)");
2134 assert!(
2135 a.chars().all(|c| c.is_ascii_hexdigit()),
2136 "redacted hash must be lowercase hex: {a}"
2137 );
2138 }
2139
2140 #[test]
2141 fn redact_with_salt_differs_across_salts() {
2142 let v = "the-same-value";
2143 let h1 = redact_with_salt(b"salt-one", v);
2144 let h2 = redact_with_salt(b"salt-two", v);
2145 assert_ne!(
2146 h1, h2,
2147 "different salts must produce different hashes for the same value"
2148 );
2149 }
2150
2151 #[test]
2152 fn redact_with_salt_distinguishes_values() {
2153 let salt = b"k";
2154 let h1 = redact_with_salt(salt, "alpha");
2155 let h2 = redact_with_salt(salt, "beta");
2156 assert_ne!(h1, h2, "different values must produce different hashes");
2158 }
2159
2160 #[test]
2161 fn policy_with_configured_salt_redacts_consistently() {
2162 let cfg = RbacConfig {
2163 enabled: true,
2164 roles: vec![],
2165 redaction_salt: Some(SecretString::from("my-stable-salt")),
2166 };
2167 let p1 = RbacPolicy::new(&cfg);
2168 let p2 = RbacPolicy::new(&cfg);
2169 assert_eq!(
2170 p1.redact_arg("payload"),
2171 p2.redact_arg("payload"),
2172 "policies built from the same configured salt must agree"
2173 );
2174 }
2175
2176 #[test]
2177 fn policy_without_configured_salt_uses_process_salt() {
2178 let cfg = RbacConfig {
2179 enabled: true,
2180 roles: vec![],
2181 redaction_salt: None,
2182 };
2183 let p1 = RbacPolicy::new(&cfg);
2184 let p2 = RbacPolicy::new(&cfg);
2185 assert_eq!(
2187 p1.redact_arg("payload"),
2188 p2.redact_arg("payload"),
2189 "process-wide salt must be consistent within one process"
2190 );
2191 }
2192
2193 #[test]
2194 fn redact_arg_is_fast_enough() {
2195 let salt = b"perf-sanity-salt-32-bytes-padded";
2199 let value = "x".repeat(256);
2200 let start = std::time::Instant::now();
2201 let _ = redact_with_salt(salt, &value);
2202 let elapsed = start.elapsed();
2203 assert!(
2204 elapsed < Duration::from_millis(5),
2205 "single redact_with_salt took {elapsed:?}, expected <5 ms even in debug"
2206 );
2207 }
2208
2209 #[tokio::test]
2221 async fn deny_path_uses_explicit_identity_not_task_local() {
2222 let policy = Arc::new(test_policy());
2223 let id = AuthIdentity {
2224 method: crate::auth::AuthMethod::BearerToken,
2225 name: "alice-the-auditor".into(),
2226 role: "viewer".into(),
2227 raw_token: None,
2228 sub: None,
2229 };
2230 let app = rbac_router_with_identity(policy, id);
2231 let body = tool_call_body("resource_delete", &serde_json::json!({}));
2233 let req = Request::builder()
2234 .method(Method::POST)
2235 .uri("/mcp")
2236 .header("content-type", "application/json")
2237 .body(Body::from(body))
2238 .unwrap();
2239 let resp = app.oneshot(req).await.unwrap();
2240 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2241 }
2242
2243 fn restricted_exec_identity() -> AuthIdentity {
2246 AuthIdentity {
2247 method: crate::auth::AuthMethod::BearerToken,
2248 name: "carol".into(),
2249 role: "restricted-exec".into(),
2250 raw_token: None,
2251 sub: None,
2252 }
2253 }
2254
2255 #[test]
2256 fn has_argument_allowlist_matches_configured_tool_argument() {
2257 let policy = test_policy();
2258 assert!(policy.has_argument_allowlist("restricted-exec", "resource_exec", "cmd"));
2259 assert!(!policy.has_argument_allowlist("restricted-exec", "resource_exec", "host"));
2260 assert!(!policy.has_argument_allowlist("restricted-exec", "other_tool", "cmd"));
2261 assert!(!policy.has_argument_allowlist("ops", "resource_exec", "cmd"));
2262 }
2263
2264 #[tokio::test]
2265 async fn array_arg_with_matching_allowlist_is_denied() {
2266 let policy = Arc::new(test_policy());
2267 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2268 let body = tool_call_body(
2269 "resource_exec",
2270 &serde_json::json!({ "host": "dev-1", "cmd": ["bash", "-c", "evil"] }),
2271 );
2272 let req = Request::builder()
2273 .method(Method::POST)
2274 .uri("/mcp")
2275 .header("content-type", "application/json")
2276 .body(Body::from(body))
2277 .unwrap();
2278 let resp = app.oneshot(req).await.unwrap();
2279 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2280 }
2281
2282 #[tokio::test]
2283 async fn object_arg_with_matching_allowlist_is_denied() {
2284 let policy = Arc::new(test_policy());
2285 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2286 let body = tool_call_body(
2287 "resource_exec",
2288 &serde_json::json!({ "host": "dev-1", "cmd": { "raw": "sh" } }),
2289 );
2290 let req = Request::builder()
2291 .method(Method::POST)
2292 .uri("/mcp")
2293 .header("content-type", "application/json")
2294 .body(Body::from(body))
2295 .unwrap();
2296 let resp = app.oneshot(req).await.unwrap();
2297 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2298 }
2299
2300 #[tokio::test]
2301 async fn number_arg_with_matching_allowlist_is_denied() {
2302 let policy = Arc::new(test_policy());
2303 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2304 let body = tool_call_body(
2305 "resource_exec",
2306 &serde_json::json!({ "host": "dev-1", "cmd": 42 }),
2307 );
2308 let req = Request::builder()
2309 .method(Method::POST)
2310 .uri("/mcp")
2311 .header("content-type", "application/json")
2312 .body(Body::from(body))
2313 .unwrap();
2314 let resp = app.oneshot(req).await.unwrap();
2315 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2316 }
2317
2318 #[tokio::test]
2319 async fn bool_arg_with_matching_allowlist_is_denied() {
2320 let policy = Arc::new(test_policy());
2321 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2322 let body = tool_call_body(
2323 "resource_exec",
2324 &serde_json::json!({ "host": "dev-1", "cmd": true }),
2325 );
2326 let req = Request::builder()
2327 .method(Method::POST)
2328 .uri("/mcp")
2329 .header("content-type", "application/json")
2330 .body(Body::from(body))
2331 .unwrap();
2332 let resp = app.oneshot(req).await.unwrap();
2333 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2334 }
2335
2336 #[tokio::test]
2337 async fn null_arg_with_matching_allowlist_is_denied() {
2338 let policy = Arc::new(test_policy());
2339 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2340 let body = tool_call_body(
2341 "resource_exec",
2342 &serde_json::json!({ "host": "dev-1", "cmd": null }),
2343 );
2344 let req = Request::builder()
2345 .method(Method::POST)
2346 .uri("/mcp")
2347 .header("content-type", "application/json")
2348 .body(Body::from(body))
2349 .unwrap();
2350 let resp = app.oneshot(req).await.unwrap();
2351 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
2352 }
2353
2354 #[tokio::test]
2355 async fn non_string_arg_without_allowlist_is_passthrough() {
2356 let policy = Arc::new(test_policy());
2360 let id = AuthIdentity {
2361 method: crate::auth::AuthMethod::BearerToken,
2362 name: "olivia".into(),
2363 role: "ops".into(),
2364 raw_token: None,
2365 sub: None,
2366 };
2367 let app = rbac_router_with_identity(policy, id);
2368 let body = tool_call_body(
2369 "resource_exec",
2370 &serde_json::json!({ "host": "dev-1", "cmd": ["bash"] }),
2371 );
2372 let req = Request::builder()
2373 .method(Method::POST)
2374 .uri("/mcp")
2375 .header("content-type", "application/json")
2376 .body(Body::from(body))
2377 .unwrap();
2378 let resp = app.oneshot(req).await.unwrap();
2379 assert_ne!(resp.status(), StatusCode::FORBIDDEN);
2380 }
2381
2382 #[tokio::test]
2383 async fn string_arg_in_allowlist_still_passes() {
2384 let policy = Arc::new(test_policy());
2385 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2386 let body = tool_call_body(
2387 "resource_exec",
2388 &serde_json::json!({ "host": "dev-1", "cmd": "bash" }),
2389 );
2390 let req = Request::builder()
2391 .method(Method::POST)
2392 .uri("/mcp")
2393 .header("content-type", "application/json")
2394 .body(Body::from(body))
2395 .unwrap();
2396 let resp = app.oneshot(req).await.unwrap();
2397 assert_ne!(resp.status(), StatusCode::FORBIDDEN);
2398 }
2399
2400 async fn exec_status(args: &serde_json::Value) -> StatusCode {
2409 let policy = Arc::new(test_policy());
2410 let app = rbac_router_with_identity(policy, restricted_exec_identity());
2411 let body = tool_call_body("resource_exec", args);
2412 let req = Request::builder()
2413 .method(Method::POST)
2414 .uri("/mcp")
2415 .header("content-type", "application/json")
2416 .body(Body::from(body))
2417 .unwrap();
2418 app.oneshot(req).await.unwrap().status()
2419 }
2420
2421 #[tokio::test]
2422 async fn non_string_host_is_denied_for_every_json_type() {
2423 for host in [
2424 serde_json::json!(["prod-1"]),
2425 serde_json::json!({ "name": "prod-1" }),
2426 serde_json::json!(42),
2427 serde_json::json!(true),
2428 serde_json::json!(null),
2429 ] {
2430 let args = serde_json::json!({ "host": host, "cmd": "sh" });
2431 assert_eq!(
2432 exec_status(&args).await,
2433 StatusCode::FORBIDDEN,
2434 "non-string host must not bypass host globs: {host:?}"
2435 );
2436 }
2437 }
2438
2439 #[tokio::test]
2440 async fn string_host_outside_globs_still_denied() {
2441 let args = serde_json::json!({ "host": "prod-1", "cmd": "sh" });
2442 assert_eq!(exec_status(&args).await, StatusCode::FORBIDDEN);
2443 }
2444
2445 #[tokio::test]
2446 async fn string_host_inside_globs_still_allowed() {
2447 let args = serde_json::json!({ "host": "dev-1", "cmd": "sh" });
2448 assert_ne!(exec_status(&args).await, StatusCode::FORBIDDEN);
2449 }
2450
2451 #[tokio::test]
2455 async fn absent_host_still_routes_to_check_operation() {
2456 let args = serde_json::json!({ "cmd": "sh" });
2457 assert_ne!(exec_status(&args).await, StatusCode::FORBIDDEN);
2458 }
2459
2460 fn required_policy(allowed: Vec<String>, required: bool) -> RbacPolicy {
2469 let role = RoleConfig::new("viewer", vec!["run".into()], vec!["*".into()])
2470 .with_argument_allowlists(vec![
2471 ArgumentAllowlist::new("run", "cmd", allowed).with_required(required),
2472 ]);
2473 let mut config = RbacConfig::with_roles(vec![role]);
2474 config.enabled = true;
2475 RbacPolicy::new(&config)
2476 }
2477
2478 fn viewer_identity() -> AuthIdentity {
2479 AuthIdentity {
2480 method: crate::auth::AuthMethod::BearerToken,
2481 name: "viewer-1".into(),
2482 role: "viewer".into(),
2483 raw_token: None,
2484 sub: None,
2485 }
2486 }
2487
2488 async fn run_status(policy: RbacPolicy, params: &serde_json::Value) -> StatusCode {
2489 let app = rbac_router_with_identity(Arc::new(policy), viewer_identity());
2490 let body = serde_json::json!({
2491 "jsonrpc": "2.0",
2492 "id": 1,
2493 "method": "tools/call",
2494 "params": params
2495 })
2496 .to_string();
2497 let req = Request::builder()
2498 .method(Method::POST)
2499 .uri("/mcp")
2500 .header("content-type", "application/json")
2501 .body(Body::from(body))
2502 .unwrap();
2503 app.oneshot(req).await.unwrap().status()
2504 }
2505
2506 #[tokio::test]
2507 async fn required_false_still_allows_omitting_the_argument() {
2508 let params = serde_json::json!({ "name": "run", "arguments": {} });
2509 assert_ne!(
2510 run_status(required_policy(vec!["ls".into()], false), ¶ms).await,
2511 StatusCode::FORBIDDEN,
2512 "default behaviour must be unchanged"
2513 );
2514 }
2515
2516 #[tokio::test]
2517 async fn required_true_denies_omitted_argument() {
2518 let params = serde_json::json!({ "name": "run", "arguments": {} });
2519 assert_eq!(
2520 run_status(required_policy(vec!["ls".into()], true), ¶ms).await,
2521 StatusCode::FORBIDDEN
2522 );
2523 }
2524
2525 #[tokio::test]
2526 async fn required_true_allows_permitted_value() {
2527 let params = serde_json::json!({ "name": "run", "arguments": { "cmd": "ls -la" } });
2528 assert_ne!(
2529 run_status(required_policy(vec!["ls".into()], true), ¶ms).await,
2530 StatusCode::FORBIDDEN
2531 );
2532 }
2533
2534 #[tokio::test]
2535 async fn required_true_still_denies_disallowed_value() {
2536 let params = serde_json::json!({ "name": "run", "arguments": { "cmd": "rm -rf /" } });
2537 assert_eq!(
2538 run_status(required_policy(vec!["ls".into()], true), ¶ms).await,
2539 StatusCode::FORBIDDEN
2540 );
2541 }
2542
2543 #[tokio::test]
2544 async fn required_true_denies_non_string_value() {
2545 let params = serde_json::json!({ "name": "run", "arguments": { "cmd": ["ls"] } });
2546 assert_eq!(
2547 run_status(required_policy(vec!["ls".into()], true), ¶ms).await,
2548 StatusCode::FORBIDDEN
2549 );
2550 }
2551
2552 #[tokio::test]
2553 async fn required_true_denies_absent_or_non_object_arguments() {
2554 for params in [
2555 serde_json::json!({ "name": "run" }),
2556 serde_json::json!({ "name": "run", "arguments": "not-an-object" }),
2557 serde_json::json!({ "name": "run", "arguments": null }),
2558 ] {
2559 assert_eq!(
2560 run_status(required_policy(vec!["ls".into()], true), ¶ms).await,
2561 StatusCode::FORBIDDEN,
2562 "omitting the arguments object must not skip `required`: {params:?}"
2563 );
2564 }
2565 }
2566
2567 #[tokio::test]
2570 async fn required_true_with_empty_allowed_accepts_any_string() {
2571 let params =
2572 serde_json::json!({ "name": "run", "arguments": { "cmd": "anything at all" } });
2573 assert_ne!(
2574 run_status(required_policy(vec![], true), ¶ms).await,
2575 StatusCode::FORBIDDEN
2576 );
2577 }
2578
2579 #[tokio::test]
2580 async fn required_true_with_empty_allowed_denies_omitted_argument() {
2581 let params = serde_json::json!({ "name": "run", "arguments": {} });
2582 assert_eq!(
2583 run_status(required_policy(vec![], true), ¶ms).await,
2584 StatusCode::FORBIDDEN
2585 );
2586 }
2587
2588 #[tokio::test]
2589 async fn required_true_with_empty_allowed_denies_non_string() {
2590 let params = serde_json::json!({ "name": "run", "arguments": { "cmd": 42 } });
2591 assert_eq!(
2592 run_status(required_policy(vec![], true), ¶ms).await,
2593 StatusCode::FORBIDDEN
2594 );
2595 }
2596
2597 #[tokio::test]
2598 async fn required_honours_globbed_tool_patterns() {
2599 let role = RoleConfig::new("viewer", vec!["*".into()], vec!["*".into()])
2600 .with_argument_allowlists(vec![
2601 ArgumentAllowlist::new("run-*", "cmd", vec!["ls".into()]).with_required(true),
2602 ]);
2603 let mut config = RbacConfig::with_roles(vec![role]);
2604 config.enabled = true;
2605 let params = serde_json::json!({ "name": "run-foo", "arguments": {} });
2606 assert_eq!(
2607 run_status(RbacPolicy::new(&config), ¶ms).await,
2608 StatusCode::FORBIDDEN,
2609 "a globbed tool pattern must enforce presence, not just value"
2610 );
2611 }
2612
2613 #[test]
2614 fn required_defaults_to_false_when_absent_from_toml() {
2615 let cfg: RbacConfig = toml::from_str(
2616 r#"
2617 enabled = true
2618 [[roles]]
2619 name = "viewer"
2620 allow = ["run"]
2621 [[roles.argument_allowlists]]
2622 tool = "run"
2623 argument = "cmd"
2624 allowed = ["ls"]
2625 "#,
2626 )
2627 .expect("config without `required` must still deserialize");
2628 assert!(
2629 !cfg.roles[0].argument_allowlists[0].required,
2630 "omitted `required` must default to false so existing configs are unchanged"
2631 );
2632 }
2633}