refactor: apply 5-pass audit optimizations across mcp-memory codebase
This commit is contained in:
Riz Ashraf committed
1 parent
924b6d09fa
commit
5bd8b1587a
43 files changed
+1866
-1658
No files matched your search
@@ -30,33 +30,56 @@ impl McpTool for WriteClipboardHandler {
|
||||
let res = tokio::task::spawn_blocking(move || {
|
||||
let mut msgs = Vec::new();
|
||||
|
||||
if let Ok(mut clipboard) = Clipboard::new()
|
||||
&& let Some(text) = &tool_args.text
|
||||
&& clipboard.set_text(text).is_ok()
|
||||
{
|
||||
msgs.push("Wrote text");
|
||||
if let Some(text) = &tool_args.text {
|
||||
let mut written = false;
|
||||
if let Ok(mut clipboard) = Clipboard::new() {
|
||||
for _ in 0..3 {
|
||||
if clipboard.set_text(text).is_ok() {
|
||||
written = true;
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
}
|
||||
if written {
|
||||
msgs.push("Wrote text");
|
||||
}
|
||||
}
|
||||
|
||||
// Handle arboard for image
|
||||
if let Some(image_path) = &tool_args.image_path {
|
||||
match image::open(image_path) {
|
||||
Ok(img) => {
|
||||
let img = if img.width() > 2048 || img.height() > 2048 {
|
||||
img.resize(2048, 2048, FilterType::Triangle)
|
||||
} else {
|
||||
img
|
||||
};
|
||||
let rgba = img.into_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let img_data = ImageData {
|
||||
width: w as usize,
|
||||
height: h as usize,
|
||||
bytes: Cow::Owned(rgba.into_raw()),
|
||||
};
|
||||
let raw_bytes = rgba.into_raw();
|
||||
let mut written = false;
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if clipboard.set_image(img_data).is_ok() {
|
||||
msgs.push("Wrote image");
|
||||
} else {
|
||||
return Err(crate::error::AppError::Internal(
|
||||
"Failed to write image to clipboard".to_string(),
|
||||
));
|
||||
for _ in 0..3 {
|
||||
let img_data = ImageData {
|
||||
width: w as usize,
|
||||
height: h as usize,
|
||||
bytes: Cow::Borrowed(&raw_bytes),
|
||||
};
|
||||
if clipboard.set_image(img_data).is_ok() {
|
||||
written = true;
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
}
|
||||
if written {
|
||||
msgs.push("Wrote image");
|
||||
} else {
|
||||
return Err(crate::error::AppError::Internal(
|
||||
"Failed to write image to clipboard".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
@@ -85,92 +108,34 @@ impl McpTool for WriteClipboardHandler {
|
||||
}
|
||||
|
||||
pub fn get_native_clipboard_text() -> Option<String> {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(text) = clipboard.get_text() {
|
||||
if !text.trim().is_empty() {
|
||||
return Some(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut cmd_wl = std::process::Command::new("wl-paste");
|
||||
cmd_wl.arg("--no-newline");
|
||||
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
|
||||
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
|
||||
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
|
||||
}
|
||||
if let Ok(output) = cmd_wl.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(text) = String::from_utf8(output.stdout) {
|
||||
for _ in 0..3 {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(text) = clipboard.get_text() {
|
||||
if !text.trim().is_empty() {
|
||||
return Some(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
|
||||
let mut cmd_xc = std::process::Command::new("xclip");
|
||||
cmd_xc.args(["-selection", "clipboard", "-o"]);
|
||||
if std::env::var("DISPLAY").is_err() {
|
||||
cmd_xc.env("DISPLAY", ":0");
|
||||
}
|
||||
if let Ok(output) = cmd_xc.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(text) = String::from_utf8(output.stdout) {
|
||||
if !text.trim().is_empty() {
|
||||
return Some(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(image_data) = clipboard.get_image() {
|
||||
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
|
||||
image_data.width as u32,
|
||||
image_data.height as u32,
|
||||
image_data.bytes.into_owned(),
|
||||
) {
|
||||
return Some(image::DynamicImage::ImageRgba8(img));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for mime in &["image/png", "image/jpeg", "image/bmp", "image/tiff"] {
|
||||
let mut cmd_wl = std::process::Command::new("wl-paste");
|
||||
cmd_wl.args(["--type", mime]);
|
||||
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
|
||||
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
|
||||
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
|
||||
}
|
||||
if let Ok(output) = cmd_wl.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(img) = image::load_from_memory(&output.stdout) {
|
||||
return Some(img);
|
||||
for _ in 0..3 {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(image_data) = clipboard.get_image() {
|
||||
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
|
||||
image_data.width as u32,
|
||||
image_data.height as u32,
|
||||
image_data.bytes.into_owned(),
|
||||
) {
|
||||
return Some(image::DynamicImage::ImageRgba8(img));
|
||||
}
|
||||
}
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
|
||||
for mime in &["image/png", "image/jpeg", "image/bmp"] {
|
||||
let mut cmd_xc = std::process::Command::new("xclip");
|
||||
cmd_xc.args(["-selection", "clipboard", "-t", mime, "-o"]);
|
||||
if std::env::var("DISPLAY").is_err() {
|
||||
cmd_xc.env("DISPLAY", ":0");
|
||||
}
|
||||
if let Ok(output) = cmd_xc.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(img) = image::load_from_memory(&output.stdout) {
|
||||
return Some(img);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
@@ -422,5 +387,33 @@ mod tests {
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
assert!(parsed.is_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_subprocess_clipboard_regression() {
|
||||
let vision_src = include_str!("vision.rs");
|
||||
let code_only = vision_src.split("mod tests").next().unwrap_or(vision_src);
|
||||
let forbidden_cmd = format!("Command::{}{}", "n", "ew");
|
||||
let forbidden_ps = format!("power{}", "shell");
|
||||
let forbidden_wl = format!("wl-{}", "paste");
|
||||
let forbidden_xc = format!("x{}", "clip");
|
||||
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_cmd),
|
||||
"Regression detected: vision.rs must not spawn subprocesses!"
|
||||
);
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_ps),
|
||||
"Regression detected: vision.rs must not invoke powershell!"
|
||||
);
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_wl),
|
||||
"Regression detected: vision.rs must not invoke wl-paste!"
|
||||
);
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_xc),
|
||||
"Regression detected: vision.rs must not invoke xclip!"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user