#![allow(missing_docs)]
use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
use deno_core::op2;
use deno_core::OpState;
use deno_error::JsErrorBox;
use std::collections::HashSet;
use crate::stash::validate_key;
use crate::{ResourceDispatcher, StashDispatcher, ToolDispatcher};
pub(crate) struct ToolCallLimits {
pub(crate) max_calls: usize,
pub(crate) max_args_size: usize,
pub(crate) calls_made: usize,
}
pub(crate) struct StashCallLimits {
pub(crate) max_calls: Option<usize>,
pub(crate) calls_made: usize,
}
impl StashCallLimits {
pub(crate) fn check_limit(&mut self) -> Result<(), String> {
if let Some(max) = self.max_calls {
if self.calls_made >= max {
return Err(format!(
"stash operation limit reached ({max} calls per execution)"
));
}
}
self.calls_made += 1;
Ok(())
}
}
#[op2(fast)]
pub fn op_forge_log(#[string] msg: &str) {
tracing::info!(target: "forge::sandbox::js", "{}", msg);
}
#[op2(fast)]
pub fn op_forge_set_result(state: &mut OpState, #[string] json: &str) {
state.put(ExecutionResult(json.to_string()));
}
#[op2]
#[string]
pub async fn op_forge_call_tool(
op_state: Rc<RefCell<OpState>>,
#[string] server: String,
#[string] tool: String,
#[string] args_json: String,
) -> Result<String, JsErrorBox> {
tracing::debug!(
server = %server,
tool = %tool,
args_len = args_json.len(),
"tool call dispatched"
);
{
let mut st = op_state.borrow_mut();
let limits = st.borrow_mut::<ToolCallLimits>();
if limits.calls_made >= limits.max_calls {
return Err(JsErrorBox::generic(format!(
"tool call limit exceeded (max {} calls per execution)",
limits.max_calls
)));
}
if args_json.len() > limits.max_args_size {
return Err(JsErrorBox::generic(format!(
"tool call args too large ({} bytes, max {} bytes)",
args_json.len(),
limits.max_args_size
)));
}
limits.calls_made += 1;
}
let dispatcher = {
let st = op_state.borrow();
st.borrow::<Arc<dyn ToolDispatcher>>().clone()
};
let args: serde_json::Value = serde_json::from_str(&args_json)
.map_err(|e| JsErrorBox::generic(format!("invalid JSON args: {e}")))?;
let result = match dispatcher.call_tool(&server, &tool, args).await {
Ok(val) => val,
Err(e) => {
let known = {
let st = op_state.borrow();
st.try_borrow::<KnownTools>()
.map(|kt| kt.0.clone())
.unwrap_or_default()
};
let pairs: Vec<(&str, &str)> = known
.iter()
.map(|(s, t)| (s.as_str(), t.as_str()))
.collect();
let mut structured = e.to_structured_error(Some(&pairs));
crate::redact::redact_structured_error(&server, &tool, &mut structured);
return serde_json::to_string(&structured)
.map_err(|e| JsErrorBox::generic(format!("error serialization failed: {e}")));
}
};
serde_json::to_string(&result)
.map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
}
pub(crate) struct ExecutionResult(pub(crate) String);
pub(crate) struct MaxResourceSize(pub(crate) usize);
pub(crate) struct CurrentGroup(pub(crate) Option<String>);
pub(crate) struct KnownServers(pub(crate) HashSet<String>);
pub(crate) struct KnownTools(pub(crate) Vec<(String, String)>);
pub(crate) fn validate_resource_uri(uri: &str) -> Result<(), String> {
if uri.len() > 2048 {
return Err(format!(
"resource URI too long ({} bytes, max 2048 bytes)",
uri.len()
));
}
if uri.bytes().any(|b| b == 0) {
return Err("resource URI must not contain null bytes".into());
}
if uri.chars().any(|c| c.is_control()) {
return Err("resource URI must not contain control characters".into());
}
if has_path_traversal(uri) {
return Err("resource URI must not contain path traversal".into());
}
if let Some(scheme) = extract_uri_scheme(uri) {
if is_blocked_scheme(&scheme) {
return Err(format!("URI scheme '{}' is not allowed", scheme));
}
}
Ok(())
}
const BLOCKED_SCHEMES: &[&str] = &[
"data",
"javascript",
"ftp",
"gopher",
"telnet",
"ldap",
"dict",
];
fn extract_uri_scheme(uri: &str) -> Option<String> {
if let Some(pos) = uri.find("://") {
let candidate = &uri[..pos];
if !candidate.is_empty()
&& candidate.as_bytes()[0].is_ascii_alphabetic()
&& candidate
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'-' || b == b'.')
{
return Some(candidate.to_ascii_lowercase());
}
}
if let Some(pos) = uri.find(':') {
let candidate = &uri[..pos];
if !candidate.is_empty()
&& candidate.as_bytes()[0].is_ascii_alphabetic()
&& candidate
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'-' || b == b'.')
{
return Some(candidate.to_ascii_lowercase());
}
}
None
}
fn is_blocked_scheme(scheme: &str) -> bool {
BLOCKED_SCHEMES.contains(&scheme)
}
fn has_path_traversal(uri: &str) -> bool {
if has_dotdot_segment(uri) {
return true;
}
let decoded = percent_encoding::percent_decode_str(uri).decode_utf8_lossy();
if has_dotdot_segment(&decoded) {
return true;
}
let double_decoded = percent_encoding::percent_decode_str(&decoded).decode_utf8_lossy();
if double_decoded != decoded && has_dotdot_segment(&double_decoded) {
return true;
}
false
}
fn has_dotdot_segment(s: &str) -> bool {
if s == ".." {
return true;
}
if s.starts_with("../") {
return true;
}
if s.ends_with("/..") {
return true;
}
if s.contains("/../") {
return true;
}
false
}
#[op2]
#[string]
pub async fn op_forge_read_resource(
op_state: Rc<RefCell<OpState>>,
#[string] server: String,
#[string] uri: String,
) -> Result<String, JsErrorBox> {
tracing::debug!(
server = %server,
uri = %uri,
"resource read dispatched"
);
validate_resource_uri(&uri).map_err(JsErrorBox::generic)?;
{
let st = op_state.borrow();
if let Some(known) = st.try_borrow::<KnownServers>() {
if !known.0.contains(&server) {
return Err(JsErrorBox::generic(format!("unknown server: '{server}'")));
}
}
}
{
let mut st = op_state.borrow_mut();
let limits = st.borrow_mut::<ToolCallLimits>();
if limits.calls_made >= limits.max_calls {
return Err(JsErrorBox::generic(format!(
"tool call limit exceeded (max {} calls per execution)",
limits.max_calls
)));
}
limits.calls_made += 1;
}
let (dispatcher, max_resource_size) = {
let st = op_state.borrow();
let d = st.borrow::<Arc<dyn ResourceDispatcher>>().clone();
let max_size = st
.try_borrow::<MaxResourceSize>()
.map(|m| m.0)
.unwrap_or(64 * 1024 * 1024); (d, max_size)
};
let result = match dispatcher.read_resource(&server, &uri).await {
Ok(val) => val,
Err(e) => {
let known = {
let st = op_state.borrow();
st.try_borrow::<KnownTools>()
.map(|kt| kt.0.clone())
.unwrap_or_default()
};
let pairs: Vec<(&str, &str)> = known
.iter()
.map(|(s, t)| (s.as_str(), t.as_str()))
.collect();
let mut structured = e.to_structured_error(Some(&pairs));
crate::redact::redact_structured_error(&server, "readResource", &mut structured);
return serde_json::to_string(&structured)
.map_err(|e| JsErrorBox::generic(format!("error serialization failed: {e}")));
}
};
let mut json = serde_json::to_string(&result)
.map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))?;
if json.len() > max_resource_size {
let truncated = &json[..max_resource_size];
let end = truncated
.char_indices()
.last()
.map(|(i, c)| i + c.len_utf8())
.unwrap_or(0);
json = json[..end].to_string();
}
Ok(json)
}
#[op2]
#[string]
pub async fn op_forge_stash_put(
op_state: Rc<RefCell<OpState>>,
#[string] key: String,
#[string] value_json: String,
#[smi] ttl_secs: u32,
) -> Result<String, JsErrorBox> {
validate_key(&key).map_err(|e| JsErrorBox::generic(e.to_string()))?;
if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
limits.check_limit().map_err(JsErrorBox::generic)?;
}
let (dispatcher, current_group) = {
let st = op_state.borrow();
let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
let g = st.borrow::<CurrentGroup>().0.clone();
(d, g)
};
let value: serde_json::Value = serde_json::from_str(&value_json)
.map_err(|e| JsErrorBox::generic(format!("invalid JSON value: {e}")))?;
let ttl = if ttl_secs == 0 { None } else { Some(ttl_secs) };
let result = dispatcher
.put(&key, value, ttl, current_group)
.await
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
serde_json::to_string(&result)
.map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
}
#[op2]
#[string]
pub async fn op_forge_stash_get(
op_state: Rc<RefCell<OpState>>,
#[string] key: String,
) -> Result<String, JsErrorBox> {
validate_key(&key).map_err(|e| JsErrorBox::generic(e.to_string()))?;
if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
limits.check_limit().map_err(JsErrorBox::generic)?;
}
let (dispatcher, current_group) = {
let st = op_state.borrow();
let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
let g = st.borrow::<CurrentGroup>().0.clone();
(d, g)
};
let result = dispatcher
.get(&key, current_group)
.await
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
serde_json::to_string(&result)
.map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
}
#[op2]
#[string]
pub async fn op_forge_stash_delete(
op_state: Rc<RefCell<OpState>>,
#[string] key: String,
) -> Result<String, JsErrorBox> {
validate_key(&key).map_err(|e| JsErrorBox::generic(e.to_string()))?;
if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
limits.check_limit().map_err(JsErrorBox::generic)?;
}
let (dispatcher, current_group) = {
let st = op_state.borrow();
let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
let g = st.borrow::<CurrentGroup>().0.clone();
(d, g)
};
let result = dispatcher
.delete(&key, current_group)
.await
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
serde_json::to_string(&result)
.map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
}
#[op2]
#[string]
pub async fn op_forge_stash_keys(op_state: Rc<RefCell<OpState>>) -> Result<String, JsErrorBox> {
if let Some(limits) = op_state.borrow_mut().try_borrow_mut::<StashCallLimits>() {
limits.check_limit().map_err(JsErrorBox::generic)?;
}
let (dispatcher, current_group) = {
let st = op_state.borrow();
let d = st.borrow::<Arc<dyn StashDispatcher>>().clone();
let g = st.borrow::<CurrentGroup>().0.clone();
(d, g)
};
let result = dispatcher
.keys(current_group)
.await
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
serde_json::to_string(&result)
.map_err(|e| JsErrorBox::generic(format!("result serialization failed: {e}")))
}
deno_core::extension!(
forge_ext,
ops = [
op_forge_log,
op_forge_set_result,
op_forge_call_tool,
op_forge_read_resource,
op_forge_stash_put,
op_forge_stash_get,
op_forge_stash_delete,
op_forge_stash_keys
],
);
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sk_v01_stash_key_rejects_control_chars() {
let err = validate_key("key\x01value").unwrap_err();
assert!(
matches!(err, crate::stash::StashError::InvalidKey),
"expected InvalidKey, got: {err}"
);
}
#[test]
fn sk_v02_stash_key_rejects_path_separators() {
assert!(validate_key("key/value").is_err());
assert!(validate_key("key\\value").is_err());
assert!(validate_key("../etc/passwd").is_err());
}
#[test]
fn sk_v03_stash_key_rejects_empty_and_oversized() {
assert!(validate_key("").is_err());
let long_key = "a".repeat(257);
let err = validate_key(&long_key).unwrap_err();
assert!(
matches!(err, crate::stash::StashError::KeyTooLong { len: 257 }),
"expected KeyTooLong, got: {err}"
);
}
#[test]
fn sk_v04_stash_key_accepts_valid_patterns() {
assert!(validate_key("simple-key").is_ok());
assert!(validate_key("key_with.dots:colons").is_ok());
assert!(validate_key("CamelCase123").is_ok());
assert!(validate_key("a").is_ok());
let max_key = "x".repeat(256);
assert!(validate_key(&max_key).is_ok());
}
#[test]
fn sk_v05_stash_key_rejects_unicode() {
assert!(validate_key("key\u{0000}null").is_err());
assert!(validate_key("key with space").is_err());
assert!(validate_key("emoji\u{1F600}").is_err());
}
#[test]
fn rs_u04_rejects_uri_with_path_traversal() {
assert!(validate_resource_uri("file:///logs/../../../etc/passwd").is_err());
assert!(validate_resource_uri("file:///..").is_err());
assert!(validate_resource_uri("..").is_err());
assert!(validate_resource_uri("a/../../b").is_err());
assert!(validate_resource_uri("file:///logs/app.log").is_ok());
assert!(validate_resource_uri("postgres://db/table").is_ok());
}
#[test]
fn uri_v01_allows_legitimate_double_dots() {
assert!(validate_resource_uri("file:///v2..backup").is_ok());
assert!(validate_resource_uri("file:///data..2024.csv").is_ok());
assert!(validate_resource_uri("file:///config..old").is_ok());
}
#[test]
fn uri_v02_blocks_url_encoded_traversal() {
assert!(validate_resource_uri("file:///logs/%2e%2e/%2e%2e/etc/passwd").is_err());
assert!(validate_resource_uri("file:///%2e%2e/secret").is_err());
}
#[test]
fn uri_v03_blocks_double_encoded_traversal() {
assert!(validate_resource_uri("file:///logs/%252e%252e/%252e%252e/etc/passwd").is_err());
}
#[test]
fn uri_v04_blocks_mixed_case_encoded_traversal() {
assert!(validate_resource_uri("file:///logs/%2E%2E/secret").is_err());
assert!(validate_resource_uri("file:///logs/%2e%2E/secret").is_err());
}
#[test]
fn rs_u05_rejects_uri_longer_than_2048_bytes() {
let long_uri = "x".repeat(2049);
let err = validate_resource_uri(&long_uri).unwrap_err();
assert!(err.contains("too long"), "should mention too long: {err}");
let ok_uri = "x".repeat(2048);
assert!(validate_resource_uri(&ok_uri).is_ok());
}
#[test]
fn rs_u06_rejects_uri_with_null_bytes() {
let uri = "file:///logs\0/app.log";
let err = validate_resource_uri(uri).unwrap_err();
assert!(err.contains("null"), "should mention null: {err}");
}
#[test]
fn rs_u07_rejects_uri_with_control_characters() {
let err = validate_resource_uri("file:///logs\x01/app.log").unwrap_err();
assert!(err.contains("control"), "should mention control: {err}");
assert!(validate_resource_uri("file:///logs\t/app.log").is_err());
assert!(validate_resource_uri("file:///logs\n/app.log").is_err());
assert!(validate_resource_uri("file:///logs\x7f/app.log").is_err());
}
#[test]
fn rs_s04_path_traversal_attack_variants() {
assert!(validate_resource_uri("../../../etc/passwd").is_err());
assert!(validate_resource_uri("file:///logs/%2e%2e/%2e%2e/etc/passwd").is_err());
assert!(validate_resource_uri("..").is_err());
assert!(validate_resource_uri("file:///../").is_err());
assert!(validate_resource_uri("file:///a/b/../../../etc/shadow").is_err());
}
#[test]
fn uri_m2_01_allows_http_scheme() {
assert!(validate_resource_uri("http://example.com/resource").is_ok());
}
#[test]
fn uri_m2_02_allows_https_scheme() {
assert!(validate_resource_uri("https://example.com/resource").is_ok());
}
#[test]
fn uri_m2_03_allows_file_scheme() {
assert!(validate_resource_uri("file:///logs/app.log").is_ok());
}
#[test]
fn uri_m2_04_rejects_data_scheme() {
let err = validate_resource_uri("data:text/plain;base64,SGVsbG8=").unwrap_err();
assert!(
err.contains("not allowed"),
"expected 'not allowed' in error: {err}"
);
}
#[test]
fn uri_m2_05_rejects_javascript_scheme() {
let err = validate_resource_uri("javascript:alert(1)").unwrap_err();
assert!(err.contains("not allowed"), "error: {err}");
}
#[test]
fn uri_m2_06_rejects_ftp_scheme() {
let err = validate_resource_uri("ftp://evil.com/malware").unwrap_err();
assert!(err.contains("not allowed"), "error: {err}");
}
#[test]
fn uri_m2_07_rejects_gopher_scheme() {
let err = validate_resource_uri("gopher://evil.com/0").unwrap_err();
assert!(err.contains("not allowed"), "error: {err}");
}
#[test]
fn uri_m2_08_allows_custom_mcp_scheme() {
assert!(validate_resource_uri("postgres://db/table").is_ok());
assert!(validate_resource_uri("redis://localhost:6379/0").is_ok());
assert!(validate_resource_uri("mongodb://host/db").is_ok());
}
#[test]
fn uri_m2_09_allows_schemeless_uri() {
assert!(validate_resource_uri("some-resource-id").is_ok());
assert!(validate_resource_uri("table_name").is_ok());
assert!(validate_resource_uri("logs/2024/app.log").is_ok());
}
#[test]
fn uri_m2_10_case_insensitive_scheme_check() {
assert!(validate_resource_uri("JAVASCRIPT:alert(1)").is_err());
assert!(validate_resource_uri("JavaScript:void(0)").is_err());
assert!(validate_resource_uri("DATA:text/plain,hello").is_err());
assert!(validate_resource_uri("FTP://evil.com/file").is_err());
assert!(validate_resource_uri("Gopher://host/0").is_err());
assert!(validate_resource_uri("TELNET://host:23").is_err());
assert!(validate_resource_uri("LDAP://host/dc=com").is_err());
assert!(validate_resource_uri("DICT://host/define").is_err());
}
#[test]
fn stash_l4_01_stash_calls_count_against_limit() {
let mut limits = StashCallLimits {
max_calls: Some(3),
calls_made: 0,
};
assert!(limits.check_limit().is_ok());
assert!(limits.check_limit().is_ok());
assert!(limits.check_limit().is_ok());
assert!(limits.check_limit().is_err());
}
#[test]
fn stash_l4_02_stash_limit_rejection_message() {
let mut limits = StashCallLimits {
max_calls: Some(1),
calls_made: 0,
};
assert!(limits.check_limit().is_ok());
let err = limits.check_limit().unwrap_err();
assert!(err.contains("limit reached"), "should mention limit: {err}");
assert!(
err.contains("1 calls"),
"should mention the limit count: {err}"
);
}
#[test]
fn stash_l4_03_stash_limit_independent_of_tool_limit() {
let mut stash_limits = StashCallLimits {
max_calls: Some(5),
calls_made: 0,
};
let tool_limits = ToolCallLimits {
max_calls: 0, max_args_size: 1024,
calls_made: 0,
};
assert!(stash_limits.check_limit().is_ok());
let _ = tool_limits;
}
#[test]
fn stash_l4_04_stash_limit_configurable() {
let mut unlimited = StashCallLimits {
max_calls: None,
calls_made: 0,
};
for _ in 0..1000 {
assert!(unlimited.check_limit().is_ok());
}
let mut limited = StashCallLimits {
max_calls: Some(2),
calls_made: 0,
};
assert!(limited.check_limit().is_ok());
assert!(limited.check_limit().is_ok());
assert!(limited.check_limit().is_err());
}
}