From 037d5be5d34c08ebcf181f990f44ce28142da80b Mon Sep 17 00:00:00 2001 From: Igor Kantor Date: Tue, 25 Aug 2026 15:14:17 -0400 Subject: [PATCH] make invalid states unrepresentable: ScopedPath, MemoryConfig, TruncationReason, NonZero caps/rates - Scope::check now returns ScopedPath, a newtype constructible only by scope validation; grep/find/cat require it, so skipping the check is a compile error rather than a security hole - MemoryConfig enum replaces (Option, Option): the impossible 'instructions without a memory dir' state is no longer constructible - TruncationReason enum replaces freeform truncation-reason strings; wire contract (byte_cap/line_cap/max_results) unchanged - PeerLimiter::new validates capacity/refill (guards the Duration::from_secs_f64 panic on non-finite input) and per_minute takes NonZeroU32; --initialize-rate-per-min is NonZeroU32 - max_results/max_bytes/max_lines are NonZeroUsize through args -> options -> tools, so zero caps are rejected at the JSON boundary 44 tests pass; clippy and fmt clean; wire smoke-tested via streamable HTTP (initialize, grep, cat truncation reason, zero-cap and out-of-scope rejection) --- src/args.rs | 17 ++++++------ src/cli.rs | 7 +++-- src/gate.rs | 14 +++++++--- src/limiter.rs | 51 +++++++++++++++++++++++++++--------- src/main.rs | 49 +++++++++++++++++----------------- src/scope.rs | 36 ++++++++++++++++++------- src/server.rs | 61 ++++++++++++++++++++++++++----------------- src/tools/cat.rs | 37 ++++++++++++++------------ src/tools/common.rs | 4 ++- src/tools/find.rs | 20 ++++++++------ src/tools/grep.rs | 55 ++++++++++++++++++++------------------ src/tools/mod.rs | 22 +++++++++++++--- src/tools/options.rs | 18 ++++++++----- src/tools/response.rs | 20 ++++++++++++-- 14 files changed, 265 insertions(+), 146 deletions(-) diff --git a/src/args.rs b/src/args.rs index a70752a..704c56e 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,14 +1,15 @@ use crate::tools::{DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, DEFAULT_MAX_RESULTS, OutputMode}; use rmcp::schemars::{self, JsonSchema}; use serde::{Deserialize, Serialize}; +use std::num::NonZeroUsize; -fn default_max_results() -> usize { +fn default_max_results() -> NonZeroUsize { DEFAULT_MAX_RESULTS } -fn default_max_bytes() -> usize { +fn default_max_bytes() -> NonZeroUsize { DEFAULT_MAX_BYTES } -fn default_max_lines() -> usize { +fn default_max_lines() -> NonZeroUsize { DEFAULT_MAX_LINES } const fn default_true() -> bool { @@ -64,7 +65,7 @@ pub struct GrepArgs { pub after_context: usize, #[serde(default = "default_max_results")] #[schemars(description = "Maximum number of results to return (default 100)")] - pub max_results: usize, + pub max_results: NonZeroUsize, #[serde(default)] #[schemars( description = "Case-insensitive search (default false). Equivalent to prefixing pattern with (?i)." @@ -88,7 +89,7 @@ pub struct GrepArgs { #[schemars( description = "Hard cap on total response size in bytes (default ~5 MiB). Truncates with a marker." )] - pub max_bytes: usize, + pub max_bytes: NonZeroUsize, #[serde(default)] #[schemars( description = "Output mode: 'files_with_matches' (default — list file paths only), 'content' (matching lines with line numbers), 'count' (per-file match tallies as path: N)." @@ -107,7 +108,7 @@ pub struct FindArgs { pub pattern: String, #[serde(default = "default_max_results")] #[schemars(description = "Maximum number of results to return (default 100)")] - pub max_results: usize, + pub max_results: NonZeroUsize, #[serde(default)] #[schemars(description = "Include hidden files and directories (default false)")] pub include_hidden: bool, @@ -133,10 +134,10 @@ pub struct CatArgs { pub offset: usize, #[serde(default = "default_max_lines")] #[schemars(description = "Maximum number of lines to return (default 2000)")] - pub max_lines: usize, + pub max_lines: NonZeroUsize, #[serde(default = "default_max_bytes")] #[schemars(description = "Maximum number of bytes to return (default ~5 MiB)")] - pub max_bytes: usize, + pub max_bytes: NonZeroUsize, } #[derive(Debug, Serialize, Deserialize, JsonSchema)] diff --git a/src/cli.rs b/src/cli.rs index 9c8f3ef..d7c254b 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,7 +1,10 @@ use clap::Parser; use std::net::SocketAddr; +use std::num::NonZeroU32; use std::path::PathBuf; +const DEFAULT_INITIALIZE_RATE: NonZeroU32 = NonZeroU32::new(12).unwrap(); + /// Command-line arguments for `code-mcp`. /// /// Parsed via clap. All fields are `pub(crate)` because they're only consumed @@ -42,8 +45,8 @@ pub struct Args { /// per-minute rate (token bucket of capacity = rate, refilling over /// 60s). When exhausted, new initializes from that peer get 429 + /// Retry-After. Existing-session traffic is unaffected. - #[arg(long, default_value_t = 12)] - pub(crate) initialize_rate_per_min: u32, + #[arg(long, default_value_t = DEFAULT_INITIALIZE_RATE)] + pub(crate) initialize_rate_per_min: NonZeroU32, /// Trust the rightmost entry of `X-Forwarded-For` as the peer IP /// instead of the TCP socket address. Assumes a single trusted proxy diff --git a/src/gate.rs b/src/gate.rs index 2643c9b..38e5bb6 100644 --- a/src/gate.rs +++ b/src/gate.rs @@ -117,6 +117,12 @@ mod tests { use axum::Router; use axum::routing::any; use std::net::Ipv4Addr; + use std::num::NonZeroU32; + + fn nz(n: u32) -> NonZeroU32 { + NonZeroU32::new(n).unwrap() + } + use tower::ServiceExt; fn dummy_addr() -> SocketAddr { @@ -148,7 +154,7 @@ mod tests { let ctx = Arc::new(GateCtx { sessions: Arc::new(LocalSessionManager::default()), max_sessions: 0, // would block POSTs - limiter: PeerLimiter::per_minute(1), + limiter: PeerLimiter::per_minute(nz(1)), trust_forwarded_for: false, activity: Arc::new(ActivityTracker::new()), }); @@ -166,7 +172,7 @@ mod tests { let ctx = Arc::new(GateCtx { sessions: Arc::new(LocalSessionManager::default()), max_sessions: 0, - limiter: PeerLimiter::per_minute(1), + limiter: PeerLimiter::per_minute(nz(1)), trust_forwarded_for: false, activity: Arc::new(ActivityTracker::new()), }); @@ -183,7 +189,7 @@ mod tests { let ctx = Arc::new(GateCtx { sessions: Arc::new(LocalSessionManager::default()), max_sessions: 0, - limiter: PeerLimiter::per_minute(100), + limiter: PeerLimiter::per_minute(nz(100)), trust_forwarded_for: false, activity: Arc::new(ActivityTracker::new()), }); @@ -201,7 +207,7 @@ mod tests { let ctx = Arc::new(GateCtx { sessions: Arc::new(LocalSessionManager::default()), max_sessions: 1000, - limiter: PeerLimiter::new(2.0, 0.001, 1024), // ~no refill + limiter: PeerLimiter::new(2.0, 0.001, 1024).unwrap(), // ~no refill trust_forwarded_for: false, activity: Arc::new(ActivityTracker::new()), }); diff --git a/src/limiter.rs b/src/limiter.rs index e42feff..da08cf4 100644 --- a/src/limiter.rs +++ b/src/limiter.rs @@ -1,7 +1,18 @@ use std::collections::HashMap; use std::net::IpAddr; +use std::num::NonZeroU32; use std::sync::Mutex; use std::time::{Duration, Instant}; +use thiserror::Error; + +/// Returned by [`PeerLimiter::new`] when `capacity` or `refill_per_sec` is +/// zero, negative, or NaN. `try_consume` divides by `refill_per_sec` and +/// feeds the quotient to `Duration::from_secs_f64`, which panics on +/// non-finite input — this error makes that state unconstructible instead +/// of merely `debug_assert`-ed. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[error("capacity and refill_per_sec must be positive and finite")] +pub struct InvalidRate; /// Per-peer token-bucket limiter. Each peer (keyed by `IpAddr`) gets a /// bucket that refills continuously toward `capacity` at `refill_per_sec`. @@ -22,20 +33,27 @@ struct Bucket { } impl PeerLimiter { - pub fn new(capacity: f64, refill_per_sec: f64, evict_threshold: usize) -> Self { - debug_assert!(capacity > 0.0 && refill_per_sec > 0.0); - Self { + pub fn new( + capacity: f64, + refill_per_sec: f64, + evict_threshold: usize, + ) -> Result { + // NaN fails both comparisons, so it lands here too. + if !(capacity > 0.0 && refill_per_sec > 0.0) { + return Err(InvalidRate); + } + Ok(Self { capacity, refill_per_sec, buckets: Mutex::new(HashMap::new()), evict_threshold, - } + }) } /// Convenience: a per-minute rate (capacity = `rate`, refill = `rate`/60s). - pub fn per_minute(rate: u32) -> Self { - let cap = f64::from(rate.max(1)); - Self::new(cap, cap / 60.0, 4096) + pub fn per_minute(rate: NonZeroU32) -> Self { + let cap = f64::from(rate.get()); + Self::new(cap, cap / 60.0, 4096).expect("derived from a NonZeroU32 rate") } /// Try to consume one token. Returns the time until the next token would @@ -91,7 +109,7 @@ mod tests { #[test] fn allows_within_capacity() { - let l = PeerLimiter::new(3.0, 0.001, 1024); // refill effectively zero + let l = PeerLimiter::new(3.0, 0.001, 1024).unwrap(); // refill effectively zero for _ in 0..3 { assert!(l.try_consume(ip(1)).is_ok()); } @@ -99,7 +117,7 @@ mod tests { #[test] fn blocks_after_exhaustion() { - let l = PeerLimiter::new(2.0, 0.001, 1024); + let l = PeerLimiter::new(2.0, 0.001, 1024).unwrap(); l.try_consume(ip(1)).unwrap(); l.try_consume(ip(1)).unwrap(); let err = l @@ -110,16 +128,25 @@ mod tests { #[test] fn separate_ips_get_separate_buckets() { - let l = PeerLimiter::new(1.0, 0.001, 1024); + let l = PeerLimiter::new(1.0, 0.001, 1024).unwrap(); assert!(l.try_consume(ip(1)).is_ok()); assert!(l.try_consume(ip(2)).is_ok()); assert!(l.try_consume(ip(1)).is_err()); } + #[test] + fn new_rejects_non_positive_rates() { + assert!(PeerLimiter::new(0.0, 1.0, 8).is_err()); + assert!(PeerLimiter::new(1.0, 0.0, 8).is_err()); + assert!(PeerLimiter::new(-1.0, 1.0, 8).is_err()); + assert!(PeerLimiter::new(1.0, f64::NAN, 8).is_err()); + assert!(PeerLimiter::new(1.0, 1.0, 8).is_ok()); + } + #[test] fn refills_over_time() { // 10 tokens/sec → one token per 100ms. - let l = PeerLimiter::new(1.0, 10.0, 1024); + let l = PeerLimiter::new(1.0, 10.0, 1024).unwrap(); l.try_consume(ip(1)).unwrap(); assert!(l.try_consume(ip(1)).is_err()); std::thread::sleep(Duration::from_millis(150)); @@ -130,7 +157,7 @@ mod tests { fn evicts_stale_entries_when_threshold_exceeded() { // Tiny threshold + fast refill so the prior entry is "stale" by the // time we hit the threshold. - let l = PeerLimiter::new(1.0, 1000.0, 4); + let l = PeerLimiter::new(1.0, 1000.0, 4).unwrap(); for n in 0..4 { l.try_consume(ip(n)).unwrap(); } diff --git a/src/main.rs b/src/main.rs index 6ea0b85..92770b6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -64,7 +64,7 @@ use crate::cli::Args; use crate::gate::{GateCtx, gate}; use crate::limiter::PeerLimiter; use crate::scope::Scope; -use crate::server::CodeMcpServer; +use crate::server::{CodeMcpServer, MemoryConfig}; use crate::error::AppError; @@ -79,29 +79,34 @@ async fn main() -> Result<(), AppError> { let args = Args::parse(); - // If a memory dir is configured, load /instructions.md once at startup. - // It's appended to the InitializeResult.instructions payload. - let extra_instructions = if let Some(dir) = args.memory_dir.as_ref() { - let path = dir.join("instructions.md"); - match tokio::fs::read_to_string(&path).await { - Ok(s) => { - tracing::info!(path = %path.display(), "loaded extra instructions"); - Some(s) - } - Err(e) if e.kind() == std::io::ErrorKind::NotFound => None, - Err(e) => { - tracing::warn!(path = %path.display(), error = %e, "could not read instructions.md"); - None + // If a memory dir is configured, load /instructions.md once at + // startup. It's appended to the InitializeResult.instructions payload, + // so the dir and its extra instructions travel together in one type. + let memory = match args.memory_dir.as_ref() { + Some(dir) => { + let path = dir.join("instructions.md"); + let extra = match tokio::fs::read_to_string(&path).await { + Ok(s) => { + tracing::info!(path = %path.display(), "loaded extra instructions"); + Some(s) + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => None, + Err(e) => { + tracing::warn!(path = %path.display(), error = %e, "could not read instructions.md"); + None + } + }; + MemoryConfig::Enabled { + dir: dir.clone(), + extra, } } - } else { - None + None => MemoryConfig::Disabled, }; let scope = Scope::new(args.project.clone())?; tracing::info!(root = %scope.root().display(), "project scope active"); - let memory_dir = args.memory_dir.clone(); let cancel = CancellationToken::new(); let session_manager = Arc::new(LocalSessionManager::default()); let sessions_for_gate = session_manager.clone(); @@ -118,13 +123,7 @@ async fn main() -> Result<(), AppError> { ..Default::default() }; let service = StreamableHttpService::new( - move || { - Ok(CodeMcpServer::new( - memory_dir.clone(), - extra_instructions.clone(), - scope.clone(), - )) - }, + move || Ok(CodeMcpServer::new(memory.clone(), scope.clone())), session_manager, config, ); @@ -138,7 +137,7 @@ async fn main() -> Result<(), AppError> { }); tracing::info!( max_sessions = args.max_sessions, - initialize_rate_per_min = args.initialize_rate_per_min, + initialize_rate_per_min = args.initialize_rate_per_min.get(), trust_forwarded_for = args.trust_forwarded_for, session_idle_timeout_secs = args.session_idle_timeout_secs, session_sweep_interval_secs = args.session_sweep_interval_secs, diff --git a/src/scope.rs b/src/scope.rs index 644be76..b783152 100644 --- a/src/scope.rs +++ b/src/scope.rs @@ -30,11 +30,12 @@ impl Scope { &self.root } - /// Validate that `input` is within the scope, returning its canonical - /// path. Symlinks in `input` are resolved before the containment - /// check, so a symlink inside the project that points outside it is - /// rejected. - pub fn check>(&self, input: P) -> Result { + /// Validate that `input` is within the scope, returning a [`ScopedPath`] + /// — the canonical path, with the in-scope proof carried by the type so + /// downstream tool entry points can require it. Symlinks in `input` are + /// resolved before the containment check, so a symlink inside the + /// project that points outside it is rejected. + pub fn check>(&self, input: P) -> Result { let input = input.as_ref(); let canon = input .canonicalize() @@ -46,7 +47,20 @@ impl Scope { self.root.display() ))); } - Ok(canon) + Ok(ScopedPath(canon)) + } +} + +/// A path proven — via [`Scope::check`] — to canonicalize inside a scope's +/// root. The only constructor is `Scope::check`, and tool entry points take +/// `&ScopedPath`, so a path that skipped scope validation cannot reach them: +/// the "forgot to call `check`" bug is a compile error, not a security hole. +#[derive(Debug, Clone)] +pub struct ScopedPath(PathBuf); + +impl AsRef for ScopedPath { + fn as_ref(&self) -> &Path { + &self.0 } } @@ -64,7 +78,11 @@ mod tests { fs::write(td.path().join("a.txt"), "x")?; let s = Scope::new(td.path())?; let canon = s.check(td.path().join("a.txt"))?; - assert!(canon.ends_with("a.txt"), "got {}", canon.display()); + assert!( + canon.as_ref().ends_with("a.txt"), + "got {}", + canon.as_ref().display() + ); Ok(()) } @@ -109,10 +127,10 @@ mod tests { // Either way, reading /etc/passwd via traversal should not succeed. Err(AppError::NotFound(_)) => Ok(()), Ok(p) => { - if p.starts_with(td.path()) { + if p.as_ref().starts_with(td.path()) { Ok(()) } else { - Err(format!("escaped root: {}", p.display()).into()) + Err(format!("escaped root: {}", p.as_ref().display()).into()) } } other => Err(format!("unexpected: {:?}", other).into()), diff --git a/src/server.rs b/src/server.rs index 7e58034..b24596a 100644 --- a/src/server.rs +++ b/src/server.rs @@ -12,8 +12,8 @@ use crate::memory::load_memory; use crate::scope::Scope; use crate::tools; -/// The MCP server handler. Owns the tool router, the optional memory dir, -/// the extra instructions loaded at startup, and the filesystem [`Scope`]. +/// The MCP server handler. Owns the tool router, the memory configuration +/// ([`MemoryConfig`]), and the filesystem [`Scope`]. /// /// Constructed once per session by the `StreamableHttpService` closure in /// `main` (so each session gets its own cheap `Clone` of the `Scope` and @@ -22,24 +22,35 @@ use crate::tools; #[derive(Clone)] pub struct CodeMcpServer { tool_router: rmcp::handler::server::router::tool::ToolRouter, - memory_dir: Option, - extra_instructions: Option, + memory: MemoryConfig, scope: Scope, } +/// Whether and how the memory subsystem is configured. +/// +/// This replaces the old `(Option, Option)` pair +/// (`memory_dir`, `extra_instructions`), which could represent the +/// impossible state "extra instructions present but no memory dir". +#[derive(Debug, Clone, Default)] +pub enum MemoryConfig { + /// `--memory-dir` not supplied: the `memories` tool errors and no + /// memory instructions are advertised in `get_info`. + #[default] + Disabled, + /// `--memory-dir `: the `memories` tool reads from `dir`, and the + /// optional contents of `/instructions.md` (loaded once at + /// startup) are appended to the initialize payload. + Enabled { dir: PathBuf, extra: Option }, +} + impl CodeMcpServer { - /// Construct a new server instance. `extra_instructions` is the contents - /// of `/instructions.md` (loaded once at startup) and is - /// appended to the `InitializeResult.instructions` payload. - pub fn new( - memory_dir: Option, - extra_instructions: Option, - scope: Scope, - ) -> Self { + /// Construct a new server instance. `memory` carries the configured + /// memory dir and the startup-loaded contents of + /// `/instructions.md` (see [`MemoryConfig::Enabled`]). + pub fn new(memory: MemoryConfig, scope: Scope) -> Self { Self { tool_router: Self::tool_router(), - memory_dir, - extra_instructions, + memory, scope, } } @@ -135,11 +146,13 @@ impl CodeMcpServer { &self, Parameters(args): Parameters, ) -> ToolResult { - let dir = match self.memory_dir.clone() { - Some(d) => d, - None => return Ok(tool_error(AppError::InvalidRequest( - "memory dir not configured; start server with --memory-dir ".into(), - ))), + let dir = match &self.memory { + MemoryConfig::Enabled { dir, .. } => dir.clone(), + MemoryConfig::Disabled => { + return Ok(tool_error(AppError::InvalidRequest( + "memory dir not configured; start server with --memory-dir ".into(), + ))); + } }; let res = tokio::task::spawn_blocking(move || load_memory(&dir, args.name.as_deref())) @@ -180,14 +193,11 @@ least one match. This is the most token-efficient mode for broad reconnaissance `find` matches the basename of each path by default. Set match_basename: false to \ match against the full path instead. -`.gitignore` files are respected by default for both grep and find. \ -Set respect_gitignore: false to walk the entire tree, including ignored paths. - `cat` supports an `offset` (0-based line number) for paginating long files, plus \ optional `max_lines` and `max_bytes` caps.", ); - if self.memory_dir.is_some() { + if matches!(self.memory, MemoryConfig::Enabled { .. }) { instructions.push_str( "\n\nThis server has a memory directory configured. \ Call the `memories` tool at the start of a session to load persisted context \ @@ -197,7 +207,10 @@ Individual memory files referenced in the index can also be read with `cat`.", ); } - if let Some(extra) = &self.extra_instructions { + if let MemoryConfig::Enabled { + extra: Some(extra), .. + } = &self.memory + { instructions.push_str("\n\n--- project instructions ---\n\n"); instructions.push_str(extra); } diff --git a/src/tools/cat.rs b/src/tools/cat.rs index 028d4b3..fb4ec65 100644 --- a/src/tools/cat.rs +++ b/src/tools/cat.rs @@ -1,10 +1,11 @@ //! The `cat` tool: read file contents with line/byte pagination. -use super::response::ToolResponse; +use super::response::{ToolResponse, TruncationReason}; use crate::error::AppError; +use crate::scope::ScopedPath; use std::fs::File; use std::io::{BufRead, BufReader}; -use std::path::Path; +use std::num::NonZeroUsize; /// Read file contents with pagination. /// @@ -17,18 +18,20 @@ use std::path::Path; /// Returns [`AppError::InvalidRequest`] if the target is missing or not a /// regular file. pub fn cat( - file_path: &Path, + file_path: &ScopedPath, offset: usize, - max_lines: usize, - max_bytes: usize, + max_lines: NonZeroUsize, + max_bytes: NonZeroUsize, ) -> Result { - if !file_path.is_file() { + let (max_lines, max_bytes) = (max_lines.get(), max_bytes.get()); + if !file_path.as_ref().is_file() { return Err(AppError::InvalidRequest( "Target is not a file or does not exist".to_string(), )); } - let file = File::open(file_path)?; + let file = File::open(file_path.as_ref())?; + let mut reader = BufReader::new(file); // Skip `offset` lines. @@ -44,7 +47,7 @@ pub fn cat( let mut output = String::new(); let mut line_count = 0usize; let mut truncated = false; - let mut truncation_reason: Option = None; + let mut truncation_reason: Option = None; let mut buf = String::new(); loop { buf.clear(); @@ -57,7 +60,7 @@ pub fn cat( output.push('\n'); } truncated = true; - truncation_reason = Some("line_cap".to_string()); + truncation_reason = Some(TruncationReason::LineCap); break; } if output.len() + buf.len() > max_bytes { @@ -71,7 +74,7 @@ pub fn cat( output.push('\n'); } truncated = true; - truncation_reason = Some("byte_cap".to_string()); + truncation_reason = Some(TruncationReason::ByteCap); break; } output.push_str(&buf); @@ -92,7 +95,7 @@ pub fn cat( #[cfg(test)] mod tests { use super::*; - use crate::tools::testutil::TestResult; + use crate::tools::testutil::{TestResult, nz, scoped_in}; use crate::tools::{DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES}; use std::fs; @@ -102,16 +105,16 @@ mod tests { let path = td.path().join("a.txt"); fs::write(&path, "L1\nL2\nL3\nL4\nL5\nL6\nL7\n")?; - let res = cat(&path, 2, 3, DEFAULT_MAX_BYTES)?; + let res = cat(&scoped_in(td.path(), &path), 2, nz(3), DEFAULT_MAX_BYTES)?; assert!( res.content.starts_with("L3\nL4\nL5\n"), "got {:?}", res.content ); assert!(res.truncated, "expected truncated=true"); - assert_eq!(res.truncation_reason, Some("line_cap".to_string())); + assert_eq!(res.truncation_reason, Some(TruncationReason::LineCap)); - let res = cat(&path, 4, 3, DEFAULT_MAX_BYTES)?; + let res = cat(&scoped_in(td.path(), &path), 4, nz(3), DEFAULT_MAX_BYTES)?; assert_eq!(res.content, "L5\nL6\nL7\n", "got {:?}", res.content); assert!(!res.truncated); Ok(()) @@ -124,9 +127,9 @@ mod tests { let body = "abcdefghijklmnopqrstuvwxyz\n".repeat(20); fs::write(&path, &body)?; - let res = cat(&path, 0, DEFAULT_MAX_LINES, 50)?; + let res = cat(&scoped_in(td.path(), &path), 0, DEFAULT_MAX_LINES, nz(50))?; assert!(res.truncated, "expected truncated=true, got {:?}", res); - assert_eq!(res.truncation_reason, Some("byte_cap".to_string())); + assert_eq!(res.truncation_reason, Some(TruncationReason::ByteCap)); assert!( res.content.len() < body.len(), "expected truncation, got len {}", @@ -139,7 +142,7 @@ mod tests { fn cat_errors_when_path_is_directory() -> TestResult { let td = tempfile::TempDir::new()?; match cat( - td.path(), + &scoped_in(td.path(), td.path()), 0, DEFAULT_MAX_LINES, DEFAULT_MAX_BYTES, diff --git a/src/tools/common.rs b/src/tools/common.rs index 001ac05..16c9432 100644 --- a/src/tools/common.rs +++ b/src/tools/common.rs @@ -2,6 +2,7 @@ //! extension filtering, error capture, and byte-capped channel draining. use ignore::WalkBuilder; +use std::num::NonZeroUsize; use std::path::Path; use std::sync::Mutex; use std::sync::mpsc::Receiver; @@ -55,7 +56,8 @@ pub(crate) fn extension_matches(path: &Path, extensions: &[String]) -> bool { /// a UTF-8 character boundary and a trailing newline is ensured; remaining /// chunks are discarded. Returns the assembled output and whether the cap was /// hit (`true` => truncated). -pub(crate) fn drain_capped(rx: &Receiver, max_bytes: usize) -> (String, bool) { +pub(crate) fn drain_capped(rx: &Receiver, max_bytes: NonZeroUsize) -> (String, bool) { + let max_bytes = max_bytes.get(); let mut output = String::new(); let mut byte_cap_hit = false; while let Ok(chunk) = rx.recv() { diff --git a/src/tools/find.rs b/src/tools/find.rs index 47e9095..5aab0ef 100644 --- a/src/tools/find.rs +++ b/src/tools/find.rs @@ -4,10 +4,10 @@ use super::common::{build_parallel_walker, record_first}; use super::options::FindOptions; use super::response::ToolResponse; use crate::error::AppError; +use crate::scope::ScopedPath; use ignore::WalkState; use regex::Regex; use std::mem; -use std::path::Path; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::mpsc::channel; use std::sync::{Arc, Mutex}; @@ -18,14 +18,18 @@ use std::sync::{Arc, Mutex}; /// Uses a parallel `ignore` walker (gitignore-aware) and an `AtomicUsize` /// counter for exact `max_results` capping. Walker entry errors are tallied /// and surfaced in the returned [`ToolResponse`] metadata. -pub fn find(directory: &Path, pattern: &str, opts: FindOptions) -> Result { +pub fn find( + directory: &ScopedPath, + pattern: &str, + opts: FindOptions, +) -> Result { let re = Regex::new(pattern)?; - let max_results = opts.max_results; + let max_results = opts.max_results.get(); let count = Arc::new(AtomicUsize::new(0)); let entry_errors = Arc::new(AtomicUsize::new(0)); let first_error: Arc>> = Arc::new(Mutex::new(None)); - let walker = build_parallel_walker(directory, &opts); + let walker = build_parallel_walker(directory.as_ref(), &opts); let (tx, rx) = channel::(); let match_basename = opts.match_basename; @@ -113,7 +117,7 @@ pub fn find(directory: &Path, pattern: &str, opts: FindOptions) -> Result TestResult { @@ -123,7 +127,7 @@ mod tests { write_file(root, "sub/bar.rs", "")?; let basename = find( - root, + &scoped_in(root, root), "^foo", FindOptions { match_basename: true, @@ -143,7 +147,7 @@ mod tests { ); let fullpath_anchored = find( - root, + &scoped_in(root, root), "^foo", FindOptions { match_basename: false, @@ -159,7 +163,7 @@ mod tests { ); let fullpath_ok = find( - root, + &scoped_in(root, root), r"sub.*foo\.rs$", FindOptions { match_basename: false, diff --git a/src/tools/grep.rs b/src/tools/grep.rs index b5f7f1a..4dec92d 100644 --- a/src/tools/grep.rs +++ b/src/tools/grep.rs @@ -2,9 +2,10 @@ use super::common::{build_parallel_walker, drain_capped, extension_matches, record_first}; use super::options::{GrepOptions, OutputMode}; -use super::response::ToolResponse; +use super::response::{ToolResponse, TruncationReason}; use super::sinks::{CountSink, FileMatchSink, MatchSink}; use crate::error::AppError; +use crate::scope::ScopedPath; use grep_regex::{RegexMatcher, RegexMatcherBuilder}; use grep_searcher::{BinaryDetection, SearcherBuilder}; use ignore::WalkState; @@ -75,7 +76,11 @@ impl ErrorState { /// Walker entry errors and per-file search errors are tallied and surfaced in /// the returned [`ToolResponse`] metadata rather than aborting the search. #[allow(clippy::needless_pass_by_value)] -pub fn grep(directory: &Path, pattern: &str, opts: GrepOptions) -> Result { +pub fn grep( + directory: &ScopedPath, + pattern: &str, + opts: GrepOptions, +) -> Result { match opts.output_mode { OutputMode::Content => grep_streamed(directory, pattern, &opts, StreamMode::Content), OutputMode::FilesWithMatches => { @@ -97,19 +102,19 @@ enum StreamMode { /// Unified implementation for `content` and `files_with_matches` modes. /// Both use the mpsc pipeline with early-quit on `max_results`. fn grep_streamed( - directory: &Path, + directory: &ScopedPath, pattern: &str, opts: &GrepOptions, mode: StreamMode, ) -> Result { let matcher = build_matcher(pattern, opts)?; let searcher_proto = build_searcher(opts, &mode); - let max_results = opts.max_results; - let max_bytes = opts.max_bytes; + let max_results = opts.max_results.get(); + let max_bytes = opts.max_bytes.get(); let errors = ErrorState::new(); let count = Arc::new(AtomicUsize::new(0)); let extensions = opts.file_extensions.clone(); - let walker = build_parallel_walker(directory, opts); + let walker = build_parallel_walker(directory.as_ref(), opts); let (tx, rx) = channel::(); walker.run(|| { @@ -183,7 +188,7 @@ fn grep_streamed( drop(tx); - let (output, byte_cap_hit) = drain_capped(&rx, max_bytes); + let (output, byte_cap_hit) = drain_capped(&rx, opts.max_bytes); let (entry_err_n, search_err_n, first_error) = errors.into_metadata(); let match_count = count.load(Ordering::Relaxed); @@ -191,7 +196,7 @@ fn grep_streamed( content: output, truncated: byte_cap_hit, truncation_reason: if byte_cap_hit { - Some("byte_cap".to_string()) + Some(TruncationReason::ByteCap) } else { None }, @@ -215,19 +220,19 @@ fn flush(tx: &Sender, buf: &mut String) { /// `count` mode — tally matches per file, output as `path: N` lines. /// Does not use a channel; collects into a shared HashMap instead. fn grep_count( - directory: &Path, + directory: &ScopedPath, pattern: &str, opts: &GrepOptions, ) -> Result { let matcher = build_matcher(pattern, opts)?; let searcher_proto = build_searcher(opts, &StreamMode::FilesWithMatches); // no context, no line numbers - let max_results = opts.max_results; - let max_bytes = opts.max_bytes; + let max_results = opts.max_results.get(); + let max_bytes = opts.max_bytes.get(); let errors = ErrorState::new(); let extensions = opts.file_extensions.clone(); let file_counts: Arc>> = Arc::new(Mutex::new(HashMap::new())); let file_count = Arc::new(AtomicUsize::new(0)); - let walker = build_parallel_walker(directory, opts); + let walker = build_parallel_walker(directory.as_ref(), opts); walker.run(|| { let errors = errors.clone(); @@ -310,7 +315,7 @@ fn grep_count( content: output, truncated, truncation_reason: if truncated { - Some("max_results".to_string()) + Some(TruncationReason::MaxResults) } else { None }, @@ -351,7 +356,7 @@ fn build_searcher(opts: &GrepOptions, mode: &StreamMode) -> grep_searcher::Searc #[cfg(test)] mod tests { use super::*; - use crate::tools::testutil::{TestResult, write_file}; + use crate::tools::testutil::{TestResult, nz, scoped_in, write_file}; use std::fs; #[test] @@ -362,11 +367,11 @@ mod tests { write_file(root, &format!("f{}.txt", i), "needle here\n")?; } let opts = GrepOptions { - max_results: 10, + max_results: nz(10), respect_gitignore: false, ..Default::default() }; - let res = grep(root, "needle", opts)?; + let res = grep(&scoped_in(root, root), "needle", opts)?; assert!( res.match_count.unwrap() <= 15, "expected match_count <= 15, got {:?}", @@ -383,7 +388,7 @@ mod tests { write_file(root, "a.txt", "Hello World\n")?; let case_sensitive = grep( - root, + &scoped_in(root, root), "hello", GrepOptions { output_mode: OutputMode::Content, @@ -399,7 +404,7 @@ mod tests { ); let case_insensitive = grep( - root, + &scoped_in(root, root), "hello", GrepOptions { case_insensitive: true, @@ -424,7 +429,7 @@ mod tests { write_file(root, "b.txt", "fn target() {}\n")?; let res = grep( - root, + &scoped_in(root, root), "target", GrepOptions { file_extensions: vec!["rs".to_string()], @@ -447,7 +452,7 @@ mod tests { write_file(root, "open.txt", "needle\n")?; let respected = grep( - root, + &scoped_in(root, root), "needle", GrepOptions { respect_gitignore: true, @@ -466,7 +471,7 @@ mod tests { ); let ignored = grep( - root, + &scoped_in(root, root), "needle", GrepOptions { respect_gitignore: false, @@ -490,7 +495,7 @@ mod tests { write_file(root, "c.rs", "needle here\n")?; let res = grep( - root, + &scoped_in(root, root), "needle", GrepOptions { output_mode: OutputMode::FilesWithMatches, @@ -519,11 +524,11 @@ mod tests { } let res = grep( - root, + &scoped_in(root, root), "needle", GrepOptions { output_mode: OutputMode::FilesWithMatches, - max_results: 5, + max_results: nz(5), respect_gitignore: false, ..Default::default() }, @@ -545,7 +550,7 @@ mod tests { write_file(root, "c.rs", "needle here\n")?; let res = grep( - root, + &scoped_in(root, root), "needle", GrepOptions { output_mode: OutputMode::Count, diff --git a/src/tools/mod.rs b/src/tools/mod.rs index e808856..1037fed 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -25,19 +25,35 @@ pub use grep::grep; pub use options::{FindOptions, GrepOptions, OutputMode}; pub use response::ToolResponse; -pub(crate) const DEFAULT_MAX_BYTES: usize = 5 * 1024 * 1024; // 5 MiB -pub(crate) const DEFAULT_MAX_RESULTS: usize = 100; -pub(crate) const DEFAULT_MAX_LINES: usize = 2000; +use std::num::NonZeroUsize; +pub(crate) const DEFAULT_MAX_BYTES: NonZeroUsize = NonZeroUsize::new(5 * 1024 * 1024).unwrap(); // 5 MiB +pub(crate) const DEFAULT_MAX_RESULTS: NonZeroUsize = NonZeroUsize::new(100).unwrap(); +pub(crate) const DEFAULT_MAX_LINES: NonZeroUsize = NonZeroUsize::new(2000).unwrap(); /// Shared test helpers used across the per-tool test modules. #[cfg(test)] pub(crate) mod testutil { use std::fs; use std::io::Write; + use std::num::NonZeroUsize; use std::path::Path; pub(crate) type TestResult = Result<(), Box>; + /// Build a [`NonZeroUsize`] from a test literal. + pub(crate) fn nz(n: usize) -> NonZeroUsize { + NonZeroUsize::new(n).expect("test literal must be nonzero") + } + + /// Obtain a [`ScopedPath`] for a tool entry point by running `target` + /// through a real [`Scope`] rooted at `root` — the same validation the + /// server performs per request. + pub(crate) fn scoped_in(root: &Path, target: &Path) -> crate::scope::ScopedPath { + crate::scope::Scope::new(root) + .and_then(|s| s.check(target)) + .expect("tempdir path should canonicalize inside its own scope") + } + pub(crate) fn write_file(dir: &Path, name: &str, contents: &str) -> std::io::Result<()> { let path = dir.join(name); if let Some(parent) = path.parent() { diff --git a/src/tools/options.rs b/src/tools/options.rs index a542010..39ad5ba 100644 --- a/src/tools/options.rs +++ b/src/tools/options.rs @@ -3,6 +3,7 @@ use super::{DEFAULT_MAX_BYTES, DEFAULT_MAX_RESULTS}; use rmcp::schemars::{self, JsonSchema}; use serde::{Deserialize, Serialize}; +use std::num::NonZeroUsize; /// Controls what the `grep` tool emits for each match. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize, Serialize, JsonSchema)] @@ -28,7 +29,9 @@ pub struct GrepOptions { pub after_context: usize, /// Exact cap on results. For `files_with_matches`/`count` this caps the /// number of files; for `content` it caps the number of matching lines. - pub max_results: usize, + /// Nonzero: a zero cap is unrepresentable, so `0` is rejected at the + /// JSON boundary before any walking starts. + pub max_results: NonZeroUsize, /// Case-insensitive matching (equivalent to a `(?i)` prefix on the pattern). pub case_insensitive: bool, /// Include hidden files and directories in the walk. @@ -39,8 +42,8 @@ pub struct GrepOptions { pub respect_gitignore: bool, /// Restrict search to files with these extensions (empty = all files). pub file_extensions: Vec, - /// Hard cap on total response size in bytes. - pub max_bytes: usize, + /// Hard cap on total response size in bytes (nonzero). + pub max_bytes: NonZeroUsize, /// What to emit for each match (see [`OutputMode`]). pub output_mode: OutputMode, } @@ -77,8 +80,8 @@ impl super::common::WalkerConfig for GrepOptions { /// Configuration for the `find` tool. #[derive(Clone, Copy)] pub struct FindOptions { - /// Exact cap on the number of matching paths returned. - pub max_results: usize, + /// Exact cap on the number of matching paths returned (nonzero). + pub max_results: NonZeroUsize, /// Include hidden files and directories in the walk. pub include_hidden: bool, /// Respect `.gitignore` / global / exclude gitignore rules. @@ -120,7 +123,10 @@ mod tests { #[test] fn grep_output_mode_rejects_unknown() -> TestResult { let result: Result = serde_json::from_str(r#""bogus""#); - assert!(result.is_err(), "expected deserialization error for unknown output_mode"); + assert!( + result.is_err(), + "expected deserialization error for unknown output_mode" + ); Ok(()) } diff --git a/src/tools/response.rs b/src/tools/response.rs index 65aee4b..b815700 100644 --- a/src/tools/response.rs +++ b/src/tools/response.rs @@ -4,6 +4,21 @@ use rmcp::model::CallToolResult; use serde::Serialize; use serde_json::json; +/// Why a tool response was truncated. Serialized `snake_case`, so the wire +/// contract (`"byte_cap"` / `"line_cap"` / `"max_results"`) is unchanged from +/// the previous freeform strings — but a typo'd or forgotten variant is now a +/// compile error instead of a silent client-side mismatch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum TruncationReason { + /// Total response size hit `max_bytes`. + ByteCap, + /// Line count hit `max_lines` (`cat` only). + LineCap, + /// Match/file count hit `max_results` (`grep` count mode). + MaxResults, +} + /// Structured metadata returned alongside the text content of every tool call. /// /// Serialized as the `structured_content` field of an MCP `CallToolResult`, so @@ -16,8 +31,9 @@ pub struct ToolResponse { pub content: String, /// Whether the output was truncated due to a size cap. pub truncated: bool, - /// If truncated, the reason (e.g. "`byte_cap`", "`line_cap`"). - pub truncation_reason: Option, + /// If `truncated`, why. Set together with `truncated` by construction: + /// every producer sets both or neither. + pub truncation_reason: Option, /// Number of matches found (grep / find). pub match_count: Option, /// Number of walker entry errors encountered.