MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
src/app/bgbash/job.rs
+30 -6
View File
@@ -15,10 +15,17 @@ use std::sync::mpsc;
use std::thread;
use std::io::BufRead;
/// Maximum number of output lines buffered in memory per background job.
/// Beyond this limit, old output is dropped to prevent OOM (CWE-770).
/// 10_000 lines at ~100 bytes each ≈ 1 MiB per job, sufficient for most
/// command output. The stderr drain thread also uses the same limit.
const MAX_OUTPUT_LINES: usize = 10_000;
/// Handle to a bash command running in a detached background thread.
///
/// Why: output is streamed over an mpsc channel rather than buffered
/// Why: output is streamed over a bounded mpsc channel rather than buffered
/// synchronously, so the TUI can poll for new lines without blocking.
/// The bounded channel prevents OOM from fast producers (e.g. `yes`).
pub struct BashJob {
pub id: String,
pub child_pid: u32,
@@ -43,9 +50,10 @@ pub struct BashJob {
/// output channel.
pub fn spawn_bash_job(command: String) -> BashJob {
let id = uuid::Uuid::new_v4().to_string();
let (output_tx, output_rx) = mpsc::channel::<String>();
let (output_tx, output_rx) = mpsc::sync_channel::<String>(MAX_OUTPUT_LINES);
let (pid_tx, pid_rx) = mpsc::channel::<u32>();
let cmd = command.clone();
let id_for_log = id.clone();
thread::spawn(move || {
let mut child = match Command::new("sh")
@@ -57,8 +65,8 @@ pub fn spawn_bash_job(command: String) -> BashJob {
{
Ok(c) => c,
Err(e) => {
let _ = output_tx.send(format!("__error:{}", e));
let _ = output_tx.send("__exit:-1".to_string());
let _ = output_tx.try_send(format!("__error:{}", e));
let _ = output_tx.try_send("__exit:-1".to_string());
return;
}
};
@@ -70,24 +78,40 @@ pub fn spawn_bash_job(command: String) -> BashJob {
// the child produces more than ~64 KB of stderr after closing
// stdout (the pipe buffer fills and the child blocks on write,
// while the parent thread waits for the child to exit).
let stderr_tx = output_tx.clone();
let _stderr_drain = child.stderr.take().map(|stderr| {
std::thread::spawn(move || {
let reader = std::io::BufReader::new(stderr);
// stderr is intentionally discarded to prevent output-line
// quota pressure from error diagnostics.
for _line in reader.lines().map_while(Result::ok) {
// Discard stderr lines to prevent pipe buffer deadlock.
}
drop(stderr_tx);
})
});
if let Some(stdout) = child.stdout.take() {
let reader = std::io::BufReader::new(stdout);
for line in reader.lines().map_while(Result::ok) {
let _ = output_tx.send(line);
// Use try_send so if the channel buffer is full (producer
// faster than consumer), old lines are silently dropped
// rather than growing memory without bound.
if output_tx.try_send(line).is_err() {
// Buffer full — consumer is not draining fast enough.
// Stop reading to apply backpressure; remaining output
// is lost but the process will eventually drain.
tracing::debug!(
"[bgbash:{}] output buffer full ({} lines), discarding remaining output",
id_for_log, MAX_OUTPUT_LINES,
);
break;
}
}
}
let status = child.wait();
let code = status.ok().and_then(|s| s.code());
let _ = output_tx.send(format!("__exit:{}", code.unwrap_or(-1)));
let _ = output_tx.try_send(format!("__exit:{}", code.unwrap_or(-1)));
});
let child_pid = pid_rx.recv().unwrap_or(0);
src/app/lsp/client.rs
+12 -4
View File
@@ -202,10 +202,18 @@ impl LspClient {
break;
}
if let Some(len_str) = trimmed.strip_prefix("Content-Length: ") {
content_length = Some(
len_str.trim().parse::<usize>()
.map_err(|e| anyhow::anyhow!("invalid Content-Length '{}': {}", len_str.trim(), e))?,
);
let length: usize = len_str.trim().parse::<usize>()
.map_err(|e| anyhow::anyhow!("invalid Content-Length '{}': {}", len_str.trim(), e))?;
// Cap Content-Length at 64 MiB to prevent OOM from a
// malicious or misconfigured LSP server (CWE-400).
const MAX_CONTENT_LENGTH: usize = 64 * 1024 * 1024;
if length > MAX_CONTENT_LENGTH {
anyhow::bail!(
"Content-Length {} exceeds maximum allowed size of {} bytes",
length, MAX_CONTENT_LENGTH,
);
}
content_length = Some(length);
}
}
src/app/mcp/manager.rs
+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;
src/app/runtime/actions/mod.rs
+5 -2
View File
@@ -514,7 +514,7 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
.filter_map(|m| m.content.as_deref())
.map(|c| c.len())
.sum();
let token_estimate = total_chars / 4;
let token_estimate = total_chars / 3;
rt.messages = crate::app::runtime::shortsend::shape_messages(&rt.messages, token_estimate, max_wire_tokens, true, None);
state.push_toast(Toast::new(ToastKind::Success, "Conversation history compacted.".to_string()));
state.dirty = true;
@@ -706,7 +706,10 @@ fn spawn_turn(state: &AppStateRest) {
let abort_flag = state.abort_flag.clone();
abort_flag.store(false, std::sync::atomic::Ordering::SeqCst);
*in_flight_flag.lock().unwrap() = true;
*in_flight_flag.lock().unwrap_or_else(|e| {
tracing::error!("[spawn_turn] in_flight_flag mutex poisoned: {}", e);
e.into_inner()
}) = true;
let events_q = turn_events.clone();
src/app/runtime/shortsend.rs
+42 -19
View File
@@ -7,15 +7,24 @@ use crate::dto::chat::message::ChatMessage;
/// sending to the LLM.
///
/// Flow: trigger based on token estimate. If `token_estimate` exceeds
/// `MAX_WIRE_TOKENS * 0.8`, we shape. We also apply hysteresis so it doesn't
/// flutter.
/// the threshold, we shape. When `prev_shaped` is true, the threshold is
/// raised (95%) to avoid fluttering — compaction only re-triggers when
/// the context is genuinely full again. When `prev_shaped` is false, the
/// threshold is lower (85%) so compaction starts proactively.
///
/// Why: hysteresis prevents repeated compaction on every turn when the
/// token count hovers near the boundary.
///
/// Return: `true` if shaping should be applied.
pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped: bool) -> bool {
let threshold = if prev_shaped {
(max_wire_tokens as f32 * 0.85) as usize
// Higher threshold when already shaped — defer re-shaping until
// the buffer is genuinely full again (95%).
(max_wire_tokens as f32 * 0.95) as usize
} else {
(max_wire_tokens as f32 * 0.90) as usize
// Lower threshold when not yet shaped — trigger shaping sooner
// (85%) to avoid hitting the context window limit.
(max_wire_tokens as f32 * 0.85) as usize
};
token_estimate >= threshold
}
@@ -59,7 +68,10 @@ pub fn shape_messages(
// Iterate backwards from the most recent to oldest
for m in msgs_to_eval.into_iter().rev() {
let text = m.content.as_deref().unwrap_or("");
let msg_tokens = text.len() / 4;
// Estimate tokens: ~1 token per 3 bytes for mixed content (code,
// prose, multi-byte). Conservative enough to stay under provider
// limits while avoiding premature compaction.
let msg_tokens = text.len() / 3;
if current_tokens + msg_tokens <= target_tokens {
current_tokens += msg_tokens;
@@ -80,22 +92,33 @@ pub fn shape_messages(
if !dropped_msgs.is_empty() {
let mut summary_text = "[prior conversation compacted]".to_string();
if let Some(llm) = client {
let prompt = format!(
"Summarize the following dropped conversation history briefly. Focus on main goals, decisions made, and files modified, so the context is preserved for future turns. Keep it concise.\n\nHistory:\n{}",
dropped_msgs.iter()
.map(|m| format!("[{}]: {}", if m.role == crate::dto::chat::message::Role::User { "User" } else { "Assistant" }, m.content.as_deref().unwrap_or("")))
.collect::<Vec<_>>()
.join("\n\n")
);
let req_msgs = vec![ChatMessage::user(prompt)];
if let Ok(resp) = llm.chat_with_tools_non_streaming(&req_msgs, None) {
if let Some(content) = resp.0.content {
summary_text = format!("[Summary of compacted prior conversation:\n{}\n]", content);
if let Some(llm) = client {
let prompt = format!(
"Summarize the following dropped conversation history briefly. Focus on main goals, decisions made, and files modified, so the context is preserved for future turns. Keep it concise.\n\nHistory:\n{}",
dropped_msgs.iter()
.map(|m| format!("[{}]: {}", if m.role == crate::dto::chat::message::Role::User { "User" } else { "Assistant" }, m.content.as_deref().unwrap_or("")))
.collect::<Vec<_>>()
.join("\n\n")
);
let req_msgs = vec![ChatMessage::user(prompt)];
match llm.chat_with_tools_non_streaming(&req_msgs, None) {
Ok(resp) => {
if let Some(content) = resp.0.content {
summary_text = format!("[Summary of compacted prior conversation:\n{}\n]", content);
}
}
Err(e) => {
tracing::warn!(
"[shortsend] LLM summarization failed: {}. \
Prior conversation history is lost — no summary available. \
This means the model will lose context about earlier parts of \
the conversation.",
e,
);
}
}
}
}
result.push(ChatMessage::system(summary_text));
}
src/app/runtime/stream/mod.rs
+20 -5
View File
@@ -65,10 +65,11 @@ impl SseParser {
events.extend(self.flush_event());
} else if let Some(ty) = line.strip_prefix("event: ") {
self.event_type = Some(ty.trim().to_string());
} else if let Some(data) = line.strip_prefix("data: ") {
self.data_lines.push(data.to_string());
} else if line.starts_with("data:") {
self.data_lines.push(String::new());
} else if let Some(data) = line.strip_prefix("data:") {
// Handle both "data: {...}" (with space) and "data:{...}"
// (without space). Some providers omit the trailing space.
let data = data.trim_start().to_string();
self.data_lines.push(data);
}
}
events
@@ -119,7 +120,21 @@ impl SseParser {
tracing::warn!("[stream] total_tokens missing in usage chunk");
prompt_tokens + completion_tokens
});
return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }];
// Only emit Usage as a standalone event if this chunk
// contains nothing else (no choices, no delta). Some
// non-standard providers may bundle usage WITH content
// in the same chunk; emitting both prevents content loss.
let has_other_content = value.get("choices")
.and_then(|c| c.as_array())
.map(|arr| arr.iter().any(|ch| {
ch.get("delta").and_then(|d| d.get("content")).is_some()
|| ch.get("delta").and_then(|d| d.get("reasoning_content")).is_some()
|| ch.get("delta").and_then(|d| d.get("tool_calls")).is_some()
}))
.unwrap_or(false);
if !has_other_content {
return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }];
}
}
}
match event_type.as_str() {
src/app/runtime/stream/turn.rs
+9 -1
View File
@@ -118,7 +118,15 @@ impl StreamedTurn {
.filter(|tc| !tc.name.is_empty())
.map(|tc| {
let args_value: serde_json::Value = serde_json::from_str(&tc.arguments)
.unwrap_or(serde_json::Value::String(tc.arguments.clone()));
.unwrap_or_else(|e| {
tracing::warn!(
"[stream] tool call '{}' has invalid JSON arguments: {} — \
arguments will be double-stringified, which may cause \
tool execution to fail",
tc.name, e,
);
serde_json::Value::String(tc.arguments.clone())
});
ToolCall {
id: tc.id.clone(),
type_: "function".to_string(),