agentmesh_mcp/mcp/
session.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Session authentication for MCP traffic.
5
6use crate::mcp::clock::{Clock, NonceGenerator};
7use crate::mcp::error::McpError;
8use base64::engine::general_purpose::URL_SAFE_NO_PAD;
9use base64::Engine;
10use hmac::{Hmac, Mac};
11use serde::{Deserialize, Serialize};
12use sha2::{Digest, Sha256};
13use std::collections::HashMap;
14use std::sync::{Arc, Mutex};
15use std::time::{Duration, SystemTime, UNIX_EPOCH};
16
17type HmacSha256 = Hmac<Sha256>;
18
19/// Persisted MCP session metadata.
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct McpSession {
22    pub id: String,
23    pub agent_id: String,
24    pub issued_at_secs: u64,
25    pub expires_at_secs: u64,
26    pub token_digest: String,
27}
28
29/// Result of issuing a session.
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31pub struct McpIssuedSession {
32    pub token: String,
33    pub session: McpSession,
34}
35
36#[derive(Debug, Serialize, Deserialize)]
37struct SessionTokenPayload {
38    sid: String,
39    aid: String,
40    exp: u64,
41}
42
43/// Session storage seam.
44pub trait McpSessionStore: Send + Sync {
45    fn get(&self, id: &str) -> Result<Option<McpSession>, McpError>;
46    fn set(&self, session: McpSession) -> Result<(), McpError>;
47    fn delete(&self, id: &str) -> Result<(), McpError>;
48    fn insert_if_below_limit(
49        &self,
50        session: McpSession,
51        concurrent_limit: usize,
52        now_secs: u64,
53    ) -> Result<bool, McpError>;
54}
55
56/// In-memory session storage for defaults and tests.
57#[derive(Debug, Default)]
58pub struct InMemorySessionStore {
59    sessions: Mutex<HashMap<String, McpSession>>,
60}
61
62impl McpSessionStore for InMemorySessionStore {
63    fn get(&self, id: &str) -> Result<Option<McpSession>, McpError> {
64        let sessions = self
65            .sessions
66            .lock()
67            .map_err(|_| McpError::store("session", "session store lock poisoned"))?;
68        Ok(sessions.get(id).cloned())
69    }
70
71    fn set(&self, session: McpSession) -> Result<(), McpError> {
72        let mut sessions = self
73            .sessions
74            .lock()
75            .map_err(|_| McpError::store("session", "session store lock poisoned"))?;
76        sessions.insert(session.id.clone(), session);
77        Ok(())
78    }
79
80    fn delete(&self, id: &str) -> Result<(), McpError> {
81        let mut sessions = self
82            .sessions
83            .lock()
84            .map_err(|_| McpError::store("session", "session store lock poisoned"))?;
85        sessions.remove(id);
86        Ok(())
87    }
88
89    fn insert_if_below_limit(
90        &self,
91        session: McpSession,
92        concurrent_limit: usize,
93        now_secs: u64,
94    ) -> Result<bool, McpError> {
95        let mut sessions = self
96            .sessions
97            .lock()
98            .map_err(|_| McpError::store("session", "session store lock poisoned"))?;
99        sessions.retain(|_, existing| existing.expires_at_secs > now_secs);
100        let active_sessions = sessions
101            .values()
102            .filter(|existing| existing.agent_id == session.agent_id)
103            .count();
104        if active_sessions >= concurrent_limit {
105            return Ok(false);
106        }
107        sessions.insert(session.id.clone(), session);
108        Ok(true)
109    }
110}
111
112/// HMAC-signed session authenticator with TTL and concurrency enforcement.
113#[derive(Clone)]
114pub struct McpSessionAuthenticator {
115    secret: Vec<u8>,
116    clock: Arc<dyn Clock>,
117    nonce_generator: Arc<dyn NonceGenerator>,
118    store: Arc<dyn McpSessionStore>,
119    ttl: Duration,
120    concurrent_limit: usize,
121}
122
123impl McpSessionAuthenticator {
124    pub fn new(
125        secret: Vec<u8>,
126        clock: Arc<dyn Clock>,
127        nonce_generator: Arc<dyn NonceGenerator>,
128        store: Arc<dyn McpSessionStore>,
129        ttl: Duration,
130        concurrent_limit: usize,
131    ) -> Result<Self, McpError> {
132        if secret.is_empty() {
133            return Err(McpError::InvalidConfig("session secret must not be empty"));
134        }
135        if ttl.is_zero() {
136            return Err(McpError::InvalidConfig("session ttl must be > 0"));
137        }
138        if concurrent_limit == 0 {
139            return Err(McpError::InvalidConfig(
140                "concurrent session limit must be > 0",
141            ));
142        }
143        Ok(Self {
144            secret,
145            clock,
146            nonce_generator,
147            store,
148            ttl,
149            concurrent_limit,
150        })
151    }
152
153    pub fn issue_session(&self, agent_id: &str) -> Result<McpIssuedSession, McpError> {
154        let now = unix_secs(self.clock.now())?;
155        let session_id = self.nonce_generator.generate()?;
156        let expires_at_secs = now
157            .checked_add(self.ttl.as_secs())
158            .ok_or(McpError::InvalidConfig("session ttl overflow"))?;
159        let payload = SessionTokenPayload {
160            sid: session_id.clone(),
161            aid: agent_id.to_string(),
162            exp: expires_at_secs,
163        };
164        let encoded_payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload)?);
165        let signature = self.sign(&encoded_payload)?;
166        let token = format!("{encoded_payload}.{signature}");
167        let session = McpSession {
168            id: session_id,
169            agent_id: agent_id.to_string(),
170            issued_at_secs: now,
171            expires_at_secs,
172            token_digest: sha256_hex(&token),
173        };
174        if !self
175            .store
176            .insert_if_below_limit(session.clone(), self.concurrent_limit, now)?
177        {
178            return Err(McpError::SessionLimitExceeded {
179                agent_id: agent_id.to_string(),
180                limit: self.concurrent_limit,
181            });
182        }
183        Ok(McpIssuedSession { token, session })
184    }
185
186    pub fn authenticate(&self, token: &str, agent_id: &str) -> Result<McpSession, McpError> {
187        let (encoded_payload, signature) =
188            token.split_once('.').ok_or(McpError::InvalidTokenFormat)?;
189        self.verify_signature(encoded_payload, signature)?;
190        let payload: SessionTokenPayload = serde_json::from_slice(
191            &URL_SAFE_NO_PAD
192                .decode(encoded_payload)
193                .map_err(|_| McpError::InvalidTokenFormat)?,
194        )?;
195        if payload.aid != agent_id {
196            return Err(McpError::AccessDenied {
197                reason: "session token does not match requested agent".to_string(),
198            });
199        }
200        let now = unix_secs(self.clock.now())?;
201        if payload.exp <= now {
202            self.store.delete(&payload.sid)?;
203            return Err(McpError::SessionExpired);
204        }
205        let session = self
206            .store
207            .get(&payload.sid)?
208            .ok_or_else(|| McpError::AccessDenied {
209                reason: "unknown session".to_string(),
210            })?;
211        if session.agent_id != payload.aid || session.expires_at_secs != payload.exp {
212            return Err(McpError::AccessDenied {
213                reason: "session metadata mismatch".to_string(),
214            });
215        }
216        // Bind the presented token to the stored session record: even with a
217        // valid HMAC signature and matching payload metadata, the token bytes
218        // must hash to the digest captured at issue time. Compared in constant
219        // time so the failure path doesn't leak digest bytes via timing.
220        if !constant_time_eq(
221            sha256_hex(token).as_bytes(),
222            session.token_digest.as_bytes(),
223        ) {
224            return Err(McpError::AccessDenied {
225                reason: "session token digest mismatch".to_string(),
226            });
227        }
228        Ok(session)
229    }
230
231    pub fn revoke(&self, session_id: &str) -> Result<(), McpError> {
232        self.store.delete(session_id)
233    }
234    fn sign(&self, payload: &str) -> Result<String, McpError> {
235        let mut mac =
236            HmacSha256::new_from_slice(&self.secret).map_err(|_| McpError::InvalidHmacKey)?;
237        mac.update(payload.as_bytes());
238        Ok(URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes()))
239    }
240
241    fn verify_signature(&self, payload: &str, signature: &str) -> Result<(), McpError> {
242        let provided = URL_SAFE_NO_PAD
243            .decode(signature)
244            .map_err(|_| McpError::InvalidSignature)?;
245        let mut mac =
246            HmacSha256::new_from_slice(&self.secret).map_err(|_| McpError::InvalidHmacKey)?;
247        mac.update(payload.as_bytes());
248        mac.verify_slice(&provided)
249            .map_err(|_| McpError::InvalidSignature)
250    }
251}
252
253fn sha256_hex(input: &str) -> String {
254    let mut hasher = Sha256::new();
255    hasher.update(input.as_bytes());
256    hasher
257        .finalize()
258        .iter()
259        .map(|byte| format!("{byte:02x}"))
260        .collect()
261}
262
263/// Constant-time byte comparison.
264///
265/// Returns `false` immediately on length mismatch (the length of a session
266/// digest is fixed and public — 64 hex characters — so this leaks nothing),
267/// and otherwise compares every byte via XOR accumulation so the timing
268/// signal carries no information about which byte differed.
269fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
270    if a.len() != b.len() {
271        return false;
272    }
273    let mut diff: u8 = 0;
274    for (x, y) in a.iter().zip(b.iter()) {
275        diff |= x ^ y;
276    }
277    diff == 0
278}
279
280fn unix_secs(time: SystemTime) -> Result<u64, McpError> {
281    Ok(time
282        .duration_since(UNIX_EPOCH)
283        .map_err(|_| McpError::AccessDenied {
284            reason: "system clock before unix epoch".to_string(),
285        })?
286        .as_secs())
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::mcp::clock::{DeterministicNonceGenerator, FixedClock};
293    use std::time::SystemTime;
294
295    #[test]
296    fn issues_and_authenticates_sessions() {
297        let auth = McpSessionAuthenticator::new(
298            b"session-secret".to_vec(),
299            Arc::new(FixedClock::new(SystemTime::UNIX_EPOCH)),
300            Arc::new(DeterministicNonceGenerator::from_values(vec![
301                "session-1".into()
302            ])),
303            Arc::new(InMemorySessionStore::default()),
304            Duration::from_secs(60),
305            1,
306        )
307        .unwrap();
308        let issued = auth.issue_session("did:agentmesh:test").unwrap();
309        let session = auth
310            .authenticate(&issued.token, "did:agentmesh:test")
311            .unwrap();
312        assert_eq!(session.id, "session-1");
313    }
314
315    #[test]
316    fn rejects_session_with_mismatched_token_digest() {
317        // The stored `token_digest` is the issue-time SHA-256 of the issued
318        // token. If a valid-looking token presents against a session record
319        // whose digest differs (e.g. tampered store, regenerated token under
320        // the same session id), authenticate() must refuse it.
321        let store = Arc::new(InMemorySessionStore::default());
322        let auth = McpSessionAuthenticator::new(
323            b"session-secret".to_vec(),
324            Arc::new(FixedClock::new(SystemTime::UNIX_EPOCH)),
325            Arc::new(DeterministicNonceGenerator::from_values(vec![
326                "session-1".into()
327            ])),
328            Arc::clone(&store) as Arc<dyn McpSessionStore>,
329            Duration::from_secs(60),
330            1,
331        )
332        .unwrap();
333        let issued = auth.issue_session("did:agentmesh:test").unwrap();
334
335        // Tamper: replace the stored digest with a different valid SHA-256.
336        let mut tampered = issued.session.clone();
337        tampered.token_digest =
338            "0000000000000000000000000000000000000000000000000000000000000000".to_string();
339        store.set(tampered).unwrap();
340
341        let result = auth.authenticate(&issued.token, "did:agentmesh:test");
342        assert!(matches!(
343            result,
344            Err(McpError::AccessDenied { ref reason }) if reason == "session token digest mismatch"
345        ));
346    }
347
348    #[test]
349    fn constant_time_eq_matches_regular_eq() {
350        assert!(constant_time_eq(b"", b""));
351        assert!(constant_time_eq(b"abc", b"abc"));
352        assert!(!constant_time_eq(b"abc", b"abd"));
353        assert!(!constant_time_eq(b"abc", b"abcd"));
354        assert!(!constant_time_eq(b"abcd", b"abc"));
355        assert!(!constant_time_eq(b"", b"a"));
356    }
357
358    #[test]
359    fn enforces_concurrent_limit() {
360        let auth = McpSessionAuthenticator::new(
361            b"session-secret".to_vec(),
362            Arc::new(FixedClock::new(SystemTime::UNIX_EPOCH)),
363            Arc::new(DeterministicNonceGenerator::from_values(vec![
364                "s1".into(),
365                "s2".into(),
366            ])),
367            Arc::new(InMemorySessionStore::default()),
368            Duration::from_secs(60),
369            1,
370        )
371        .unwrap();
372        auth.issue_session("did:agentmesh:test").unwrap();
373        assert!(matches!(
374            auth.issue_session("did:agentmesh:test"),
375            Err(McpError::SessionLimitExceeded { .. })
376        ));
377    }
378}
agentmesh_mcp/mcp/session.rs

agentmesh_mcp/mcp/
session.rs
