vcon_server/mcp/
gate.rs

1//! The dispatch gate's rules, each a plain function or type a test can
2//! drive: the scope a tool needs, the per-key rate, the concurrency cap and
3//! the refusal sent when one says no.
4
5use std::collections::HashMap;
6use std::sync::{Arc, Mutex};
7use std::time::{Duration, Instant};
8
9use rmcp::ErrorData;
10use rmcp::model::{ErrorCode, ToolAnnotations};
11use tokio::sync::{OwnedSemaphorePermit, Semaphore};
12
13use super::Limits;
14use crate::access::Scope;
15use crate::api::Problem;
16use crate::ids::KeyId;
17
18/// The scope a tool needs, from its name's area and its hints: `vcon_`
19/// tools need read when read-only, delete when destructive, else write;
20/// `legal_` tools legal; `audit_` tools audit; `admin_` tools admin. None,
21/// which refuses every call, for a tool that does not state both hints,
22/// states both true, or has no known area.
23#[must_use]
24pub fn scope(name: &str, hints: Option<&ToolAnnotations>) -> Option<Scope> {
25    let hints = hints?;
26    let (read_only, destructive) = (hints.read_only_hint?, hints.destructive_hint?);
27    if read_only && destructive {
28        return None;
29    }
30    match name.split_once('_')?.0 {
31        "vcon" if read_only => Some(Scope::Read),
32        "vcon" if destructive => Some(Scope::Delete),
33        "vcon" => Some(Scope::Write),
34        "legal" => Some(Scope::Legal),
35        "audit" => Some(Scope::Audit),
36        "admin" => Some(Scope::Admin),
37        _ => None,
38    }
39}
40
41/// The length of one rate window.
42const WINDOW: Duration = Duration::from_secs(1);
43
44/// Past this many keys in the table, windows that have ended are dropped.
45const KEPT: usize = 1024;
46
47/// Each key's calls in its current one-second window.
48#[derive(Debug, Default)]
49pub(super) struct Windows(Mutex<HashMap<KeyId, (Instant, u32)>>);
50
51impl Windows {
52    /// Counts one call by `key` at `now` against `limit`, or says how long
53    /// until its window ends when the window is full. A key's window starts
54    /// at its first call and lasts one second.
55    pub(super) fn take(&self, key: KeyId, limit: u32, now: Instant) -> Result<(), Duration> {
56        let mut windows = self.0.lock().expect("rate windows");
57        if windows.len() > KEPT {
58            windows.retain(|_, (start, _)| now.saturating_duration_since(*start) < WINDOW);
59        }
60        let (start, calls) = windows.entry(key).or_insert((now, 0));
61        if now.saturating_duration_since(*start) >= WINDOW {
62            *start = now;
63            *calls = 0;
64        }
65        if *calls >= limit {
66            return Err((*start + WINDOW).saturating_duration_since(now));
67        }
68        *calls += 1;
69        Ok(())
70    }
71}
72
73/// The server's state the gate checks calls against.
74#[derive(Debug)]
75pub(super) struct Gate {
76    pub(super) limits: Limits,
77    pub(super) windows: Windows,
78    slots: Arc<Semaphore>,
79}
80
81impl Gate {
82    pub(super) fn new(limits: Limits) -> Gate {
83        Gate {
84            limits,
85            windows: Windows::default(),
86            slots: Arc::new(Semaphore::new(limits.concurrent as usize)),
87        }
88    }
89
90    /// A slot for one call, held until the permit drops; None at once when
91    /// every slot is taken, never a wait.
92    pub(super) fn slot(&self) -> Option<OwnedSemaphorePermit> {
93        self.slots.clone().try_acquire_owned().ok()
94    }
95}
96
97/// The JSON-RPC error a refusal is sent as: its code by kind, the problem's
98/// title as its message, and its problem document as its data, with
99/// `retry_after_ms` when the wait is known.
100pub(super) fn refused(problem: Problem, retry: Option<Duration>) -> ErrorData {
101    let code = match problem {
102        Problem::UnknownTool | Problem::InvalidArguments => ErrorCode::INVALID_PARAMS,
103        Problem::Forbidden => ErrorCode(-32001),
104        _ => ErrorCode(-32000),
105    };
106    let mut data = problem.document();
107    if let Some(retry) = retry {
108        // At least 1, so a wait shorter than a millisecond is not read as
109        // none.
110        let millis = u64::try_from(retry.as_millis()).unwrap_or(u64::MAX).max(1);
111        data["retry_after_ms"] = millis.into();
112    }
113    ErrorData::new(code, problem.name().1, Some(data))
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    fn hints(read_only: Option<bool>, destructive: Option<bool>) -> ToolAnnotations {
121        let mut hints = ToolAnnotations::new();
122        hints.read_only_hint = read_only;
123        hints.destructive_hint = destructive;
124        hints
125    }
126
127    #[test]
128    fn the_scope_comes_from_the_area_and_the_hints() {
129        let (t, f) = (Some(true), Some(false));
130        for (name, read_only, destructive, want) in [
131            ("vcon_get", t, f, Some(Scope::Read)),
132            ("vcon_delete", f, t, Some(Scope::Delete)),
133            ("vcon_store", f, f, Some(Scope::Write)),
134            ("legal_get", t, f, Some(Scope::Legal)),
135            ("legal_restore", f, f, Some(Scope::Legal)),
136            ("legal_hold", f, t, Some(Scope::Legal)),
137            ("audit_read", t, f, Some(Scope::Audit)),
138            ("admin_tenant_list", t, f, Some(Scope::Admin)),
139            ("admin_tenant_remove", f, t, Some(Scope::Admin)),
140            // Both hints, and only consistent ones, are needed.
141            ("vcon_get", None, f, None),
142            ("vcon_get", t, None, None),
143            ("vcon_get", t, t, None),
144            ("admin_tenant_list", None, None, None),
145            // An area no scope belongs to.
146            ("tenant_list", t, f, None),
147            ("vcons", t, f, None),
148            ("", t, f, None),
149        ] {
150            assert_eq!(
151                scope(name, Some(&hints(read_only, destructive))),
152                want,
153                "{name} {read_only:?} {destructive:?}"
154            );
155        }
156        assert_eq!(scope("vcon_get", None), None, "no hints at all");
157    }
158
159    #[test]
160    fn a_key_has_its_calls_per_window_and_then_waits_for_its_end() {
161        let windows = Windows::default();
162        let (a, b) = (KeyId([1; 8]), KeyId([2; 8]));
163        let start = Instant::now();
164        assert_eq!(windows.take(a, 2, start), Ok(()));
165        let later = start + Duration::from_millis(300);
166        assert_eq!(windows.take(a, 2, later), Ok(()));
167        assert_eq!(
168            windows.take(a, 2, later),
169            Err(Duration::from_millis(700)),
170            "full until the window ends"
171        );
172        assert_eq!(
173            windows.take(b, 2, later),
174            Ok(()),
175            "another key's own window"
176        );
177        let next = start + WINDOW;
178        assert_eq!(windows.take(a, 2, next), Ok(()), "a new window");
179        assert_eq!(windows.take(a, 2, next), Ok(()));
180        assert!(windows.take(a, 2, next).is_err());
181    }
182
183    #[test]
184    fn ended_windows_are_dropped_once_the_table_is_large() {
185        let windows = Windows::default();
186        let start = Instant::now();
187        for n in 0..=KEPT as u64 {
188            let key = KeyId(n.to_be_bytes());
189            windows.take(key, 1, start).expect("a first call");
190        }
191        let later = start + WINDOW;
192        windows.take(KeyId([0xff; 8]), 1, later).expect("a call");
193        assert_eq!(windows.0.lock().expect("windows").len(), 1);
194    }
195
196    #[test]
197    fn a_full_server_refuses_a_slot_at_once_and_frees_it_on_drop() {
198        let gate = Gate::new(Limits {
199            calls_per_second: 1,
200            concurrent: 2,
201            ..Limits::DEFAULT
202        });
203        let first = gate.slot().expect("a slot");
204        let second = gate.slot().expect("a slot");
205        assert!(gate.slot().is_none(), "no third slot");
206        drop(first);
207        assert!(gate.slot().is_some(), "freed");
208        drop(second);
209    }
210
211    #[test]
212    fn a_refusal_carries_its_problem_document_and_its_code() {
213        for (problem, code) in [
214            (Problem::UnknownTool, -32602),
215            (Problem::InvalidArguments, -32602),
216            (Problem::Forbidden, -32001),
217            (Problem::RateLimited, -32000),
218            (Problem::AtCapacity, -32000),
219            (Problem::AuditFailed, -32000),
220        ] {
221            let error = refused(problem, None);
222            assert_eq!(error.code.0, code, "{problem:?}");
223            assert_eq!(error.data, Some(problem.document()), "{problem:?}");
224            assert_eq!(error.message, problem.name().1);
225        }
226        let error = refused(Problem::RateLimited, Some(Duration::from_micros(10)));
227        let data = error.data.expect("data");
228        assert_eq!(data["retry_after_ms"], 1, "never 0");
229        assert_eq!(data["status"], 429);
230    }
231}
vcon_server/mcp/gate.rs

vcon_server/mcp/
gate.rs
