1use std::time::{SystemTime, UNIX_EPOCH};
19
20use base64::Engine;
21use hmac::{Hmac, KeyInit, Mac};
22use sha2::Sha256;
23use thiserror::Error;
24
25type HmacSha256 = Hmac<Sha256>;
26
27const B64: base64::engine::general_purpose::GeneralPurpose =
28 base64::engine::general_purpose::URL_SAFE_NO_PAD;
29
30pub const MCP_TOKEN_HKDF_LABEL: &[u8] = b"ados/mcp-token/v1";
33
34pub const SCOPE_GROUPS: [&str; 6] = [
36 "read",
37 "safe_write",
38 "admin",
39 "flight",
40 "destructive",
41 "secret_read",
42];
43
44#[derive(Debug, Error, PartialEq, Eq)]
45pub enum McpTokenError {
46 #[error("malformed mcp token")]
47 Malformed,
48 #[error("mcp token base64 is invalid")]
49 BadBase64,
50 #[error("mcp token HMAC mismatch")]
51 HmacMismatch,
52 #[error("mcp token claims are not valid JSON")]
53 BadClaims,
54 #[error("mcp token carries an unknown scope")]
55 UnknownScope,
56 #[error("mcp token expired")]
57 Expired,
58 #[error("mcp token issued for a different node")]
59 WrongNode,
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
64pub struct McpClaims {
65 #[serde(rename = "tokenId")]
66 pub token_id: String,
67 #[serde(rename = "operatorId")]
68 pub operator_id: String,
69 pub iss: String,
70 pub scopes: Vec<String>,
71 #[serde(rename = "allowedNodes", default)]
72 pub allowed_nodes: Vec<String>,
73 #[serde(rename = "expiresAt")]
74 pub expires_at: i64, #[serde(default)]
76 pub label: String,
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub enum ScopeClass {
84 Read,
85 SecretRead,
86 SafeWrite,
87 Admin,
88 Flight,
89 Destructive,
90}
91
92impl ScopeClass {
93 pub fn group_name(self) -> &'static str {
94 match self {
95 ScopeClass::Read => "read",
96 ScopeClass::SecretRead => "secret_read",
97 ScopeClass::SafeWrite => "safe_write",
98 ScopeClass::Admin => "admin",
99 ScopeClass::Flight => "flight",
100 ScopeClass::Destructive => "destructive",
101 }
102 }
103}
104
105pub fn scope_allows_class(required: ScopeClass, granted: &[String]) -> bool {
108 let want = required.group_name();
109 granted.iter().any(|g| g == want)
110}
111
112fn hkdf_sha256_32(salt: &[u8], ikm: &[u8], info: &[u8]) -> [u8; 32] {
115 let mut extract = HmacSha256::new_from_slice(salt).expect("HMAC accepts any key length");
116 extract.update(ikm);
117 let prk = extract.finalize().into_bytes();
118 let mut expand = HmacSha256::new_from_slice(&prk).expect("HMAC accepts any key length");
119 expand.update(info);
120 expand.update(&[0x01]);
121 let okm = expand.finalize().into_bytes();
122 let mut out = [0u8; 32];
123 out.copy_from_slice(&okm);
124 out
125}
126
127fn now_unix_ms() -> i64 {
128 SystemTime::now()
129 .duration_since(UNIX_EPOCH)
130 .map(|d| d.as_millis() as i64)
131 .unwrap_or(0)
132}
133
134#[derive(Clone)]
138pub struct McpTokenIssuer {
139 key: [u8; 32],
140}
141
142impl McpTokenIssuer {
143 pub fn from_api_key(api_key: &str, revocation_salt: &[u8]) -> Self {
147 Self {
148 key: hkdf_sha256_32(MCP_TOKEN_HKDF_LABEL, api_key.as_bytes(), revocation_salt),
149 }
150 }
151
152 pub fn verify(
157 &self,
158 token: &str,
159 now_ms: i64,
160 expected_node_id: Option<&str>,
161 ) -> Result<McpClaims, McpTokenError> {
162 let (b_blob, b_sig) = token.rsplit_once('.').ok_or(McpTokenError::Malformed)?;
163 if b_blob.is_empty() || b_sig.is_empty() {
164 return Err(McpTokenError::Malformed);
165 }
166 let blob = B64.decode(b_blob).map_err(|_| McpTokenError::BadBase64)?;
167 let sig = B64.decode(b_sig).map_err(|_| McpTokenError::BadBase64)?;
168 let mut mac = HmacSha256::new_from_slice(&self.key).expect("HMAC accepts any key length");
169 mac.update(&blob);
170 mac.verify_slice(&sig)
171 .map_err(|_| McpTokenError::HmacMismatch)?;
172 let claims: McpClaims =
173 serde_json::from_slice(&blob).map_err(|_| McpTokenError::BadClaims)?;
174 for s in &claims.scopes {
175 if !SCOPE_GROUPS.contains(&s.as_str()) {
176 return Err(McpTokenError::UnknownScope);
177 }
178 }
179 if claims.expires_at <= now_ms {
180 return Err(McpTokenError::Expired);
181 }
182 if let Some(node) = expected_node_id {
183 if let Some(subj) = claims.iss.strip_prefix("agent:") {
184 if subj != node {
185 return Err(McpTokenError::WrongNode);
186 }
187 }
188 }
189 Ok(claims)
190 }
191
192 pub fn verify_now(
194 &self,
195 token: &str,
196 expected_node_id: Option<&str>,
197 ) -> Result<McpClaims, McpTokenError> {
198 self.verify(token, now_unix_ms(), expected_node_id)
199 }
200
201 pub fn mint(&self, claims: &McpClaims) -> String {
205 let blob = serde_json::to_vec(claims).expect("claims serialize");
206 let mut mac = HmacSha256::new_from_slice(&self.key).expect("HMAC accepts any key length");
207 mac.update(&blob);
208 let sig = mac.finalize().into_bytes();
209 format!("{}.{}", B64.encode(&blob), B64.encode(sig))
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216
217 fn claims(exp_ms: i64) -> McpClaims {
218 McpClaims {
219 token_id: "mct_abc".into(),
220 operator_id: "cloud:usr_1".into(),
221 iss: "agent:dev-01".into(),
222 scopes: vec!["read".into(), "admin".into()],
223 allowed_nodes: vec![],
224 expires_at: exp_ms,
225 label: "claude".into(),
226 }
227 }
228
229 #[test]
230 fn mint_round_trips_and_verifies() {
231 let iss = McpTokenIssuer::from_api_key("ados_secret", &[]);
232 let tok = iss.mint(&claims(2_000_000_000_000));
233 let c = iss.verify(&tok, 1_000_000_000_000, Some("dev-01")).unwrap();
234 assert_eq!(c.token_id, "mct_abc");
235 assert_eq!(c.scopes, vec!["read", "admin"]);
236 }
237
238 #[test]
239 fn rejects_expired_at_the_boundary() {
240 let iss = McpTokenIssuer::from_api_key("k", &[]);
241 let tok = iss.mint(&claims(1_000));
242 assert_eq!(iss.verify(&tok, 1_000, None), Err(McpTokenError::Expired));
243 assert!(iss.verify(&tok, 999, None).is_ok());
244 }
245
246 #[test]
247 fn rejects_wrong_key_and_tampered_blob() {
248 let iss = McpTokenIssuer::from_api_key("k1", &[]);
249 let tok = iss.mint(&claims(2_000_000_000_000));
250 assert_eq!(
251 McpTokenIssuer::from_api_key("k2", &[]).verify(&tok, 1, None),
252 Err(McpTokenError::HmacMismatch)
253 );
254 let (blob, sig) = tok.rsplit_once('.').unwrap();
256 let mut raw = B64.decode(blob).unwrap();
257 raw[0] ^= 0xff;
258 let tampered = format!("{}.{}", B64.encode(&raw), sig);
259 assert_eq!(
260 iss.verify(&tampered, 1, None),
261 Err(McpTokenError::HmacMismatch)
262 );
263 }
264
265 #[test]
266 fn rejects_unknown_scope_and_wrong_node() {
267 let iss = McpTokenIssuer::from_api_key("k", &[]);
268 let mut c = claims(2_000_000_000_000);
269 c.scopes = vec!["read".into(), "superuser".into()];
270 assert_eq!(
271 iss.verify(&iss.mint(&c), 1, None),
272 Err(McpTokenError::UnknownScope)
273 );
274 let ok = iss.mint(&claims(2_000_000_000_000));
275 assert_eq!(
276 iss.verify(&ok, 1, Some("other")),
277 Err(McpTokenError::WrongNode)
278 );
279 }
280
281 #[test]
282 fn rejects_malformed() {
283 let iss = McpTokenIssuer::from_api_key("k", &[]);
284 for bad in ["", "nodot", "a.", ".b", "!!.??"] {
285 assert!(iss.verify(bad, 1, None).is_err());
286 }
287 }
288
289 #[test]
290 fn revocation_salt_fold_revokes_all() {
291 let a = McpTokenIssuer::from_api_key("k", b"salt-a");
292 let b = McpTokenIssuer::from_api_key("k", b"salt-b");
293 let tok = a.mint(&claims(2_000_000_000_000));
294 assert!(a.verify(&tok, 1, None).is_ok());
295 assert_eq!(b.verify(&tok, 1, None), Err(McpTokenError::HmacMismatch));
296 }
297
298 #[test]
299 fn scope_class_membership() {
300 let granted = vec!["read".to_string(), "admin".to_string()];
301 assert!(scope_allows_class(ScopeClass::Read, &granted));
302 assert!(scope_allows_class(ScopeClass::Admin, &granted));
303 assert!(!scope_allows_class(ScopeClass::Flight, &granted));
304 assert!(!scope_allows_class(ScopeClass::SafeWrite, &granted));
305 }
306
307 #[test]
311 fn known_answer_vector_for_connector_interop() {
312 let iss = McpTokenIssuer::from_api_key("ados_kat_key", &[]);
313 let c = McpClaims {
315 token_id: "mct_kat".into(),
316 operator_id: "cloud:kat".into(),
317 iss: "agent:kat-node".into(),
318 scopes: vec!["read".into()],
319 allowed_nodes: vec![],
320 expires_at: 4_102_444_800_000, label: "kat".into(),
322 };
323 let tok = iss.mint(&c);
324 assert!(iss.verify(&tok, 1, Some("kat-node")).is_ok());
327 assert!(tok.contains('.'));
328 }
329}