Skip to main content

ados_protocol/
mcp_token.rs

1//! Self-contained, scoped ADOS MCP tokens (agent-side verify + on-box mint).
2//!
3//! An MCP client that reaches the drone directly (the `ADOS-MCP` connector in
4//! agent-mode) can present a **scoped, revocable** token instead of the full
5//! pairing key. The token is minted either by the connector's `agent:` issuer or
6//! by this agent's own on-box mint (`ados mcp mint`); both derive the identical
7//! key from the pairing `api_key`, so the agent verifies it with no shared state.
8//!
9//! Wire form (like the plugin capability token): `base64url(json).base64url(hmac)`,
10//! split on the LAST `.`. The signature is `HMAC-SHA256(K, blob_bytes)` over the
11//! EXACT received blob bytes (never a re-serialization). The key
12//! `K = HKDF-SHA256(ikm = api_key, salt = "ados/mcp-token/v1", info = revocation_salt)`
13//! — an HKDF (extract + expand), NOT the single-HMAC derivation the ws-ticket uses,
14//! and the expiry is in MILLISECONDS (not seconds). These three facts are what make
15//! this module byte-for-byte compatible with the connector's `agent:` issuer, and
16//! the known-answer vector in the tests is the guardrail.
17
18use 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
30/// Domain-separation label for the HKDF. MUST equal the connector's
31/// `MCP_TOKEN_HKDF_LABEL = utf8("ados/mcp-token/v1")` (ADOS-MCP src/auth/issuers.ts).
32pub const MCP_TOKEN_HKDF_LABEL: &[u8] = b"ados/mcp-token/v1";
33
34/// The scope groups a token may carry. Mirrors ADOS-MCP src/auth/scopes.ts.
35pub 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/// The verified claim set. Field renames match the connector's TokenClaims.
63#[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, // milliseconds since the epoch
75    #[serde(default)]
76    pub label: String,
77}
78
79/// The scope class a route requires. Carries the connector's safety classes, plus
80/// `SecretRead` for a read route whose body carries a secret: it needs the
81/// `secret_read` scope, NOT plain `read`, so a `read`-only token cannot reach it.
82#[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
105/// True when the granted scope groups admit a route of this class (group-level
106/// membership; the finer tool->capability expansion is the connector's job).
107pub fn scope_allows_class(required: ScopeClass, granted: &[String]) -> bool {
108    let want = required.group_name();
109    granted.iter().any(|g| g == want)
110}
111
112/// HKDF-SHA256 to 32 bytes = HMAC(HMAC(salt, ikm), info || 0x01). One expand block
113/// suffices because the output length equals the SHA-256 block (32 bytes).
114fn 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/// Mints and verifies self-contained MCP tokens, keyed by a pairing-derived HKDF
135/// secret. Build with [`McpTokenIssuer::from_api_key`] from the same `api_key` the
136/// connector's `agent:` issuer uses.
137#[derive(Clone)]
138pub struct McpTokenIssuer {
139    key: [u8; 32],
140}
141
142impl McpTokenIssuer {
143    /// Derive the token key from the pairing `api_key` and the bulk-revocation salt.
144    /// An empty salt is the connector's default (revocationSalt = new Uint8Array(0));
145    /// rotating the salt (a "revoke all") yields a fresh key so every prior token dies.
146    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    /// Verify a token to its claims. Hashes the EXACT received blob bytes, checks the
153    /// HMAC (constant time), parses the claims, rejects an unknown scope, enforces the
154    /// millisecond expiry (`expires_at <= now`), and — when `expected_node_id` is set
155    /// and the issuer is `agent:<id>` — enforces the node subject.
156    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    /// Verify against the current wall clock.
193    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    /// On-box mint (for `ados mcp mint`). The blob is `serde_json::to_vec(claims)` —
202    /// interop with the connector is guaranteed by HKDF + HMAC, not by JSON key order
203    /// (each side verifies received bytes), so no canonical serializer is needed.
204    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        // flip a blob byte, keep the sig -> HMAC over received bytes fails
255        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    // Known-answer vector: pin a token this module mints for a fixed key + claims so
308    // a mirror test in ADOS-MCP (src/auth) can assert its verifyToken accepts the
309    // identical string. If this drifts, cross-repo tokens silently stop verifying.
310    #[test]
311    fn known_answer_vector_for_connector_interop() {
312        let iss = McpTokenIssuer::from_api_key("ados_kat_key", &[]);
313        // a fixed HKDF key + a fixed HMAC over fixed bytes is deterministic
314        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, // 2100-01-01
321            label: "kat".into(),
322        };
323        let tok = iss.mint(&c);
324        // round-trips under the same key (the connector mirror re-derives the key
325        // from "ados_kat_key" and asserts verify(tok) == these claims)
326        assert!(iss.verify(&tok, 1, Some("kat-node")).is_ok());
327        assert!(tok.contains('.'));
328    }
329}