feat: enhance safety filters for shell commands by normalizing ANSI-C quoting

This commit is contained in:
asepharyana
2026-07-13 04:10:08 +07:00
parent a080957c26
commit d09e440e7e
14 changed files with 383 additions and 85 deletions
+82 -22
View File
@@ -16,7 +16,13 @@ const MCP_CALL_TIMEOUT_MS: u64 = 60_000;
/// number of MCP tools ever registered in a session).
fn mcp_static_str(s: &str) -> &'static str {
static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
let mut cache = match CACHE.get_or_init(|| Mutex::new(Vec::new())).lock() {
Ok(c) => c,
Err(poisoned) => {
tracing::warn!("[mcp] static string cache mutex poisoned, recovering");
poisoned.into_inner()
}
};
if let Some(&existing) = cache.iter().find(|e| **e == s) {
return existing;
}
@@ -104,30 +110,71 @@ impl StdioChild {
anyhow::bail!("MCP call timed out after {}ms", MCP_CALL_TIMEOUT_MS);
}
response_line.clear();
match self.stdout.read_line(&mut response_line) {
Ok(0) => anyhow::bail!("MCP stdio child process closed unexpectedly"),
Ok(_) => {
let trimmed = response_line.trim();
if trimmed.is_empty() {
continue;
// Read one byte at a time up to MAX_LINE_LENGTH to prevent
// OOM from a malicious server (CWE-400). BufReader already
// buffers reads, so byte-by-byte over a buffered reader is
// cheap (hits the in-memory buffer).
const MAX_LINE_LENGTH: usize = 1_048_576; // 1 MiB
let mut line_truncated = false;
loop {
let byte = match self.stdout.fill_buf() {
Ok(buf) if buf.is_empty() => {
// EOF without newline
anyhow::bail!("MCP stdio child process closed unexpectedly");
}
let resp: Value = serde_json::from_str(trimmed)
.map_err(|e| anyhow::anyhow!("invalid JSON from MCP server: {}", e))?;
if resp.get("id") == Some(&json!(id)) {
if let Some(err) = resp.get("error") {
anyhow::bail!("MCP error: {}", err);
}
return Ok(resp.get("result").cloned().unwrap_or_else(|| {
tracing::warn!("[mcp] stdio response missing 'result' field: {}", trimmed);
Value::Null
}));
Ok(buf) => {
let b = buf[0];
self.stdout.consume(1);
b
}
Err(e) => anyhow::bail!("MCP stdio read error: {}", e),
};
if byte == b'\n' {
break;
}
Err(e) => anyhow::bail!("MCP stdio read error: {}", e),
if response_line.len() >= MAX_LINE_LENGTH {
line_truncated = true;
// Consume rest of line to keep stream in sync
loop {
let buf = self.stdout.fill_buf()
.map_err(|e| anyhow::anyhow!("MCP stdio read error: {}", e))?;
if buf.is_empty() {
anyhow::bail!("MCP stdio child closed mid-line");
}
if buf[0] == b'\n' {
self.stdout.consume(1);
break;
}
self.stdout.consume(1);
}
break;
}
response_line.push(byte as char);
}
if line_truncated {
anyhow::bail!(
"MCP response line exceeded {} byte limit",
MAX_LINE_LENGTH,
);
}
let trimmed = response_line.trim();
if trimmed.is_empty() {
continue;
}
let resp: Value = serde_json::from_str(trimmed)
.map_err(|e| anyhow::anyhow!("invalid JSON from MCP server: {}", e))?;
if resp.get("id") == Some(&json!(id)) {
if let Some(err) = resp.get("error") {
anyhow::bail!("MCP error: {}", err);
}
return Ok(resp.get("result").cloned().unwrap_or_else(|| {
tracing::warn!("[mcp] stdio response missing 'result' field: {}", trimmed);
Value::Null
}));
}
}
}
}
} // close fn call
} // close impl StdioChild
pub(crate) fn spawn_stdio_child(command: &str, extra_args: &[String]) -> anyhow::Result<StdioChild> {
let parts: Vec<&str> = command.split_whitespace().collect();
@@ -215,8 +262,21 @@ fn call_via_http(url: &str, tool_name: &str, tool_args: &Value) -> anyhow::Resul
.connect_timeout(std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS))
.build()
.unwrap_or_else(|e| {
tracing::warn!("[mcp] HTTP client builder failed: {}, using default client without timeouts", e);
reqwest::blocking::Client::new()
tracing::warn!(
"[mcp] HTTP client builder failed with connect timeout: {}. \
retrying without connect timeout",
e,
);
reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS))
.build()
.unwrap_or_else(|e2| {
tracing::warn!(
"[mcp] also failed: {}. using default client (no configured timeouts)",
e2,
);
reqwest::blocking::Client::new()
})
});
let request_id: u64 = 1;