From 94b448ffab521d82ff8226e8eb26887b6ed00155 Mon Sep 17 00:00:00 2001 From: Gabe Hamilton Date: Wed, 11 Mar 2026 17:29:59 -0600 Subject: [PATCH] fix(security): load WASM tool description and schema from capabilities.json (#520) The extract_tool_description and extract_tool_schema stubs in runtime.rs returned permissive fallbacks ("WASM sandboxed tool" and additionalProperties:true) for every WASM tool, defeating parameter validation and preventing the LLM from using tools correctly. Add optional `description` and `parameters` fields to CapabilitiesFile so tool authors can declare proper metadata in their sidecar JSON. The WasmToolLoader now extracts these fields and passes them through to the tool registry as overrides. Tools without a capabilities.json or without these fields get a tracing::warn and fall back to the old stubs. Co-authored-by: Claude Sonnet 4.6 --- src/tools/wasm/capabilities_schema.rs | 124 ++++++++++++++++++++++++++ src/tools/wasm/loader.rs | 89 ++++++++++++------ src/tools/wasm/runtime.rs | 23 +++-- 3 files changed, 202 insertions(+), 34 deletions(-) diff --git a/src/tools/wasm/capabilities_schema.rs b/src/tools/wasm/capabilities_schema.rs index 99007943..60cae7c2 100644 --- a/src/tools/wasm/capabilities_schema.rs +++ b/src/tools/wasm/capabilities_schema.rs @@ -41,6 +41,18 @@ use crate::tools::wasm::{ /// Root schema for a capabilities JSON file. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct CapabilitiesFile { + /// Human-readable description of what the tool does. + /// Used as the `Tool::description()` return value. + /// If omitted, a generic fallback is used (with a warning). + #[serde(default)] + pub description: Option, + + /// JSON Schema for the tool's input parameters. + /// Used as the `Tool::parameters_schema()` return value. + /// If omitted, a permissive fallback is used (with a warning). + #[serde(default)] + pub parameters: Option, + /// Extension version (semver). #[serde(default)] pub version: Option, @@ -107,6 +119,8 @@ impl CapabilitiesFile { fn resolve_nested(mut self) -> Self { if let Some(inner) = self.capabilities.take() { let inner = inner.resolve_nested(); + self.description = self.description.or(inner.description); + self.parameters = self.parameters.or(inner.parameters); self.http = self.http.or(inner.http); self.secrets = self.secrets.or(inner.secrets); self.tool_invoke = self.tool_invoke.or(inner.tool_invoke); @@ -1259,4 +1273,114 @@ mod tests { "Empty inner capabilities should not clobber outer http" ); } + + // ── Tool description and parameters schema ────────────────────────── + + #[test] + fn test_parse_description_and_parameters() { + let json = r#"{ + "description": "Search the web using Brave Search API", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query" + }, + "count": { + "type": "integer", + "description": "Number of results" + } + }, + "required": ["query"] + } + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + assert_eq!( + caps.description.as_deref(), + Some("Search the web using Brave Search API") + ); + let params = caps.parameters.unwrap(); + assert_eq!(params["type"], "object"); + assert!(params["properties"]["query"].is_object()); + assert_eq!(params["required"][0], "query"); + } + + #[test] + fn test_parse_description_only() { + let json = r#"{ + "description": "A tool without explicit parameters schema" + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + assert_eq!( + caps.description.as_deref(), + Some("A tool without explicit parameters schema") + ); + assert!(caps.parameters.is_none()); + } + + #[test] + fn test_parse_without_description_or_parameters() { + let json = r#"{ + "http": { + "allowlist": [{ "host": "api.example.com" }] + } + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + assert!( + caps.description.is_none(), + "description should be None when not provided" + ); + assert!( + caps.parameters.is_none(), + "parameters should be None when not provided" + ); + } + + #[test] + fn test_resolve_nested_description_promoted() { + let json = r#"{ + "capabilities": { + "description": "Inner tool description", + "parameters": { + "type": "object", + "properties": { + "input": { "type": "string" } + }, + "required": ["input"] + } + } + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + assert_eq!( + caps.description.as_deref(), + Some("Inner tool description"), + "description should be promoted from inner capabilities" + ); + assert!( + caps.parameters.is_some(), + "parameters should be promoted from inner capabilities" + ); + } + + #[test] + fn test_resolve_nested_outer_description_takes_precedence() { + let json = r#"{ + "description": "Outer description wins", + "capabilities": { + "description": "Inner description loses" + } + }"#; + + let caps = CapabilitiesFile::from_json(json).unwrap(); + assert_eq!( + caps.description.as_deref(), + Some("Outer description wins"), + "Outer description should take precedence over inner" + ); + } } diff --git a/src/tools/wasm/loader.rs b/src/tools/wasm/loader.rs index afa471a1..a96fc9bb 100644 --- a/src/tools/wasm/loader.rs +++ b/src/tools/wasm/loader.rs @@ -123,34 +123,73 @@ impl WasmToolLoader { } let wasm_bytes = fs::read(wasm_path).await?; - // Read capabilities (optional) and extract OAuth refresh config - let (capabilities, oauth_refresh) = if let Some(cap_path) = capabilities_path { - if cap_path.exists() { - let cap_bytes = fs::read(cap_path).await?; - let cap_file = CapabilitiesFile::from_bytes(&cap_bytes) - .map_err(|e| WasmLoadError::InvalidCapabilities(e.to_string()))?; - cap_file.validate(name); + // Read capabilities (optional) and extract OAuth refresh config, + // tool description, and parameter schema. + let (capabilities, oauth_refresh, description, schema) = + if let Some(cap_path) = capabilities_path { + if cap_path.exists() { + let cap_bytes = fs::read(cap_path).await?; + let cap_file = CapabilitiesFile::from_bytes(&cap_bytes) + .map_err(|e| WasmLoadError::InvalidCapabilities(e.to_string()))?; + cap_file.validate(name); - // Check WIT version compatibility - check_wit_version_compat( - name, - cap_file.wit_version.as_deref(), - crate::tools::wasm::WIT_TOOL_VERSION, - )?; + // Check WIT version compatibility + check_wit_version_compat( + name, + cap_file.wit_version.as_deref(), + crate::tools::wasm::WIT_TOOL_VERSION, + )?; - let caps = cap_file.to_capabilities(); - let oauth = resolve_oauth_refresh_config(&cap_file); - (caps, oauth) + let caps = cap_file.to_capabilities(); + let oauth = resolve_oauth_refresh_config(&cap_file); + let desc = cap_file.description.clone(); + // Validate parameters schema before accepting it. + let params = cap_file.parameters.clone().and_then(|p| { + let errors = crate::tools::validate_tool_schema(&p, name); + if errors.is_empty() { + Some(p) + } else { + tracing::warn!( + tool = name, + ?errors, + "Invalid parameters schema in capabilities.json, \ + using permissive fallback" + ); + None + } + }); + if desc.is_none() { + tracing::warn!( + tool = name, + path = %cap_path.display(), + "Capabilities file missing \"description\" field; \ + tool will use generic fallback description" + ); + } + if params.is_none() && cap_file.parameters.is_none() { + tracing::warn!( + tool = name, + path = %cap_path.display(), + "Capabilities file missing \"parameters\" field; \ + tool will accept any JSON object (permissive fallback)" + ); + } + (caps, oauth, desc, params) + } else { + tracing::warn!( + path = %cap_path.display(), + "Capabilities file not found, using default (no permissions)" + ); + (Capabilities::default(), None, None, None) + } } else { tracing::warn!( - path = %cap_path.display(), - "Capabilities file not found, using default (no permissions)" + tool = name, + "No capabilities file for WASM tool; \ + tool will use generic fallback description and accept any JSON object" ); - (Capabilities::default(), None) - } - } else { - (Capabilities::default(), None) - }; + (Capabilities::default(), None, None, None) + }; // Register the tool self.registry @@ -160,8 +199,8 @@ impl WasmToolLoader { runtime: &self.runtime, capabilities, limits: None, - description: None, - schema: None, + description: description.as_deref(), + schema, secrets_store: self.secrets_store.clone(), oauth_refresh, }) diff --git a/src/tools/wasm/runtime.rs b/src/tools/wasm/runtime.rs index 05e20de5..7645af3e 100644 --- a/src/tools/wasm/runtime.rs +++ b/src/tools/wasm/runtime.rs @@ -323,27 +323,32 @@ impl WasmToolRuntime { /// Extract tool description from a compiled component. /// -/// In a full implementation, this would use WIT bindgen to call the description() export. -/// For now, we return a placeholder since we can't easily introspect without more setup. +/// Returns a generic fallback. Callers should prefer loading the description +/// from the sidecar `*.capabilities.json` file and overriding via +/// `WasmToolWrapper::with_description()` or the `WasmToolRegistration::description` field. fn extract_tool_description( _engine: &Engine, _component: &wasmtime::component::Component, ) -> Result { - // TODO: Use WIT bindgen to properly extract description - // This requires instantiating with a linker, which needs host functions. - // For now, tools should have their description set externally. + // WIT bindgen extraction is not yet implemented (see TODO #4 in CLAUDE.md). + // Real descriptions come from the capabilities.json sidecar file, which is + // loaded by the WasmToolLoader and passed as an override at registration time. Ok("WASM sandboxed tool".to_string()) } -/// Extract tool schema from a compiled component. +/// Extract tool parameter schema from a compiled component. /// -/// In a full implementation, this would use WIT bindgen to call the schema() export. +/// Returns a permissive fallback that accepts any JSON object. Callers should +/// prefer loading the schema from the sidecar `*.capabilities.json` file and +/// overriding via `WasmToolWrapper::with_schema()` or the +/// `WasmToolRegistration::schema` field. fn extract_tool_schema( _engine: &Engine, _component: &wasmtime::component::Component, ) -> Result { - // TODO: Use WIT bindgen to properly extract schema - // For now, return a minimal schema that accepts any object. + // WIT bindgen extraction is not yet implemented (see TODO #4 in CLAUDE.md). + // Real schemas come from the capabilities.json sidecar file, which is + // loaded by the WasmToolLoader and passed as an override at registration time. Ok(serde_json::json!({ "type": "object", "properties": {},