diff --git a/crates/fetchkit-cli/src/main.rs b/crates/fetchkit-cli/src/main.rs index 6bf0530..734339f 100644 --- a/crates/fetchkit-cli/src/main.rs +++ b/crates/fetchkit-cli/src/main.rs @@ -234,25 +234,29 @@ fn print_md_with_frontmatter(response: &fetchkit::FetchResponse) { writeln_safe(&format_md_with_frontmatter(response)); } +fn yaml_quote(value: &str) -> String { + serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string()) +} + /// Format response as markdown with YAML frontmatter fn format_md_with_frontmatter(response: &fetchkit::FetchResponse) -> String { let mut output = String::new(); // Build frontmatter output.push_str("---\n"); - output.push_str(&format!("url: {}\n", response.url)); + output.push_str(&format!("url: {}\n", yaml_quote(&response.url))); output.push_str(&format!("status_code: {}\n", response.status_code)); if let Some(ref ct) = response.content_type { - output.push_str(&format!("source_content_type: {}\n", ct)); + output.push_str(&format!("source_content_type: {}\n", yaml_quote(ct))); } if let Some(size) = response.size { output.push_str(&format!("source_size: {}\n", size)); } if let Some(ref lm) = response.last_modified { - output.push_str(&format!("last_modified: {}\n", lm)); + output.push_str(&format!("last_modified: {}\n", yaml_quote(lm))); } if let Some(ref filename) = response.filename { - output.push_str(&format!("filename: {}\n", filename)); + output.push_str(&format!("filename: {}\n", yaml_quote(filename))); } if let Some(truncated) = response.truncated { if truncated { @@ -302,9 +306,9 @@ mod tests { let output = format_md_with_frontmatter(&response); assert!(output.starts_with("---\n")); - assert!(output.contains("url: https://example.com\n")); + assert!(output.contains("url: \"https://example.com\"\n")); assert!(output.contains("status_code: 200\n")); - assert!(output.contains("source_content_type: text/html\n")); + assert!(output.contains("source_content_type: \"text/html\"\n")); assert!(output.contains("---\n# Hello World")); } @@ -325,8 +329,8 @@ mod tests { let output = format_md_with_frontmatter(&response); assert!(output.contains("source_size: 1234\n")); - assert!(output.contains("last_modified: Wed, 01 Jan 2025 00:00:00 GMT\n")); - assert!(output.contains("filename: page.html\n")); + assert!(output.contains("last_modified: \"Wed, 01 Jan 2025 00:00:00 GMT\"\n")); + assert!(output.contains("filename: \"page.html\"\n")); assert!(output.contains("truncated: true\n")); } @@ -362,4 +366,21 @@ mod tests { // truncated: false should not appear assert!(!output.contains("truncated")); } + + #[test] + fn test_format_md_quotes_untrusted_scalars() { + let response = FetchResponse { + url: "https://example.com/a\nforged: true".to_string(), + status_code: 200, + filename: Some("*alias".to_string()), + content: Some("ok".to_string()), + ..Default::default() + }; + + let output = format_md_with_frontmatter(&response); + + assert!(output.contains("url: \"https://example.com/a\\nforged: true\"\n")); + assert!(output.contains("filename: \"*alias\"\n")); + assert!(!output.contains("\nforged: true\n")); + } } diff --git a/crates/fetchkit-cli/src/mcp.rs b/crates/fetchkit-cli/src/mcp.rs index 18f520d..783382c 100644 --- a/crates/fetchkit-cli/src/mcp.rs +++ b/crates/fetchkit-cli/src/mcp.rs @@ -197,19 +197,19 @@ fn format_md_with_frontmatter(response: &fetchkit::FetchResponse) -> String { // Build frontmatter output.push_str("---\n"); - output.push_str(&format!("url: {}\n", response.url)); + output.push_str(&format!("url: {}\n", yaml_quote(&response.url))); output.push_str(&format!("status_code: {}\n", response.status_code)); if let Some(ref ct) = response.content_type { - output.push_str(&format!("source_content_type: {}\n", ct)); + output.push_str(&format!("source_content_type: {}\n", yaml_quote(ct))); } if let Some(size) = response.size { output.push_str(&format!("source_size: {}\n", size)); } if let Some(ref lm) = response.last_modified { - output.push_str(&format!("last_modified: {}\n", lm)); + output.push_str(&format!("last_modified: {}\n", yaml_quote(lm))); } if let Some(ref filename) = response.filename { - output.push_str(&format!("filename: {}\n", filename)); + output.push_str(&format!("filename: {}\n", yaml_quote(filename))); } if let Some(truncated) = response.truncated { if truncated { @@ -228,6 +228,10 @@ fn format_md_with_frontmatter(response: &fetchkit::FetchResponse) -> String { output } +fn yaml_quote(value: &str) -> String { + serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string()) +} + /// Run the MCP server over stdio pub async fn run_server(tool: Tool) { let server = McpServer::new(tool); diff --git a/crates/fetchkit-cli/tests/cli_integration.rs b/crates/fetchkit-cli/tests/cli_integration.rs index 136bde1..ffdb45d 100644 --- a/crates/fetchkit-cli/tests/cli_integration.rs +++ b/crates/fetchkit-cli/tests/cli_integration.rs @@ -38,7 +38,7 @@ fn test_fetch_markdown_output() { "Expected YAML frontmatter, got: {}", &stdout[..80.min(stdout.len())] ); - assert!(stdout.contains("url: https://example.com")); + assert!(stdout.contains("url: \"https://example.com")); assert!(stdout.contains("status_code: 200")); assert!(stdout.contains("source_content_type:"));