barbacane_lib/mcp/
session.rs

1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4use parking_lot::Mutex;
5use uuid::Uuid;
6
7/// Maximum number of concurrent MCP sessions. `initialize` is unauthenticated,
8/// so without a cap a remote client could create sessions without bound and
9/// exhaust memory before TTL eviction runs (MCP-2). When the cap is reached we
10/// first drop expired sessions; if the store is still full, `create` fails
11/// closed rather than growing unbounded.
12const MAX_SESSIONS: usize = 10_000;
13
14/// MCP session state.
15///
16/// Deliberately holds no client-supplied data: the `clientInfo` blob from
17/// `initialize` is attacker-controlled and was previously retained here unused,
18/// amplifying the unauthenticated-session memory-exhaustion vector (MCP-2). Only
19/// the activity timestamp needed for TTL eviction is kept.
20struct McpSession {
21    last_active: Instant,
22}
23
24/// Thread-safe MCP session store with TTL eviction.
25pub struct SessionStore {
26    sessions: Mutex<HashMap<String, McpSession>>,
27    ttl: Duration,
28}
29
30impl SessionStore {
31    pub fn new(ttl: Duration) -> Self {
32        Self {
33            sessions: Mutex::new(HashMap::new()),
34            ttl,
35        }
36    }
37
38    /// Create a new session and return its ID, or `None` if the session cap is
39    /// reached even after evicting expired sessions (fail closed).
40    pub fn create(&self) -> Option<String> {
41        let mut sessions = self.sessions.lock();
42        if sessions.len() >= MAX_SESSIONS {
43            // Reclaim expired sessions before giving up.
44            let now = Instant::now();
45            sessions.retain(|_, session| now.duration_since(session.last_active) < self.ttl);
46            if sessions.len() >= MAX_SESSIONS {
47                return None;
48            }
49        }
50        let id = Uuid::new_v4().to_string();
51        sessions.insert(
52            id.clone(),
53            McpSession {
54                last_active: Instant::now(),
55            },
56        );
57        Some(id)
58    }
59
60    /// Touch a session (update last_active). Returns false if session doesn't exist.
61    pub fn touch(&self, id: &str) -> bool {
62        let mut sessions = self.sessions.lock();
63        if let Some(session) = sessions.get_mut(id) {
64            session.last_active = Instant::now();
65            true
66        } else {
67            false
68        }
69    }
70
71    /// Remove a session.
72    pub fn remove(&self, id: &str) {
73        self.sessions.lock().remove(id);
74    }
75
76    /// Evict expired sessions. Call periodically from a background task.
77    pub fn evict_expired(&self) {
78        let now = Instant::now();
79        self.sessions
80            .lock()
81            .retain(|_, session| now.duration_since(session.last_active) < self.ttl);
82    }
83
84    /// Number of active sessions.
85    #[cfg(test)]
86    pub fn len(&self) -> usize {
87        self.sessions.lock().len()
88    }
89
90    /// Whether the session store is empty.
91    #[cfg(test)]
92    pub fn is_empty(&self) -> bool {
93        self.sessions.lock().is_empty()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn create_and_touch() {
103        let store = SessionStore::new(Duration::from_secs(60));
104        let id = store.create().expect("under cap");
105        assert!(store.touch(&id));
106        assert!(!store.touch("nonexistent"));
107        assert_eq!(store.len(), 1);
108    }
109
110    #[test]
111    fn remove_session() {
112        let store = SessionStore::new(Duration::from_secs(60));
113        let id = store.create().expect("under cap");
114        assert_eq!(store.len(), 1);
115        store.remove(&id);
116        assert_eq!(store.len(), 0);
117        assert!(!store.touch(&id));
118    }
119
120    #[test]
121    fn evict_expired() {
122        let store = SessionStore::new(Duration::from_millis(1));
123        store.create().expect("under cap");
124        store.create().expect("under cap");
125        assert_eq!(store.len(), 2);
126        std::thread::sleep(Duration::from_millis(10));
127        store.evict_expired();
128        assert_eq!(store.len(), 0);
129    }
130
131    #[test]
132    fn create_fails_closed_at_cap_but_reclaims_expired() {
133        // A tiny TTL so every existing session is expired and reclaimable.
134        let store = SessionStore::new(Duration::from_millis(1));
135        {
136            let mut sessions = store.sessions.lock();
137            for _ in 0..MAX_SESSIONS {
138                sessions.insert(
139                    Uuid::new_v4().to_string(),
140                    McpSession {
141                        last_active: Instant::now(),
142                    },
143                );
144            }
145        }
146        std::thread::sleep(Duration::from_millis(10));
147        // At the cap, but the expired entries are reclaimed so create succeeds.
148        assert!(store.create().is_some());
149
150        // Now fill to the cap with fresh (non-expired) sessions: create fails closed.
151        let store = SessionStore::new(Duration::from_secs(3600));
152        {
153            let mut sessions = store.sessions.lock();
154            for _ in 0..MAX_SESSIONS {
155                sessions.insert(
156                    Uuid::new_v4().to_string(),
157                    McpSession {
158                        last_active: Instant::now(),
159                    },
160                );
161            }
162        }
163        assert!(store.create().is_none());
164    }
165}
barbacane_lib/mcp/session.rs

barbacane_lib/mcp/
session.rs
