Skip to main content

forge_sandbox/
ops.rs

1//! deno_core op definitions for the Forge sandbox.
2//!
3//! The `#[op2]` macro generates additional public items (v8 function pointers,
4//! metadata structs) that cannot carry doc comments. We suppress `missing_docs`
5//! at the module level — all actual functions and types are documented below.
6#![allow(missing_docs)]
7
8use std::cell::RefCell;
9use std::rc::Rc;
10use std::sync::Arc;
11
12use deno_core::op2;
13use deno_core::OpState;
14use deno_error::JsErrorBox;
15
16use std::collections::HashSet;
17
18use crate::stash::validate_key;
19use crate::{ResourceDispatcher, StashDispatcher, ToolDispatcher};
20
21/// Rate limiting state for tool calls within a single execution.
22pub(crate) struct ToolCallLimits {
23    /// Maximum number of tool calls allowed.
24    pub(crate) max_calls: usize,
25    /// Maximum size of serialized arguments per call.
26    pub(crate) max_args_size: usize,
27    /// Number of tool calls made so far.
28    pub(crate) calls_made: usize,
29}
30
31/// Rate limiting state for stash operations within a single execution.
32pub(crate) struct StashCallLimits {
33    /// Maximum number of stash operations allowed (None = unlimited).
34    pub(crate) max_calls: Option<usize>,
35    /// Number of stash operations made so far.
36    pub(crate) calls_made: usize,
37}
38
39impl StashCallLimits {
40    /// Check if the limit has been reached. Returns an error message if so.
41    pub(crate) fn check_limit(&mut self) -> Result<(), String> {
42        if let Some(max) = self.max_calls {
43            if self.calls_made >= max {
44                return Err(format!(
45                    "stash operation limit reached ({max} calls per execution)"
46                ));
47            }
48        }
49        self.calls_made += 1;
50        Ok(())
51    }
52}
53
54/// Log a message from sandbox code.
55#[op2(fast)]
56pub fn op_forge_log(#[string] msg: &str) {
57    tracing::info!(target: "forge::sandbox::js", "{}", msg);
58}
59
60/// Store the execution result in OpState.
61#[op2(fast)]
62pub fn op_forge_set_result(state: &mut OpState, #[string] json: &str) {
63    state.put(ExecutionResult(json.to_string()));
64}
65
66/// Call a tool on a downstream server via the ToolDispatcher.
67///
68/// Enforces per-execution rate limiting and argument size limits via
69/// [`ToolCallLimits`] stored in OpState.
70#[op2]
71#[string]
72pub async fn op_forge_call_tool(
73    op_state: Rc<RefCell<OpState>>,
74    #[string] server: String,
75    #[string] tool: String,
76    #[string] args_json: String,
77) -> Result<String, JsErrorBox> {
78    tracing::debug!(
79        server = %server,
80        tool = %tool,
81        args_len = args_json.len(),
82        "tool call dispatched"
83    );
84
85    // Check and increment tool call limits
86    {
87        let mut st = op_state.borrow_mut();
88        let limits = st.borrow_mut::<ToolCallLimits>();
89        if limits.calls_made >= limits.max_calls {
90            return Err(JsErrorBox::generic(format!(
91                "tool call limit exceeded (max {} calls per execution)",
92                limits.max_calls
93            )));
94        }
95        if args_json.len() > limits.max_args_size {
96            return Err(JsErrorBox::generic(format!(
97                "tool call args too large ({} bytes, max {} bytes)",
98                args_json.len(),
99                limits.max_args_size
100            )));
101        }
102        limits.calls_made += 1;
103    }
104
105    let dispatcher = {
106        let st = op_state.borrow();
107        st.borrow::<Arc<dyn ToolDispatcher>>().clone()
108    };
109
110    let args: serde_json::Value = serde_json::from_str(&args_json)
111        .map_err(|e| JsErrorBox::generic(format!("invalid JSON args: {e}")))?;
112
113    let result = match dispatcher.call_tool(&server, &tool, args).await {
114        Ok(val) => val,
115        Err(e) => {
116            let known = {
117                let st = op_state.borrow();
118                st.try_borrow::<KnownTools>()
119                    .map(|kt| kt.0.clone())
120                    .unwrap_or_default()
121            };
122            let pairs: Vec<(&str, &str)> = known
123                .iter()
124                .map(|(s, t)| (s.as_str(), t.as_str()))
125                .collect();
126            let mut structured = e.to_structured_error(Some(&pairs));
127            crate::redact::redact_structured_error(&server, &tool, &mut structured);
128            return serde_json::to_string(&structured)
129                .map_err(|e| JsErrorBox::generic(format!("error serialization failed: {e}")));
130        }
131    };
132
133    serde_json::to_string(&result)
134        .map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
135}
136
137/// Wrapper for execution results stored in OpState.
138pub(crate) struct ExecutionResult(pub(crate) String);
139
140/// Wrapper for the maximum resource content size, stored in OpState.
141pub(crate) struct MaxResourceSize(pub(crate) usize);
142
143/// Wrapper for the current server group, stored in OpState.
144///
145/// Used by stash ops to enforce group isolation. `None` means ungrouped.
146pub(crate) struct CurrentGroup(pub(crate) Option<String>);
147
148/// Set of known server names for SR-R6 validation.
149///
150/// Stored in OpState so `op_forge_read_resource` can reject unknown servers
151/// before any dispatch machinery runs.
152pub(crate) struct KnownServers(pub(crate) HashSet<String>);
153
154/// Known (server, tool) pairs for structured error fuzzy matching.
155///
156/// Stored in OpState so tool/resource dispatch errors can produce
157/// `{error:true, code:"TOOL_NOT_FOUND", suggested_fix:"Did you mean 'find_symbols'?"}`.
158pub(crate) struct KnownTools(pub(crate) Vec<(String, String)>);
159
160/// Validate a resource URI for security.
161///
162/// Rejects:
163/// - URIs with path traversal (`..` as a path segment, including URL-encoded variants)
164/// - URIs longer than 2048 bytes
165/// - URIs containing null bytes
166/// - URIs containing control characters (U+0000..U+001F, U+007F)
167pub(crate) fn validate_resource_uri(uri: &str) -> Result<(), String> {
168    if uri.len() > 2048 {
169        return Err(format!(
170            "resource URI too long ({} bytes, max 2048 bytes)",
171            uri.len()
172        ));
173    }
174    if uri.bytes().any(|b| b == 0) {
175        return Err("resource URI must not contain null bytes".into());
176    }
177    if uri.chars().any(|c| c.is_control()) {
178        return Err("resource URI must not contain control characters".into());
179    }
180    if has_path_traversal(uri) {
181        return Err("resource URI must not contain path traversal".into());
182    }
183    if let Some(scheme) = extract_uri_scheme(uri) {
184        if is_blocked_scheme(&scheme) {
185            return Err(format!("URI scheme '{}' is not allowed", scheme));
186        }
187    }
188    Ok(())
189}
190
191/// Blocked URI schemes that should never be passed to resource dispatchers.
192const BLOCKED_SCHEMES: &[&str] = &[
193    "data",
194    "javascript",
195    "ftp",
196    "gopher",
197    "telnet",
198    "ldap",
199    "dict",
200];
201
202/// Extract the scheme from a URI (text before `://` or `:`).
203/// Returns `None` for schemeless URIs like `some-resource-id`.
204fn extract_uri_scheme(uri: &str) -> Option<String> {
205    // Check for `://` first (most common)
206    if let Some(pos) = uri.find("://") {
207        let candidate = &uri[..pos];
208        // A valid scheme contains only alphanumeric, +, -, . and starts with a letter
209        if !candidate.is_empty()
210            && candidate.as_bytes()[0].is_ascii_alphabetic()
211            && candidate
212                .bytes()
213                .all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'-' || b == b'.')
214        {
215            return Some(candidate.to_ascii_lowercase());
216        }
217    }
218    // Check for `:` without `//` (e.g., `data:text/plain,...`, `javascript:...`)
219    if let Some(pos) = uri.find(':') {
220        let candidate = &uri[..pos];
221        if !candidate.is_empty()
222            && candidate.as_bytes()[0].is_ascii_alphabetic()
223            && candidate
224                .bytes()
225                .all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'-' || b == b'.')
226        {
227            return Some(candidate.to_ascii_lowercase());
228        }
229    }
230    None
231}
232
233/// Check if a scheme is in the blocklist (case-insensitive, scheme already lowercased).
234fn is_blocked_scheme(scheme: &str) -> bool {
235    BLOCKED_SCHEMES.contains(&scheme)
236}
237
238/// Check if a URI contains path traversal (`..`) as a path segment.
239///
240/// Checks the raw URI, then percent-decodes and checks again (to catch `%2e%2e`),
241/// then double-decodes and checks again (to catch `%252e%252e`).
242fn has_path_traversal(uri: &str) -> bool {
243    if has_dotdot_segment(uri) {
244        return true;
245    }
246    // Decode once and check
247    let decoded = percent_encoding::percent_decode_str(uri).decode_utf8_lossy();
248    if has_dotdot_segment(&decoded) {
249        return true;
250    }
251    // Double-decode and check (catches %252e%252e → %2e%2e → ..)
252    let double_decoded = percent_encoding::percent_decode_str(&decoded).decode_utf8_lossy();
253    if double_decoded != decoded && has_dotdot_segment(&double_decoded) {
254        return true;
255    }
256    false
257}
258
259/// Check if a string contains `..` as a path segment (not as part of a filename).
260///
261/// Matches: `..`, `../`, `/../`, `/..` at end, bare `..`
262fn has_dotdot_segment(s: &str) -> bool {
263    // Exact match
264    if s == ".." {
265        return true;
266    }
267    // Starts with ../
268    if s.starts_with("../") {
269        return true;
270    }
271    // Ends with /..
272    if s.ends_with("/..") {
273        return true;
274    }
275    // Contains /../ anywhere
276    if s.contains("/../") {
277        return true;
278    }
279    false
280}
281
282/// Read a resource by URI from a downstream server via the ResourceDispatcher.
283///
284/// Enforces URI validation (SR-R1), per-execution rate limiting (SR-R3),
285/// max resource size truncation (SR-R2), and error redaction (SR-R5).
286#[op2]
287#[string]
288pub async fn op_forge_read_resource(
289    op_state: Rc<RefCell<OpState>>,
290    #[string] server: String,
291    #[string] uri: String,
292) -> Result<String, JsErrorBox> {
293    tracing::debug!(
294        server = %server,
295        uri = %uri,
296        "resource read dispatched"
297    );
298
299    // SR-R1: Validate URI
300    validate_resource_uri(&uri).map_err(JsErrorBox::generic)?;
301
302    // SR-R6: Reject unknown server names before dispatch
303    {
304        let st = op_state.borrow();
305        if let Some(known) = st.try_borrow::<KnownServers>() {
306            if !known.0.contains(&server) {
307                return Err(JsErrorBox::generic(format!("unknown server: '{server}'")));
308            }
309        }
310    }
311
312    // SR-R3: Check and increment tool call limits (shared with tool calls)
313    {
314        let mut st = op_state.borrow_mut();
315        let limits = st.borrow_mut::<ToolCallLimits>();
316        if limits.calls_made >= limits.max_calls {
317            return Err(JsErrorBox::generic(format!(
318                "tool call limit exceeded (max {} calls per execution)",
319                limits.max_calls
320            )));
321        }
322        limits.calls_made += 1;
323    }
324
325    // Get dispatcher and max_resource_size from OpState
326    let (dispatcher, max_resource_size) = {
327        let st = op_state.borrow();
328        let d = st.borrow::<Arc<dyn ResourceDispatcher>>().clone();
329        let max_size = st
330            .try_borrow::<MaxResourceSize>()
331            .map(|m| m.0)
332            .unwrap_or(64 * 1024 * 1024); // 64 MB default
333        (d, max_size)
334    };
335
336    let result = match dispatcher.read_resource(&server, &uri).await {
337        Ok(val) => val,
338        Err(e) => {
339            let known = {
340                let st = op_state.borrow();
341                st.try_borrow::<KnownTools>()
342                    .map(|kt| kt.0.clone())
343                    .unwrap_or_default()
344            };
345            let pairs: Vec<(&str, &str)> = known
346                .iter()
347                .map(|(s, t)| (s.as_str(), t.as_str()))
348                .collect();
349            let mut structured = e.to_structured_error(Some(&pairs));
350            crate::redact::redact_structured_error(&server, "readResource", &mut structured);
351            return serde_json::to_string(&structured)
352                .map_err(|e| JsErrorBox::generic(format!("error serialization failed: {e}")));
353        }
354    };
355
356    // SR-R2: Serialize and truncate if > max_resource_size
357    let mut json = serde_json::to_string(&result)
358        .map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))?;
359
360    if json.len() > max_resource_size {
361        // Truncate to max_resource_size and ensure valid UTF-8 at boundary
362        let truncated = &json[..max_resource_size];
363        // Find the last valid UTF-8 char boundary
364        let end = truncated
365            .char_indices()
366            .last()
367            .map(|(i, c)| i + c.len_utf8())
368            .unwrap_or(0);
369        json = json[..end].to_string();
370    }
371
372    Ok(json)
373}
374
375/// Store a value in the session stash via the StashDispatcher.
376///
377/// - `key`: The stash key (alphanumeric + `_-.:`).
378/// - `value_json`: JSON-serialized value to store.
379/// - `ttl_secs`: TTL in seconds; 0 means use server default.
380#[op2]
381#[string]
382pub async fn op_forge_stash_put(
383    op_state: Rc<RefCell<OpState>>,
384    #[string] key: String,
385    #[string] value_json: String,
386    #[smi] ttl_secs: u32,
387) -> Result<String, JsErrorBox> {
388    // Defense-in-depth: validate key at op boundary before IPC traversal
389    validate_key(&key).map_err(|e| JsErrorBox::generic(e.to_string()))?;
390
391    // Check stash operation limit
392    if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
393        limits.check_limit().map_err(JsErrorBox::generic)?;
394    }
395
396    let (dispatcher, current_group) = {
397        let st = op_state.borrow();
398        let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
399        let g = st.borrow::<CurrentGroup>().0.clone();
400        (d, g)
401    };
402
403    let value: serde_json::Value = serde_json::from_str(&value_json)
404        .map_err(|e| JsErrorBox::generic(format!("invalid JSON value: {e}")))?;
405
406    let ttl = if ttl_secs == 0 { None } else { Some(ttl_secs) };
407
408    let result = dispatcher
409        .put(&key, value, ttl, current_group)
410        .await
411        .map_err(|e| JsErrorBox::generic(e.to_string()))?;
412
413    serde_json::to_string(&result)
414        .map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
415}
416
417/// Retrieve a value from the session stash via the StashDispatcher.
418#[op2]
419#[string]
420pub async fn op_forge_stash_get(
421    op_state: Rc<RefCell<OpState>>,
422    #[string] key: String,
423) -> Result<String, JsErrorBox> {
424    validate_key(&key).map_err(|e| JsErrorBox::generic(e.to_string()))?;
425
426    if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
427        limits.check_limit().map_err(JsErrorBox::generic)?;
428    }
429
430    let (dispatcher, current_group) = {
431        let st = op_state.borrow();
432        let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
433        let g = st.borrow::<CurrentGroup>().0.clone();
434        (d, g)
435    };
436
437    let result = dispatcher
438        .get(&key, current_group)
439        .await
440        .map_err(|e| JsErrorBox::generic(e.to_string()))?;
441
442    serde_json::to_string(&result)
443        .map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
444}
445
446/// Delete an entry from the session stash via the StashDispatcher.
447#[op2]
448#[string]
449pub async fn op_forge_stash_delete(
450    op_state: Rc<RefCell<OpState>>,
451    #[string] key: String,
452) -> Result<String, JsErrorBox> {
453    validate_key(&key).map_err(|e| JsErrorBox::generic(e.to_string()))?;
454
455    if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
456        limits.check_limit().map_err(JsErrorBox::generic)?;
457    }
458
459    let (dispatcher, current_group) = {
460        let st = op_state.borrow();
461        let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
462        let g = st.borrow::<CurrentGroup>().0.clone();
463        (d, g)
464    };
465
466    let result = dispatcher
467        .delete(&key, current_group)
468        .await
469        .map_err(|e| JsErrorBox::generic(e.to_string()))?;
470
471    serde_json::to_string(&result)
472        .map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
473}
474
475/// List all keys visible to the current group from the session stash.
476#[op2]
477#[string]
478pub async fn op_forge_stash_keys(op_state: Rc<RefCell<OpState>>) -> Result<String, JsErrorBox> {
479    if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
480        limits.check_limit().map_err(JsErrorBox::generic)?;
481    }
482
483    let (dispatcher, current_group) = {
484        let st = op_state.borrow();
485        let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
486        let g = st.borrow::<CurrentGroup>().0.clone();
487        (d, g)
488    };
489
490    let result = dispatcher
491        .keys(current_group)
492        .await
493        .map_err(|e| JsErrorBox::generic(e.to_string()))?;
494
495    serde_json::to_string(&result)
496        .map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
497}
498
499deno_core::extension!(
500    forge_ext,
501    ops = [
502        op_forge_log,
503        op_forge_set_result,
504        op_forge_call_tool,
505        op_forge_read_resource,
506        op_forge_stash_put,
507        op_forge_stash_get,
508        op_forge_stash_delete,
509        op_forge_stash_keys
510    ],
511);
512
513#[cfg(test)]
514mod tests {
515    use super::*;
516
517    // --- SK-V01: stash key validation rejects control characters ---
518    #[test]
519    fn sk_v01_stash_key_rejects_control_chars() {
520        let err = validate_key("key\x01value").unwrap_err();
521        assert!(
522            matches!(err, crate::stash::StashError::InvalidKey),
523            "expected InvalidKey, got: {err}"
524        );
525    }
526
527    // --- SK-V02: stash key validation rejects path separators ---
528    #[test]
529    fn sk_v02_stash_key_rejects_path_separators() {
530        assert!(validate_key("key/value").is_err());
531        assert!(validate_key("key\\value").is_err());
532        assert!(validate_key("../etc/passwd").is_err());
533    }
534
535    // --- SK-V03: stash key validation rejects empty and oversized keys ---
536    #[test]
537    fn sk_v03_stash_key_rejects_empty_and_oversized() {
538        assert!(validate_key("").is_err());
539        let long_key = "a".repeat(257);
540        let err = validate_key(&long_key).unwrap_err();
541        assert!(
542            matches!(err, crate::stash::StashError::KeyTooLong { len: 257 }),
543            "expected KeyTooLong, got: {err}"
544        );
545    }
546
547    // --- SK-V04: stash key validation accepts valid patterns ---
548    #[test]
549    fn sk_v04_stash_key_accepts_valid_patterns() {
550        assert!(validate_key("simple-key").is_ok());
551        assert!(validate_key("key_with.dots:colons").is_ok());
552        assert!(validate_key("CamelCase123").is_ok());
553        assert!(validate_key("a").is_ok());
554        let max_key = "x".repeat(256);
555        assert!(validate_key(&max_key).is_ok());
556    }
557
558    // --- SK-V05: stash key validation rejects unicode ---
559    #[test]
560    fn sk_v05_stash_key_rejects_unicode() {
561        assert!(validate_key("key\u{0000}null").is_err());
562        assert!(validate_key("key with space").is_err());
563        assert!(validate_key("emoji\u{1F600}").is_err());
564    }
565
566    // --- RS-U04: reject URIs with path traversal (..) ---
567    #[test]
568    fn rs_u04_rejects_uri_with_path_traversal() {
569        assert!(validate_resource_uri("file:///logs/../../../etc/passwd").is_err());
570        assert!(validate_resource_uri("file:///..").is_err());
571        assert!(validate_resource_uri("..").is_err());
572        assert!(validate_resource_uri("a/../../b").is_err());
573        // Valid URI without traversal should pass
574        assert!(validate_resource_uri("file:///logs/app.log").is_ok());
575        assert!(validate_resource_uri("postgres://db/table").is_ok());
576    }
577
578    // --- URI-V01: legitimate double dots in filenames are allowed ---
579    #[test]
580    fn uri_v01_allows_legitimate_double_dots() {
581        // Double dots as part of a filename (not a path segment) should be allowed
582        assert!(validate_resource_uri("file:///v2..backup").is_ok());
583        assert!(validate_resource_uri("file:///data..2024.csv").is_ok());
584        assert!(validate_resource_uri("file:///config..old").is_ok());
585    }
586
587    // --- URI-V02: URL-encoded traversal (%2e%2e) is blocked ---
588    #[test]
589    fn uri_v02_blocks_url_encoded_traversal() {
590        // %2e = '.', so %2e%2e/%2e%2e = ../../..
591        assert!(validate_resource_uri("file:///logs/%2e%2e/%2e%2e/etc/passwd").is_err());
592        assert!(validate_resource_uri("file:///%2e%2e/secret").is_err());
593    }
594
595    // --- URI-V03: double-encoded traversal (%252e%252e) is blocked ---
596    #[test]
597    fn uri_v03_blocks_double_encoded_traversal() {
598        // %252e decodes to %2e, which decodes to '.'
599        assert!(validate_resource_uri("file:///logs/%252e%252e/%252e%252e/etc/passwd").is_err());
600    }
601
602    // --- URI-V04: mixed-case encoded traversal is blocked ---
603    #[test]
604    fn uri_v04_blocks_mixed_case_encoded_traversal() {
605        assert!(validate_resource_uri("file:///logs/%2E%2E/secret").is_err());
606        assert!(validate_resource_uri("file:///logs/%2e%2E/secret").is_err());
607    }
608
609    // --- RS-U05: reject URIs longer than 2048 bytes ---
610    #[test]
611    fn rs_u05_rejects_uri_longer_than_2048_bytes() {
612        let long_uri = "x".repeat(2049);
613        let err = validate_resource_uri(&long_uri).unwrap_err();
614        assert!(err.contains("too long"), "should mention too long: {err}");
615
616        // Exactly 2048 should be OK
617        let ok_uri = "x".repeat(2048);
618        assert!(validate_resource_uri(&ok_uri).is_ok());
619    }
620
621    // --- RS-U06: reject URIs with null bytes ---
622    #[test]
623    fn rs_u06_rejects_uri_with_null_bytes() {
624        let uri = "file:///logs\0/app.log";
625        let err = validate_resource_uri(uri).unwrap_err();
626        assert!(err.contains("null"), "should mention null: {err}");
627    }
628
629    // --- RS-U07: reject URIs with control characters ---
630    #[test]
631    fn rs_u07_rejects_uri_with_control_characters() {
632        // SOH (0x01)
633        let err = validate_resource_uri("file:///logs\x01/app.log").unwrap_err();
634        assert!(err.contains("control"), "should mention control: {err}");
635
636        // Tab (0x09)
637        assert!(validate_resource_uri("file:///logs\t/app.log").is_err());
638
639        // Newline (0x0A)
640        assert!(validate_resource_uri("file:///logs\n/app.log").is_err());
641
642        // DEL (0x7F)
643        assert!(validate_resource_uri("file:///logs\x7f/app.log").is_err());
644    }
645
646    // --- RS-S04: path traversal attack variants ---
647    #[test]
648    fn rs_s04_path_traversal_attack_variants() {
649        // Classic traversal
650        assert!(validate_resource_uri("../../../etc/passwd").is_err());
651        // Encoded traversal
652        assert!(validate_resource_uri("file:///logs/%2e%2e/%2e%2e/etc/passwd").is_err());
653        // Double dots at start
654        assert!(validate_resource_uri("..").is_err());
655        // Double dots embedded
656        assert!(validate_resource_uri("file:///../").is_err());
657        // Traversal after normal path
658        assert!(validate_resource_uri("file:///a/b/../../../etc/shadow").is_err());
659    }
660
661    // --- M2: URI Scheme Validation Tests ---
662
663    #[test]
664    fn uri_m2_01_allows_http_scheme() {
665        assert!(validate_resource_uri("http://example.com/resource").is_ok());
666    }
667
668    #[test]
669    fn uri_m2_02_allows_https_scheme() {
670        assert!(validate_resource_uri("https://example.com/resource").is_ok());
671    }
672
673    #[test]
674    fn uri_m2_03_allows_file_scheme() {
675        assert!(validate_resource_uri("file:///logs/app.log").is_ok());
676    }
677
678    #[test]
679    fn uri_m2_04_rejects_data_scheme() {
680        let err = validate_resource_uri("data:text/plain;base64,SGVsbG8=").unwrap_err();
681        assert!(
682            err.contains("not allowed"),
683            "expected 'not allowed' in error: {err}"
684        );
685    }
686
687    #[test]
688    fn uri_m2_05_rejects_javascript_scheme() {
689        let err = validate_resource_uri("javascript:alert(1)").unwrap_err();
690        assert!(err.contains("not allowed"), "error: {err}");
691    }
692
693    #[test]
694    fn uri_m2_06_rejects_ftp_scheme() {
695        let err = validate_resource_uri("ftp://evil.com/malware").unwrap_err();
696        assert!(err.contains("not allowed"), "error: {err}");
697    }
698
699    #[test]
700    fn uri_m2_07_rejects_gopher_scheme() {
701        let err = validate_resource_uri("gopher://evil.com/0").unwrap_err();
702        assert!(err.contains("not allowed"), "error: {err}");
703    }
704
705    #[test]
706    fn uri_m2_08_allows_custom_mcp_scheme() {
707        // Custom MCP resource URIs should be allowed
708        assert!(validate_resource_uri("postgres://db/table").is_ok());
709        assert!(validate_resource_uri("redis://localhost:6379/0").is_ok());
710        assert!(validate_resource_uri("mongodb://host/db").is_ok());
711    }
712
713    #[test]
714    fn uri_m2_09_allows_schemeless_uri() {
715        // Bare resource identifiers without any scheme
716        assert!(validate_resource_uri("some-resource-id").is_ok());
717        assert!(validate_resource_uri("table_name").is_ok());
718        assert!(validate_resource_uri("logs/2024/app.log").is_ok());
719    }
720
721    #[test]
722    fn uri_m2_10_case_insensitive_scheme_check() {
723        // Mixed case should be blocked
724        assert!(validate_resource_uri("JAVASCRIPT:alert(1)").is_err());
725        assert!(validate_resource_uri("JavaScript:void(0)").is_err());
726        assert!(validate_resource_uri("DATA:text/plain,hello").is_err());
727        assert!(validate_resource_uri("FTP://evil.com/file").is_err());
728        assert!(validate_resource_uri("Gopher://host/0").is_err());
729        assert!(validate_resource_uri("TELNET://host:23").is_err());
730        assert!(validate_resource_uri("LDAP://host/dc=com").is_err());
731        assert!(validate_resource_uri("DICT://host/define").is_err());
732    }
733
734    // --- Phase 7: Stash call limits tests ---
735
736    #[test]
737    fn stash_l4_01_stash_calls_count_against_limit() {
738        let mut limits = StashCallLimits {
739            max_calls: Some(3),
740            calls_made: 0,
741        };
742        assert!(limits.check_limit().is_ok());
743        assert!(limits.check_limit().is_ok());
744        assert!(limits.check_limit().is_ok());
745        // 4th call should fail
746        assert!(limits.check_limit().is_err());
747    }
748
749    #[test]
750    fn stash_l4_02_stash_limit_rejection_message() {
751        let mut limits = StashCallLimits {
752            max_calls: Some(1),
753            calls_made: 0,
754        };
755        assert!(limits.check_limit().is_ok());
756        let err = limits.check_limit().unwrap_err();
757        assert!(err.contains("limit reached"), "should mention limit: {err}");
758        assert!(
759            err.contains("1 calls"),
760            "should mention the limit count: {err}"
761        );
762    }
763
764    #[test]
765    fn stash_l4_03_stash_limit_independent_of_tool_limit() {
766        // Stash limits track separately from tool call limits
767        let mut stash_limits = StashCallLimits {
768            max_calls: Some(5),
769            calls_made: 0,
770        };
771        let tool_limits = ToolCallLimits {
772            max_calls: 0, // tool calls exhausted
773            max_args_size: 1024,
774            calls_made: 0,
775        };
776        // Stash should still work even if tool limit is at 0
777        assert!(stash_limits.check_limit().is_ok());
778        let _ = tool_limits;
779    }
780
781    #[test]
782    fn stash_l4_04_stash_limit_configurable() {
783        // None means unlimited
784        let mut unlimited = StashCallLimits {
785            max_calls: None,
786            calls_made: 0,
787        };
788        for _ in 0..1000 {
789            assert!(unlimited.check_limit().is_ok());
790        }
791
792        // Some(N) means N calls max
793        let mut limited = StashCallLimits {
794            max_calls: Some(2),
795            calls_made: 0,
796        };
797        assert!(limited.check_limit().is_ok());
798        assert!(limited.check_limit().is_ok());
799        assert!(limited.check_limit().is_err());
800    }
801}