vtcode_webmcp/
pairing.rs

1use crate::error::{Result, WebmcpError};
2use std::collections::{HashMap, HashSet};
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, Instant};
5use uuid::Uuid;
6use vtcode_commons::validation::is_valid_origin;
7
8const PAIRING_CODE_HEX_DIGITS: usize = 12;
9const MAX_FAILED_PAIRING_ATTEMPTS: u8 = 5;
10
11#[derive(Debug)]
12struct PendingPairing {
13    code: String,
14    origin: Option<String>,
15    expires_at: Instant,
16    used: bool,
17    failed_attempts: u8,
18}
19
20#[derive(Debug)]
21struct PairingState {
22    pending: Option<PendingPairing>,
23    sessions: HashMap<String, PairingSessionState>,
24}
25
26#[derive(Debug)]
27struct PairingSessionState {
28    origin: String,
29    expires_at: Instant,
30}
31
32/// A terminal-displayable, short-lived pairing code.
33#[derive(Clone)]
34pub struct PairingDisplay {
35    code: String,
36    expires_at: Instant,
37}
38
39impl PairingDisplay {
40    /// Returns the code that must be entered by the browser.
41    pub fn code(&self) -> &str {
42        &self.code
43    }
44
45    /// Returns the remaining lifetime of the code.
46    pub fn expires_in(&self) -> Duration {
47        self.expires_at.saturating_duration_since(Instant::now())
48    }
49}
50
51impl std::fmt::Debug for PairingDisplay {
52    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        formatter
54            .debug_struct("PairingDisplay")
55            .field("expires_in", &self.expires_in())
56            .finish_non_exhaustive()
57    }
58}
59
60/// An authenticated, in-memory browser session.
61#[derive(Clone)]
62pub struct PairingSession {
63    token: String,
64    origin: String,
65    expires_at: Instant,
66}
67
68impl PairingSession {
69    /// Returns the bearer token. Callers should keep it in memory only.
70    pub fn token(&self) -> &str {
71        &self.token
72    }
73
74    /// Returns the origin bound to this session.
75    pub fn origin(&self) -> &str {
76        &self.origin
77    }
78
79    /// Returns the remaining inactivity-lease lifetime.
80    pub fn expires_in(&self) -> Duration {
81        self.expires_at.saturating_duration_since(Instant::now())
82    }
83}
84
85impl std::fmt::Debug for PairingSession {
86    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
87        formatter
88            .debug_struct("PairingSession")
89            .field("origin", &self.origin)
90            .field("expires_in", &self.expires_in())
91            .finish_non_exhaustive()
92    }
93}
94
95/// Manages one-time pairing and revocable in-memory sessions.
96#[derive(Clone)]
97pub struct PairingManager {
98    allowed_origins: Arc<HashSet<String>>,
99    ttl: Duration,
100    state: Arc<Mutex<PairingState>>,
101}
102
103impl PairingManager {
104    /// Create a manager with an explicit origin allowlist and inactivity lease.
105    pub fn new<I, S>(allowed_origins: I, ttl: Duration) -> Result<Self>
106    where
107        I: IntoIterator<Item = S>,
108        S: Into<String>,
109    {
110        if ttl.is_zero() || ttl > Duration::from_secs(3600) {
111            return Err(WebmcpError::InvalidRequest("pairing ttl must be between 1 and 3600 seconds".to_string()));
112        }
113
114        let origins = allowed_origins.into_iter().map(Into::into).collect::<Vec<String>>();
115        if origins.iter().any(|origin| !is_valid_origin(origin)) {
116            return Err(WebmcpError::InvalidRequest("WebMCP origins must be exact non-wildcard origins".to_string()));
117        }
118
119        Ok(Self {
120            allowed_origins: Arc::new(origins.into_iter().collect()),
121            ttl,
122            state: Arc::new(Mutex::new(PairingState { pending: None, sessions: HashMap::new() })),
123        })
124    }
125
126    /// Starts a fresh pairing code, invalidating any older unconsumed code.
127    pub fn begin_pairing(&self) -> PairingDisplay {
128        self.begin_pairing_inner(None, false)
129    }
130
131    fn begin_pairing_inner(&self, origin: Option<String>, revoke_sessions: bool) -> PairingDisplay {
132        let code = Uuid::new_v4()
133            .simple()
134            .to_string()
135            .chars()
136            .take(PAIRING_CODE_HEX_DIGITS)
137            .collect::<String>()
138            .to_ascii_uppercase();
139        let expires_at = Instant::now() + self.ttl;
140        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
141        if revoke_sessions {
142            state.sessions.clear();
143        }
144        state.pending = Some(PendingPairing {
145            code: code.clone(),
146            origin,
147            expires_at,
148            used: false,
149            failed_attempts: 0,
150        });
151        PairingDisplay { code, expires_at }
152    }
153
154    /// Starts a code bound to a specific allowed origin.
155    pub fn begin_pairing_for_origin(&self, origin: impl Into<String>) -> Result<PairingDisplay> {
156        let origin = origin.into();
157        self.ensure_origin_allowed(&origin)?;
158        Ok(self.begin_pairing_inner(Some(origin), false))
159    }
160
161    /// Revoke all browser sessions and issue a fresh code in one state update.
162    pub fn replace_pairing_for_origin(&self, origin: impl Into<String>) -> Result<PairingDisplay> {
163        let origin = origin.into();
164        self.ensure_origin_allowed(&origin)?;
165        Ok(self.begin_pairing_inner(Some(origin), true))
166    }
167
168    /// Consumes a code and creates a session token bound to the browser origin.
169    pub fn pair(&self, code: &str, origin: &str) -> Result<PairingSession> {
170        self.ensure_origin_allowed(origin)?;
171        let now = Instant::now();
172        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
173        let pending = state.pending.as_mut().ok_or(WebmcpError::PairingExpired)?;
174        if pending.expires_at <= now {
175            state.pending = None;
176            return Err(WebmcpError::PairingExpired);
177        }
178        if pending.used {
179            return Err(WebmcpError::PairingUsed);
180        }
181        if pending.code != code {
182            pending.failed_attempts = pending.failed_attempts.saturating_add(1);
183            if pending.failed_attempts >= MAX_FAILED_PAIRING_ATTEMPTS {
184                state.pending = None;
185            }
186            return Err(WebmcpError::PairingExpired);
187        }
188        if pending.origin.as_deref().is_some_and(|expected| expected != origin) {
189            return Err(WebmcpError::OriginRejected(origin.to_string()));
190        }
191
192        pending.used = true;
193        let token = Uuid::new_v4().simple().to_string();
194        let expires_at = now + self.ttl;
195        drop(
196            state
197                .sessions
198                .insert(token.clone(), PairingSessionState { origin: origin.to_string(), expires_at }),
199        );
200        Ok(PairingSession { token, origin: origin.to_string(), expires_at })
201    }
202
203    /// Validates a session token and its origin binding without extending it.
204    pub fn validate(&self, token: &str, origin: &str) -> Result<()> {
205        self.validate_session(token, origin, false)
206    }
207
208    /// Validates a session and extends its inactivity deadline.
209    ///
210    /// Pairing codes remain one-time and expire according to the configured
211    /// TTL. An authenticated browser session may remain connected longer than
212    /// that TTL as long as it continues making authenticated requests.
213    pub fn refresh(&self, token: &str, origin: &str) -> Result<()> {
214        self.validate_session(token, origin, true)
215    }
216
217    fn validate_session(&self, token: &str, origin: &str, refresh: bool) -> Result<()> {
218        self.ensure_origin_allowed(origin)?;
219        let now = Instant::now();
220        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
221        let valid = state
222            .sessions
223            .get(token)
224            .is_some_and(|session| session.expires_at > now && session.origin == origin);
225        if !valid {
226            drop(state.sessions.remove(token));
227            return Err(WebmcpError::Unauthorized);
228        }
229        if refresh && let Some(session) = state.sessions.get_mut(token) {
230            session.expires_at = now + self.ttl;
231        }
232        Ok(())
233    }
234
235    /// Rehydrates an existing in-memory session for a reconnecting socket.
236    pub fn resume(&self, token: &str, origin: &str) -> Result<PairingSession> {
237        self.refresh(token, origin)?;
238        let state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
239        let session = state.sessions.get(token).ok_or(WebmcpError::Unauthorized)?;
240        Ok(PairingSession {
241            token: token.to_string(),
242            origin: session.origin.clone(),
243            expires_at: session.expires_at,
244        })
245    }
246
247    /// Revokes one session token.
248    pub fn revoke(&self, token: &str) -> bool {
249        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
250        state.sessions.remove(token).is_some()
251    }
252
253    /// Revokes every active browser session and pending code.
254    pub fn revoke_all(&self) {
255        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
256        state.sessions.clear();
257        state.pending = None;
258    }
259
260    /// Returns whether the origin is explicitly allowed.
261    pub fn is_origin_allowed(&self, origin: &str) -> bool {
262        self.allowed_origins.contains(origin)
263    }
264
265    fn ensure_origin_allowed(&self, origin: &str) -> Result<()> {
266        if origin.is_empty() || !self.allowed_origins.contains(origin) {
267            return Err(WebmcpError::OriginRejected(origin.to_string()));
268        }
269        Ok(())
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn pairing_is_one_time_and_origin_bound() {
279        let manager = PairingManager::new(["https://example.test"], Duration::from_secs(60)).expect("manager");
280        let display = manager.begin_pairing();
281        assert_eq!(display.code().len(), 12);
282        let session = manager.pair(display.code(), "https://example.test").expect("pair");
283        assert!(manager.validate(session.token(), "https://example.test").is_ok());
284        assert!(matches!(manager.pair(display.code(), "https://example.test"), Err(WebmcpError::PairingUsed)));
285        assert!(matches!(
286            manager.validate(session.token(), "https://other.test"),
287            Err(WebmcpError::OriginRejected(_))
288        ));
289    }
290
291    #[test]
292    fn configured_origins_can_keep_independent_sessions_active() {
293        let manager = PairingManager::new(
294            ["https://vtcode.vinhnx.chatgpt.site", "https://vinhnx.github.io"],
295            Duration::from_secs(60),
296        )
297        .expect("manager");
298
299        let chatgpt_pairing = manager
300            .begin_pairing_for_origin("https://vtcode.vinhnx.chatgpt.site")
301            .expect("ChatGPT Site pairing");
302        let chatgpt_session = manager
303            .pair(chatgpt_pairing.code(), "https://vtcode.vinhnx.chatgpt.site")
304            .expect("ChatGPT Site session");
305
306        let github_pairing = manager
307            .begin_pairing_for_origin("https://vinhnx.github.io")
308            .expect("GitHub Pages pairing");
309        let github_session = manager
310            .pair(github_pairing.code(), "https://vinhnx.github.io")
311            .expect("GitHub Pages session");
312
313        assert!(
314            manager
315                .validate(chatgpt_session.token(), "https://vtcode.vinhnx.chatgpt.site")
316                .is_ok()
317        );
318        assert!(manager.validate(github_session.token(), "https://vinhnx.github.io").is_ok());
319        assert!(matches!(
320            manager.validate(chatgpt_session.token(), "https://vinhnx.github.io"),
321            Err(WebmcpError::Unauthorized)
322        ));
323    }
324
325    #[test]
326    fn expired_pairing_cannot_be_reused() {
327        let manager = PairingManager::new(["https://example.test"], Duration::from_millis(1)).expect("manager");
328        let display = manager.begin_pairing();
329        std::thread::sleep(Duration::from_millis(3));
330        assert!(matches!(manager.pair(display.code(), "https://example.test"), Err(WebmcpError::PairingExpired)));
331    }
332
333    #[test]
334    fn revocation_invalidates_session() {
335        let manager = PairingManager::new(["https://example.test"], Duration::from_secs(60)).expect("manager");
336        let display = manager.begin_pairing();
337        let session = manager.pair(display.code(), "https://example.test").expect("pair");
338        assert!(manager.revoke(session.token()));
339        assert!(matches!(manager.validate(session.token(), "https://example.test"), Err(WebmcpError::Unauthorized)));
340    }
341
342    #[test]
343    fn replacing_pairing_revokes_sessions_and_preserves_origin_validation() {
344        let manager = PairingManager::new(["https://example.test"], Duration::from_secs(60)).expect("manager");
345        let display = manager.begin_pairing();
346        let session = manager.pair(display.code(), "https://example.test").expect("pair");
347
348        let replacement = manager
349            .replace_pairing_for_origin("https://example.test")
350            .expect("replacement pairing");
351        assert!(matches!(manager.validate(session.token(), "https://example.test"), Err(WebmcpError::Unauthorized)));
352        assert!(matches!(manager.pair(display.code(), "https://example.test"), Err(WebmcpError::PairingExpired)));
353        assert!(manager.pair(replacement.code(), "https://example.test").is_ok());
354    }
355
356    #[test]
357    fn rejected_replacement_does_not_revoke_current_session() {
358        let manager = PairingManager::new(["https://example.test"], Duration::from_secs(60)).expect("manager");
359        let display = manager.begin_pairing();
360        let session = manager.pair(display.code(), "https://example.test").expect("pair");
361
362        assert!(matches!(
363            manager.replace_pairing_for_origin("https://other.test"),
364            Err(WebmcpError::OriginRejected(_))
365        ));
366        assert!(manager.validate(session.token(), "https://example.test").is_ok());
367    }
368
369    #[test]
370    fn reconnect_can_resume_without_reusing_the_pairing_code() {
371        let manager = PairingManager::new(["https://example.test"], Duration::from_secs(60)).expect("manager");
372        let display = manager.begin_pairing();
373        let session = manager.pair(display.code(), "https://example.test").expect("pair");
374        let resumed = manager.resume(session.token(), "https://example.test").expect("resume");
375        assert_eq!(resumed.token(), session.token());
376        assert!(matches!(manager.pair(display.code(), "https://example.test"), Err(WebmcpError::PairingUsed)));
377    }
378
379    #[test]
380    fn wildcard_and_malformed_origins_are_rejected() {
381        for origin in [
382            "*",
383            "https://*",
384            " https://example.test ",
385            "localhost",
386            "ftp://example.test",
387            "https://example.test/path",
388            "https://user@example.test",
389        ] {
390            assert!(PairingManager::new([origin], Duration::from_secs(60)).is_err(), "accepted {origin}");
391        }
392    }
393
394    #[test]
395    fn repeated_invalid_codes_expire_the_pending_code() {
396        let manager = PairingManager::new(["https://example.test"], Duration::from_secs(60)).expect("manager");
397        let display = manager.begin_pairing();
398        for _ in 0..5 {
399            assert!(matches!(manager.pair("000000000000", "https://example.test"), Err(WebmcpError::PairingExpired)));
400        }
401        assert!(matches!(manager.pair(display.code(), "https://example.test"), Err(WebmcpError::PairingExpired)));
402    }
403
404    #[test]
405    fn refreshing_an_active_session_extends_only_its_inactivity_deadline() {
406        let manager = PairingManager::new(["https://example.test"], Duration::from_millis(400)).expect("manager");
407        let display = manager.begin_pairing();
408        let session = manager.pair(display.code(), "https://example.test").expect("pair");
409
410        std::thread::sleep(Duration::from_millis(150));
411        manager.refresh(session.token(), "https://example.test").expect("refresh");
412        std::thread::sleep(Duration::from_millis(150));
413        assert!(manager.validate(session.token(), "https://example.test").is_ok());
414
415        std::thread::sleep(Duration::from_millis(450));
416        assert!(matches!(manager.validate(session.token(), "https://example.test"), Err(WebmcpError::Unauthorized)));
417    }
418}
vtcode_webmcp/pairing.rs

vtcode_webmcp/
pairing.rs
