mcp_airlock/
vault.rs

1//! # OS-Native Secure Vault
2//!
3//! This module provides an abstraction over the system's native secure storage via the
4//! `keyring` crate: macOS Keychain, Windows Credential Manager, and on Linux the kernel
5//! keyring (keyutils) as a cache in front of the Secret Service, which persists across
6//! reboots.
7//!
8//! It also includes an in-memory backend for headless or testing environments.
9
10use crate::Result;
11use anyhow::Context;
12use keyring::Entry;
13use std::collections::HashMap;
14use std::sync::{Arc, Mutex};
15
16/// The kind of secret stored in the vault. Each kind lives in its own keyring entry.
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18enum Secret {
19    Token,
20    DpopKey,
21    RefreshToken,
22    Meta,
23}
24
25impl Secret {
26    /// Suffix used for the in-memory key and the keyring service name.
27    fn suffix(self) -> &'static str {
28        match self {
29            Secret::Token => "token",
30            Secret::DpopKey => "dpop",
31            Secret::RefreshToken => "refresh",
32            Secret::Meta => "meta",
33        }
34    }
35
36    fn label(self) -> &'static str {
37        match self {
38            Secret::Token => "token",
39            Secret::DpopKey => "DPoP key",
40            Secret::RefreshToken => "refresh token",
41            Secret::Meta => "credential metadata",
42        }
43    }
44}
45
46/// Extra guidance for keychain failures on Linux, where storage needs a
47/// running Secret Service.
48#[cfg(target_os = "linux")]
49const KEYCHAIN_HINT: &str = " (on Linux this needs a running Secret Service such as GNOME \
50    Keyring or KWallet; MCP_AIRLOCK_USE_MEMORY_VAULT=1 keeps credentials in memory only)";
51#[cfg(not(target_os = "linux"))]
52const KEYCHAIN_HINT: &str = "";
53
54#[derive(Clone)]
55enum Backend {
56    /// The operating system's credential store.
57    Keyring,
58    /// A process-local store. Clones of the same `Vault` share it.
59    Memory(Arc<Mutex<HashMap<String, String>>>),
60}
61
62/// What the stored credentials were issued for.
63#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
64pub struct CredentialMeta {
65    /// Issuer of the authorization server that issued the tokens. Credentials
66    /// are never used with another authorization server.
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub issuer: Option<String>,
69    /// Scopes requested when the tokens were obtained.
70    #[serde(default, skip_serializing_if = "Vec::is_empty")]
71    pub scopes: Vec<String>,
72    /// When the access token expires (seconds since the Unix epoch), if known.
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub expires_at: Option<u64>,
75}
76
77/// A secure storage abstraction for tokens and keys.
78#[derive(Clone)]
79pub struct Vault {
80    /// The service name used for isolation in the keychain.
81    pub service: String,
82    backend: Backend,
83}
84
85/// The keychain service name for a remote MCP server.
86///
87/// Credentials are scoped per server so a token issued for one resource is never
88/// sent to another (RFC 8707).
89pub fn service_name_for(remote_url: &str) -> String {
90    use sha2::{Digest, Sha256};
91    let digest = Sha256::digest(remote_url.trim_end_matches('/').as_bytes());
92    format!("mcp-airlock:{}", &hex::encode(digest)[..16])
93}
94
95impl Vault {
96    /// Creates a vault backed by the OS keychain.
97    pub fn keyring(service: &str) -> Self {
98        Self {
99            service: service.to_string(),
100            backend: Backend::Keyring,
101        }
102    }
103
104    /// Creates a vault backed by a fresh, process-local in-memory store.
105    pub fn in_memory(service: &str) -> Self {
106        Self {
107            service: service.to_string(),
108            backend: Backend::Memory(Arc::new(Mutex::new(HashMap::new()))),
109        }
110    }
111
112    /// Uses the in-memory backend when `MCP_AIRLOCK_USE_MEMORY_VAULT` is set,
113    /// otherwise the OS keychain.
114    pub fn from_env(service: &str) -> Self {
115        if std::env::var("MCP_AIRLOCK_USE_MEMORY_VAULT").is_ok() {
116            Self::in_memory(service)
117        } else {
118            Self::keyring(service)
119        }
120    }
121
122    fn make_key(&self, user_id: &str, suffix: &str) -> String {
123        format!("{}:{}:{}", self.service, user_id, suffix)
124    }
125
126    fn keyring_entry(&self, user_id: &str, secret: Secret) -> Result<Entry> {
127        let service = match secret {
128            // Kept as the bare service name for backwards compatibility.
129            Secret::Token => self.service.clone(),
130            _ => format!("{}-{}", self.service, secret.suffix()),
131        };
132        Ok(Entry::new(&service, user_id)?)
133    }
134
135    fn set(&self, user_id: &str, secret: Secret, value: &str) -> Result<()> {
136        match &self.backend {
137            Backend::Memory(store) => {
138                store
139                    .lock()
140                    .map_err(|e| anyhow::anyhow!("Mutex poisoned: {}", e))?
141                    .insert(self.make_key(user_id, secret.suffix()), value.to_string());
142                Ok(())
143            }
144            Backend::Keyring => self
145                .keyring_entry(user_id, secret)?
146                .set_password(value)
147                .with_context(|| {
148                    format!(
149                        "Failed to store {} in the OS keychain{}",
150                        secret.label(),
151                        KEYCHAIN_HINT
152                    )
153                }),
154        }
155    }
156
157    fn get(&self, user_id: &str, secret: Secret) -> Result<Option<String>> {
158        match &self.backend {
159            Backend::Memory(store) => Ok(store
160                .lock()
161                .map_err(|e| anyhow::anyhow!("Mutex poisoned: {}", e))?
162                .get(&self.make_key(user_id, secret.suffix()))
163                .cloned()),
164            Backend::Keyring => match self.keyring_entry(user_id, secret)?.get_password() {
165                Ok(v) => Ok(Some(v)),
166                Err(keyring::Error::NoEntry) => Ok(None),
167                Err(e) => Err(anyhow::anyhow!(e)
168                    .context(format!("Failed to retrieve {} from vault", secret.label()))),
169            },
170        }
171    }
172
173    fn delete(&self, user_id: &str, secret: Secret) -> Result<()> {
174        match &self.backend {
175            Backend::Memory(store) => {
176                store
177                    .lock()
178                    .map_err(|e| anyhow::anyhow!("Mutex poisoned: {}", e))?
179                    .remove(&self.make_key(user_id, secret.suffix()));
180            }
181            Backend::Keyring => {
182                let _ = self.keyring_entry(user_id, secret)?.delete_credential();
183            }
184        }
185        Ok(())
186    }
187
188    /// Stores an access token securely in the vault.
189    pub fn store_token(&self, user_id: &str, token: &str) -> Result<()> {
190        self.set(user_id, Secret::Token, token)
191    }
192
193    /// Retrieves an access token from the vault.
194    pub fn get_token(&self, user_id: &str) -> Result<Option<String>> {
195        self.get(user_id, Secret::Token)
196    }
197
198    /// Deletes an access token from the vault.
199    pub fn delete_token(&self, user_id: &str) -> Result<()> {
200        self.delete(user_id, Secret::Token)
201    }
202
203    /// Stores the DPoP private key securely.
204    pub fn store_dpop_key(&self, user_id: &str, key_bytes: &[u8]) -> Result<()> {
205        self.set(user_id, Secret::DpopKey, &hex::encode(key_bytes))
206    }
207
208    /// Retrieves the DPoP private key from the vault.
209    pub fn get_dpop_key(&self, user_id: &str) -> Result<Option<Vec<u8>>> {
210        self.get(user_id, Secret::DpopKey)?
211            .map(|h| hex::decode(h).context("Failed to decode DPoP key hex"))
212            .transpose()
213    }
214
215    /// Deletes the DPoP private key from the vault.
216    pub fn delete_dpop_key(&self, user_id: &str) -> Result<()> {
217        self.delete(user_id, Secret::DpopKey)
218    }
219
220    /// Stores a refresh token securely in the vault.
221    pub fn store_refresh_token(&self, user_id: &str, token: &str) -> Result<()> {
222        self.set(user_id, Secret::RefreshToken, token)
223    }
224
225    /// Retrieves the refresh token from the vault.
226    pub fn get_refresh_token(&self, user_id: &str) -> Result<Option<String>> {
227        self.get(user_id, Secret::RefreshToken)
228    }
229
230    /// Deletes the refresh token from the vault.
231    pub fn delete_refresh_token(&self, user_id: &str) -> Result<()> {
232        self.delete(user_id, Secret::RefreshToken)
233    }
234
235    /// Stores what the current credentials were issued for.
236    pub fn store_meta(&self, user_id: &str, meta: &CredentialMeta) -> Result<()> {
237        self.set(user_id, Secret::Meta, &serde_json::to_string(meta)?)
238    }
239
240    /// Retrieves the credential metadata. Unreadable metadata counts as none.
241    pub fn get_meta(&self, user_id: &str) -> Result<Option<CredentialMeta>> {
242        Ok(self
243            .get(user_id, Secret::Meta)?
244            .and_then(|m| serde_json::from_str(&m).ok()))
245    }
246
247    /// Deletes every credential stored for `user_id`.
248    pub fn clear(&self, user_id: &str) -> Result<()> {
249        for secret in [
250            Secret::Token,
251            Secret::DpopKey,
252            Secret::RefreshToken,
253            Secret::Meta,
254        ] {
255            self.delete(user_id, secret)?;
256        }
257        Ok(())
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn test_service_name_for() {
267        let a = service_name_for("https://a.example.com/mcp");
268        assert!(a.starts_with("mcp-airlock:"));
269        assert_eq!(a.len(), "mcp-airlock:".len() + 16);
270        assert_eq!(a, service_name_for("https://a.example.com/mcp/"));
271        assert_ne!(a, service_name_for("https://b.example.com/mcp"));
272    }
273
274    #[test]
275    fn test_make_key() {
276        let vault = Vault::in_memory("my-service");
277
278        // Happy path
279        assert_eq!(
280            vault.make_key("user123", "token"),
281            "my-service:user123:token"
282        );
283
284        // Edge cases
285        assert_eq!(vault.make_key("", "token"), "my-service::token");
286        assert_eq!(vault.make_key("user123", ""), "my-service:user123:");
287        assert_eq!(vault.make_key("", ""), "my-service::");
288
289        // Special characters
290        assert_eq!(
291            vault.make_key("user@domain.com", "dpop"),
292            "my-service:user@domain.com:dpop"
293        );
294    }
295
296    #[test]
297    fn test_vault_token_ops() -> Result<()> {
298        let vault = Vault::in_memory("mcp-airlock-test");
299        let user = "test_user_1";
300        let token = "test_token_123";
301
302        vault.store_token(user, token)?;
303        assert_eq!(vault.get_token(user)?, Some(token.to_string()));
304
305        vault.delete_token(user)?;
306        assert_eq!(vault.get_token(user)?, None);
307        Ok(())
308    }
309
310    #[test]
311    fn test_vault_refresh_token_ops() -> Result<()> {
312        let vault = Vault::in_memory("mcp-airlock-test");
313        vault.store_token("u", "access")?;
314        vault.store_refresh_token("u", "refresh")?;
315        assert_eq!(vault.get_refresh_token("u")?, Some("refresh".into()));
316        // Separate entries: the access token is untouched.
317        assert_eq!(vault.get_token("u")?, Some("access".into()));
318        vault.delete_refresh_token("u")?;
319        assert_eq!(vault.get_refresh_token("u")?, None);
320        assert_eq!(vault.get_token("u")?, Some("access".into()));
321        Ok(())
322    }
323
324    #[test]
325    fn test_vault_meta_and_clear() -> Result<()> {
326        let vault = Vault::in_memory("svc");
327        assert_eq!(vault.get_meta("u")?, None);
328        let meta = CredentialMeta {
329            issuer: Some("https://as.example.com".into()),
330            scopes: vec!["read".into()],
331            expires_at: Some(42),
332        };
333        vault.store_meta("u", &meta)?;
334        assert_eq!(vault.get_meta("u")?, Some(meta));
335
336        vault.store_token("u", "t")?;
337        vault.store_refresh_token("u", "r")?;
338        vault.clear("u")?;
339        assert_eq!(vault.get_token("u")?, None);
340        assert_eq!(vault.get_refresh_token("u")?, None);
341        assert_eq!(vault.get_meta("u")?, None);
342
343        vault.set("u", Secret::Meta, "not json")?;
344        assert_eq!(vault.get_meta("u")?, None);
345        Ok(())
346    }
347
348    #[test]
349    fn test_delete_nonexistent_token() -> Result<()> {
350        let vault = Vault::in_memory("mcp-airlock-test-nonexistent");
351        assert!(vault.delete_token("non_existent_user").is_ok());
352        Ok(())
353    }
354
355    #[test]
356    fn test_vault_dpop_ops() -> Result<()> {
357        let vault = Vault::in_memory("mcp-airlock-test");
358        let user = "test_user_dpop";
359        let key_bytes = b"test_key_bytes_123456789012345678";
360
361        vault.store_dpop_key(user, key_bytes)?;
362        assert_eq!(vault.get_dpop_key(user)?, Some(key_bytes.to_vec()));
363
364        assert_eq!(vault.get_dpop_key("non_existent")?, None);
365
366        vault.delete_dpop_key(user)?;
367        assert_eq!(vault.get_dpop_key(user)?, None);
368        Ok(())
369    }
370
371    #[test]
372    fn test_delete_nonexistent_dpop_key() -> Result<()> {
373        let vault = Vault::in_memory("mcp-airlock-test-nonexistent-dpop");
374        assert!(vault.delete_dpop_key("non_existent_user").is_ok());
375        Ok(())
376    }
377
378    #[test]
379    fn test_vault_dpop_hex_failure() -> Result<()> {
380        let vault = Vault::in_memory("mcp-airlock-test");
381        let user = "test_user_bad_hex";
382        vault.set(user, Secret::DpopKey, "invalid hex")?;
383
384        let res = vault.get_dpop_key(user);
385        assert!(res.is_err());
386        assert!(format!("{:?}", res.err().unwrap()).contains("Failed to decode DPoP key hex"));
387        Ok(())
388    }
389
390    #[test]
391    fn test_in_memory_vaults_are_isolated_but_clones_share() -> Result<()> {
392        let a = Vault::in_memory("svc");
393        let b = Vault::in_memory("svc");
394        let a2 = a.clone();
395
396        a.store_token("user", "token-a")?;
397        assert_eq!(a2.get_token("user")?, Some("token-a".into()));
398        assert_eq!(b.get_token("user")?, None);
399        Ok(())
400    }
401
402    /// Needs a real OS keychain (on Linux: a running Secret Service).
403    /// Run with `cargo test --lib keyring_roundtrip -- --ignored`.
404    #[test]
405    #[ignore]
406    fn test_keyring_roundtrip_persists_across_instances() -> Result<()> {
407        let service = format!("mcp-airlock-roundtrip-{}", uuid::Uuid::new_v4());
408        let writer = Vault::keyring(&service);
409        writer.store_token("u", "token-123")?;
410        writer.store_dpop_key("u", &[1, 2, 3])?;
411        writer.store_refresh_token("u", "refresh-456")?;
412
413        // A different instance must see what the first one stored.
414        let reader = Vault::keyring(&service);
415        let result = (
416            reader.get_token("u"),
417            reader.get_dpop_key("u"),
418            reader.get_refresh_token("u"),
419        );
420        reader.delete_token("u")?;
421        reader.delete_dpop_key("u")?;
422        reader.delete_refresh_token("u")?;
423
424        assert_eq!(result.0?, Some("token-123".into()));
425        assert_eq!(result.1?, Some(vec![1, 2, 3]));
426        assert_eq!(result.2?, Some("refresh-456".into()));
427        assert_eq!(reader.get_token("u")?, None);
428        Ok(())
429    }
430
431    #[test]
432    fn test_vault_real_keyring_attempt() {
433        let vault = Vault::keyring("mcp-airlock-unit-test-real");
434
435        // This will likely fail in CI but it's okay, we just want to cover the lines.
436        // We use a dummy user to avoid messing up real keys.
437        let _ = vault.store_token("dummy_user_test", "dummy_token");
438        let _ = vault.get_token("dummy_user_test");
439        let _ = vault.delete_token("dummy_user_test");
440        let _ = vault.store_dpop_key("dummy_user_test", b"dummy");
441        let _ = vault.get_dpop_key("dummy_user_test");
442        let _ = vault.delete_dpop_key("dummy_user_test");
443    }
444}
mcp_airlock/vault.rs

mcp_airlock/
vault.rs
