862fe9b395
Add comprehensive JSON serialization and validation to RawTimingMetrics: - Add serde::Serialize/Deserialize derives to RawTimingMetrics - Implement validate() method checking required metrics (runtime, throughput, file counts) - Add to_json() and from_json() for import/export - Implement store_temporary() for JSON file storage (benches/results/raw_metrics_<timestamp>.json) - Integrate validation and temporary storage into benchmark flow - Add validation for NaN/infinity values in floating-point metrics Closes bf-4b7pm. Verification: notes/bf-4b7pm.md, commit $(git rev-parse --short HEAD). Tests: PASS (cargo check --bench grep_1000).
2217 lines
77 KiB
Rust
2217 lines
77 KiB
Rust
#![cfg(feature = "remote")]
|
|
//! TH-05: SSRF protection tests (Phase 1.8).
|
|
//!
|
|
//! This test suite exercises SSRF payloads against the remote-source fetcher
|
|
//! and the MCP extract tool. It asserts that dangerous URLs are refused with
|
|
//! the URL_PRIVATE_NETWORK diagnostic.
|
|
//!
|
|
//! Test categories:
|
|
//! - Cloud metadata endpoints (AWS, GCP, Azure, Alibaba)
|
|
//! - RFC 1918 private IPv4 ranges
|
|
//! - Loopback addresses
|
|
//! - Link-local addresses
|
|
//! - IPv6 ULA and loopback
|
|
//! - Non-https schemes (http, ftp, file)
|
|
//!
|
|
//! Each payload is tested against:
|
|
//! - CLI: `pdftract extract --url <payload>`
|
|
//! - MCP: extract tool with URL parameter
|
|
//! - Serve: POST /extract with URL
|
|
//!
|
|
//! With --allow-private-networks set, the same URLs are accepted.
|
|
|
|
use pdftract_core::diagnostics::DiagCode;
|
|
use pdftract_core::url_validation::{validate_url, UrlValidationError};
|
|
|
|
/// Test payload categories for SSRF protection.
|
|
struct TestPayload {
|
|
/// The URL to test
|
|
url: &'static str,
|
|
/// Expected error variant
|
|
expected_error: ExpectedError,
|
|
/// Description of what this tests
|
|
description: &'static str,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
enum ExpectedError {
|
|
InvalidScheme,
|
|
PrivateNetwork,
|
|
DnsFailed,
|
|
}
|
|
|
|
impl ExpectedError {
|
|
fn matches(&self, err: &UrlValidationError) -> bool {
|
|
match (self, err) {
|
|
(ExpectedError::InvalidScheme, UrlValidationError::InvalidScheme(_)) => true,
|
|
(ExpectedError::PrivateNetwork, UrlValidationError::PrivateNetwork(_)) => true,
|
|
(ExpectedError::DnsFailed, UrlValidationError::DnsFailed(_)) => true,
|
|
_ => false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// SSRF test payloads covering all dangerous categories.
|
|
const SSRF_PAYLOADS: &[TestPayload] = &[
|
|
// === Cloud metadata endpoints ===
|
|
TestPayload {
|
|
url: "https://169.254.169.254/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "AWS metadata endpoint (169.254.169.254)",
|
|
},
|
|
TestPayload {
|
|
url: "https://169.254.169.254/latest/meta-data/identity-credentials/ec2/security-credentials/ec2-instance",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "AWS metadata endpoint (full path)",
|
|
},
|
|
TestPayload {
|
|
url: "https://metadata.google.internal/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "GCP metadata endpoint (hostname)",
|
|
},
|
|
TestPayload {
|
|
url: "https://instance-data.google.internal/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "GCP instance metadata endpoint",
|
|
},
|
|
TestPayload {
|
|
url: "https://168.63.129.16/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "Azure metadata endpoint (168.63.129.16)",
|
|
},
|
|
TestPayload {
|
|
url: "https://100.100.100.200/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "Alibaba metadata endpoint (100.100.100.200)",
|
|
},
|
|
|
|
// === RFC 1918 private IPv4 ranges ===
|
|
TestPayload {
|
|
url: "https://10.0.0.1/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "RFC 1918: 10.0.0.0/8 (lower bound)",
|
|
},
|
|
TestPayload {
|
|
url: "https://10.255.255.255/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "RFC 1918: 10.0.0.0/8 (upper bound)",
|
|
},
|
|
TestPayload {
|
|
url: "https://172.16.0.1/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "RFC 1918: 172.16.0.0/12 (lower bound)",
|
|
},
|
|
TestPayload {
|
|
url: "https://172.31.255.255/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "RFC 1918: 172.16.0.0/12 (upper bound)",
|
|
},
|
|
TestPayload {
|
|
url: "https://192.168.1.1/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "RFC 1918: 192.168.0.0/16",
|
|
},
|
|
TestPayload {
|
|
url: "https://192.168.255.255/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "RFC 1918: 192.168.0.0/16 (upper bound)",
|
|
},
|
|
|
|
// === Loopback addresses ===
|
|
TestPayload {
|
|
url: "https://127.0.0.1/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "Loopback: 127.0.0.1",
|
|
},
|
|
TestPayload {
|
|
url: "https://127.0.0.2/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "Loopback: 127.0.0.2",
|
|
},
|
|
TestPayload {
|
|
url: "https://127.255.255.255/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "Loopback: 127.255.255.255",
|
|
},
|
|
|
|
// === Link-local addresses ===
|
|
TestPayload {
|
|
url: "https://169.254.0.1/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "IPv4 link-local: 169.254.0.1",
|
|
},
|
|
|
|
// === IPv6 ULA ===
|
|
TestPayload {
|
|
url: "https://[fd00::1]/",
|
|
expected_error: ExpectedError::PrivateNetwork, // IPv6 ULA is detected as private
|
|
description: "IPv6 ULA: fd00::1",
|
|
},
|
|
TestPayload {
|
|
url: "https://[fc00::1]/",
|
|
expected_error: ExpectedError::PrivateNetwork, // IPv6 ULA is detected as private
|
|
description: "IPv6 ULA: fc00::1",
|
|
},
|
|
|
|
// === IPv6 loopback ===
|
|
TestPayload {
|
|
url: "https://[::1]/",
|
|
expected_error: ExpectedError::PrivateNetwork,
|
|
description: "IPv6 loopback: ::1",
|
|
},
|
|
|
|
// === IPv6 link-local ===
|
|
TestPayload {
|
|
url: "https://[fe80::1]/",
|
|
expected_error: ExpectedError::PrivateNetwork, // IPv6 link-local is detected as private
|
|
description: "IPv6 link-local: fe80::1",
|
|
},
|
|
|
|
// === Non-https schemes ===
|
|
TestPayload {
|
|
url: "http://example.com/",
|
|
expected_error: ExpectedError::InvalidScheme,
|
|
description: "HTTP scheme (not https)",
|
|
},
|
|
TestPayload {
|
|
url: "ftp://example.com/",
|
|
expected_error: ExpectedError::InvalidScheme,
|
|
description: "FTP scheme",
|
|
},
|
|
TestPayload {
|
|
url: "file:///etc/passwd",
|
|
expected_error: ExpectedError::InvalidScheme,
|
|
description: "file:// scheme",
|
|
},
|
|
];
|
|
|
|
/// Public URLs that should be accepted (positive test).
|
|
const PUBLIC_URLS: &[&str] = &[
|
|
"https://example.com/",
|
|
"https://www.google.com/",
|
|
"https://github.com/",
|
|
"https://8.8.8.8/", // Public DNS
|
|
"https://1.1.1.1/", // Cloudflare DNS
|
|
];
|
|
|
|
#[test]
|
|
fn test_ssrf_protection_blocks_all_dangerous_payloads() {
|
|
for payload in SSRF_PAYLOADS {
|
|
let result = validate_url(payload.url, false);
|
|
|
|
assert!(
|
|
result.is_err(),
|
|
"URL should be rejected: {} ({})",
|
|
payload.url,
|
|
payload.description
|
|
);
|
|
|
|
let err = result.unwrap_err();
|
|
assert!(
|
|
payload.expected_error.matches(&err),
|
|
"URL '{}' ({}) expected {:?}, got {:?}",
|
|
payload.url,
|
|
payload.description,
|
|
payload.expected_error,
|
|
err
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_allow_private_networks_bypass() {
|
|
for payload in SSRF_PAYLOADS {
|
|
// Skip scheme validation tests (those should always fail)
|
|
if matches!(payload.expected_error, ExpectedError::InvalidScheme) {
|
|
continue;
|
|
}
|
|
|
|
// Skip metadata endpoint tests (those should always fail for security)
|
|
if payload.description.contains("metadata") {
|
|
continue;
|
|
}
|
|
|
|
// With --allow-private-networks, private network URLs are accepted
|
|
let result = validate_url(payload.url, true);
|
|
|
|
match result {
|
|
Ok(_) => {
|
|
// URL is now accepted
|
|
}
|
|
Err(UrlValidationError::DnsFailed(_)) => {
|
|
// DNS resolution failure is OK in tests (no network)
|
|
}
|
|
Err(other) => {
|
|
panic!(
|
|
"URL '{}' ({}) should be accepted with --allow-private-networks, got: {:?}",
|
|
payload.url, payload.description, other
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_public_urls_are_accepted() {
|
|
for url in PUBLIC_URLS {
|
|
// Note: These may fail with DnsFailed in offline test environments
|
|
let result = validate_url(url, false);
|
|
|
|
match result {
|
|
Ok(_) => {
|
|
// URL accepted
|
|
}
|
|
Err(UrlValidationError::DnsFailed(_)) => {
|
|
// OK in offline tests
|
|
}
|
|
Err(other) => {
|
|
panic!("Public URL '{}' should be accepted, got: {:?}", url, other);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_http_scheme_always_rejected() {
|
|
// Even with --allow-private-networks, http:// is rejected
|
|
let result = validate_url("http://127.0.0.1/", true);
|
|
assert!(matches!(result, Err(UrlValidationError::InvalidScheme(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_file_scheme_always_rejected() {
|
|
let result = validate_url("file:///etc/passwd", true);
|
|
assert!(matches!(result, Err(UrlValidationError::InvalidScheme(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_ftp_scheme_always_rejected() {
|
|
let result = validate_url("ftp://example.com/", true);
|
|
assert!(matches!(result, Err(UrlValidationError::InvalidScheme(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_url_with_basic_auth_rejected() {
|
|
// URLs with embedded credentials should still be checked by host, not credentials
|
|
let result = validate_url("https://user:pass@127.0.0.1/", false);
|
|
assert!(matches!(result, Err(UrlValidationError::PrivateNetwork(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_ipv6_zone_id_detected_as_link_local() {
|
|
// IPv6 zone IDs indicate link-local addresses
|
|
let result = validate_url("https://[fe80::1%eth0]/", false);
|
|
assert!(matches!(result, Err(UrlValidationError::PrivateNetwork(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_metadata_subdomain_detected() {
|
|
// Subdomains of metadata endpoints should also be blocked
|
|
let result = validate_url("https://foo.metadata.google.internal/", false);
|
|
assert!(matches!(result, Err(UrlValidationError::PrivateNetwork(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_url_validation_returns_correct_diagnostic_code() {
|
|
use pdftract_core::url_validation::validate_url_with_diagnostic;
|
|
|
|
let result = validate_url_with_diagnostic("https://127.0.0.1/", false);
|
|
assert!(result.is_err());
|
|
let diag = result.unwrap_err();
|
|
assert_eq!(diag.code, DiagCode::RemoteUrlPrivateNetwork);
|
|
}
|
|
|
|
#[test]
|
|
fn test_private_ipv4_boundary_addresses() {
|
|
// Test addresses just outside the private ranges
|
|
let public_addrs = &[
|
|
"172.15.255.255", // Just below 172.16.0.0/12
|
|
"172.32.0.1", // Just above 172.16.0.0/12
|
|
"192.167.255.255", // Just below 192.168.0.0/16
|
|
"192.169.0.1", // Just above 192.168.0.0/16
|
|
];
|
|
|
|
for addr in public_addrs {
|
|
let url = format!("https://{}/", addr);
|
|
let result = validate_url(&url, false);
|
|
|
|
// These should not be rejected as private network (may fail DNS in tests)
|
|
match result {
|
|
Ok(_) => {}
|
|
Err(UrlValidationError::DnsFailed(_)) => {}
|
|
Err(UrlValidationError::PrivateNetwork(msg)) => {
|
|
panic!(
|
|
"Public address {} should not be rejected as private: {}",
|
|
addr, msg
|
|
);
|
|
}
|
|
Err(_) => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_current_network_range_blocked() {
|
|
// 0.0.0.0/8 (current network) should be blocked
|
|
let result = validate_url("https://0.0.0.0/", false);
|
|
assert!(matches!(result, Err(UrlValidationError::PrivateNetwork(_))));
|
|
|
|
let result = validate_url("https://0.0.0.8/", false);
|
|
assert!(matches!(result, Err(UrlValidationError::PrivateNetwork(_))));
|
|
}
|
|
|
|
// ============================================================================
|
|
// MCP JSON-RPC Message Construction Helpers
|
|
// ============================================================================
|
|
|
|
/// Helper module for constructing JSON-RPC MCP tool call messages.
|
|
///
|
|
/// This module provides type-safe helpers for building MCP tools/call requests,
|
|
/// particularly for SSRF testing. It uses the JSON-RPC framing types from
|
|
/// pdftract_cli::mcp::framing and provides convenient constructors for
|
|
/// common tool call patterns.
|
|
///
|
|
/// ## Example
|
|
///
|
|
/// ```rust
|
|
/// use mcp_helpers::ToolCallBuilder;
|
|
///
|
|
/// let request = ToolCallBuilder::extract()
|
|
/// .with_url("https://example.com/doc.pdf")
|
|
/// .build();
|
|
/// ```
|
|
#[cfg(feature = "remote")]
|
|
#[cfg(test)]
|
|
pub mod mcp_helpers {
|
|
use std::io::{BufRead, BufReader};
|
|
use serde_json::json;
|
|
|
|
// ============================================================================
|
|
// Local JSON-RPC Type Definitions (to avoid cross-crate imports)
|
|
// ============================================================================
|
|
|
|
/// A JSON-RPC request identifier.
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum Id {
|
|
Number(i64),
|
|
String(String),
|
|
Null,
|
|
}
|
|
|
|
impl serde::Serialize for Id {
|
|
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
|
where
|
|
S: serde::Serializer,
|
|
{
|
|
match self {
|
|
Id::Number(n) => n.serialize(serializer),
|
|
Id::String(s) => s.serialize(serializer),
|
|
Id::Null => serializer.serialize_none(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A JSON-RPC response identifier.
|
|
///
|
|
/// Can be a number, string, or null (for notifications).
|
|
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize, PartialEq)]
|
|
#[serde(untagged)]
|
|
pub enum ResponseId {
|
|
Number(i64),
|
|
String(String),
|
|
Null,
|
|
}
|
|
|
|
/// A JSON-RPC request object.
|
|
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize)]
|
|
pub struct Request {
|
|
#[serde(rename = "jsonrpc")]
|
|
pub jsonrpc_version: String,
|
|
pub method: String,
|
|
pub params: Option<serde_json::Value>,
|
|
pub id: Option<Id>,
|
|
}
|
|
|
|
impl Request {
|
|
/// Create a new request with the given method and optional params.
|
|
pub fn new(method: impl Into<String>, params: Option<serde_json::Value>, id: Option<Id>) -> Self {
|
|
Self {
|
|
jsonrpc_version: "2.0".to_string(),
|
|
method: method.into(),
|
|
params,
|
|
id,
|
|
}
|
|
}
|
|
|
|
/// Returns true if this is a notification (no id field).
|
|
pub fn is_notification(&self) -> bool {
|
|
self.id.is_none()
|
|
}
|
|
|
|
/// Get the request ID, or Id::Null for notifications.
|
|
pub fn request_id(&self) -> Id {
|
|
self.id.clone().unwrap_or(Id::Null)
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tool Call Builder
|
|
// ============================================================================
|
|
|
|
/// Builder for constructing MCP tools/call JSON-RPC requests.
|
|
///
|
|
/// Provides a fluent API for building tool call requests with proper
|
|
/// JSON-RPC structure and type-safe parameter construction.
|
|
pub struct ToolCallBuilder {
|
|
tool_name: String,
|
|
arguments: serde_json::Map<String, serde_json::Value>,
|
|
request_id: Id,
|
|
}
|
|
|
|
impl ToolCallBuilder {
|
|
/// Create a new ToolCallBuilder for the specified tool name.
|
|
fn new(tool_name: impl Into<String>) -> Self {
|
|
Self {
|
|
tool_name: tool_name.into(),
|
|
arguments: serde_json::Map::new(),
|
|
request_id: Id::Number(1),
|
|
}
|
|
}
|
|
|
|
/// Create a builder for the "extract" tool.
|
|
///
|
|
/// This is the primary tool for PDF extraction and accepts URL
|
|
/// parameters that must be validated for SSRF protection.
|
|
pub fn extract() -> Self {
|
|
Self::new("extract")
|
|
}
|
|
|
|
/// Create a builder for the "get_metadata" tool.
|
|
///
|
|
/// This tool fetches document metadata and also accepts URL parameters.
|
|
pub fn get_metadata() -> Self {
|
|
Self::new("get_metadata")
|
|
}
|
|
|
|
/// Set the request ID.
|
|
///
|
|
/// Default is Id::Number(1). This allows customizing the ID for
|
|
/// concurrent request tracking.
|
|
pub fn with_id(mut self, id: Id) -> Self {
|
|
self.request_id = id;
|
|
self
|
|
}
|
|
|
|
/// Add a URL parameter to the tool arguments.
|
|
///
|
|
/// This is the primary parameter for extract and get_metadata tools
|
|
/// when fetching remote PDFs.
|
|
pub fn with_url(mut self, url: impl Into<String>) -> Self {
|
|
self.arguments
|
|
.insert("path".to_string(), serde_json::Value::String(url.into()));
|
|
self
|
|
}
|
|
|
|
/// Add a custom argument to the tool arguments.
|
|
///
|
|
/// Allows adding arbitrary parameters like "password", "ocr", etc.
|
|
pub fn with_argument(
|
|
mut self,
|
|
key: impl Into<String>,
|
|
value: serde_json::Value,
|
|
) -> Self {
|
|
self.arguments.insert(key.into(), value);
|
|
self
|
|
}
|
|
|
|
/// Build the JSON-RPC request.
|
|
///
|
|
/// Returns a Request object that can be serialized to JSON and sent
|
|
/// to the MCP server via stdio or HTTP transport.
|
|
pub fn build(self) -> Request {
|
|
let params = json!({
|
|
"name": self.tool_name,
|
|
"arguments": self.arguments
|
|
});
|
|
|
|
Request::new("tools/call", Some(params), Some(self.request_id))
|
|
}
|
|
|
|
/// Build the request and serialize to JSON string.
|
|
///
|
|
/// Convenience method that combines build() and serde_json::to_string().
|
|
pub fn build_json(self) -> String {
|
|
let request = self.build();
|
|
// Manual serialization to include jsonrpc field
|
|
let id_value = match request.id {
|
|
Some(Id::Number(n)) => serde_json::json!(n),
|
|
Some(Id::String(s)) => serde_json::json!(s),
|
|
Some(Id::Null) => serde_json::json!(null),
|
|
None => serde_json::json!(null),
|
|
};
|
|
|
|
let output = json!({
|
|
"jsonrpc": "2.0",
|
|
"method": request.method,
|
|
"params": request.params,
|
|
"id": id_value
|
|
});
|
|
|
|
serde_json::to_string(&output)
|
|
.expect("ToolCallRequest should always be serializable")
|
|
}
|
|
}
|
|
|
|
/// Quick helper to create an extract tool call with just a URL.
|
|
///
|
|
/// This is the most common pattern for SSRF testing.
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust
|
|
/// let request = extract_call("https://127.0.0.1/");
|
|
/// ```
|
|
pub fn extract_call(url: impl Into<String>) -> Request {
|
|
ToolCallBuilder::extract().with_url(url).build()
|
|
}
|
|
|
|
/// Quick helper to create a get_metadata tool call with just a URL.
|
|
pub fn get_metadata_call(url: impl Into<String>) -> Request {
|
|
ToolCallBuilder::get_metadata().with_url(url).build()
|
|
}
|
|
|
|
// ============================================================================
|
|
// JSON-RPC Response Parsing Helpers
|
|
// ============================================================================
|
|
|
|
/// JSON-RPC 2.0 error object structure.
|
|
///
|
|
/// Represents the error field in a JSON-RPC response with code, message,
|
|
/// and optional data fields.
|
|
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
|
|
pub struct JsonRpcError {
|
|
/// The error code (negative for server errors in the -32099..-32000 range)
|
|
pub code: i64,
|
|
/// Human-readable error message
|
|
pub message: String,
|
|
/// Optional additional error data
|
|
pub data: Option<serde_json::Value>,
|
|
}
|
|
|
|
impl JsonRpcError {
|
|
/// Check if this error is an SSRF_BLOCKED error.
|
|
///
|
|
/// Returns true if the error data contains a "code" field with value "SSRF_BLOCKED"
|
|
/// or if the error message contains the substring "SSRF_BLOCKED".
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust
|
|
/// use mcp_helpers::JsonRpcError;
|
|
/// use serde_json::json;
|
|
///
|
|
/// let error = JsonRpcError {
|
|
/// code: -32001,
|
|
/// message: "SSRF_BLOCKED: URL targets private network".to_string(),
|
|
/// data: Some(json!({"code": "SSRF_BLOCKED"})),
|
|
/// };
|
|
/// assert!(error.is_ssrf_blocked());
|
|
///
|
|
/// let error2 = JsonRpcError {
|
|
/// code: -32601,
|
|
/// message: "Method not found".to_string(),
|
|
/// data: None,
|
|
/// };
|
|
/// assert!(!error2.is_ssrf_blocked());
|
|
/// ```
|
|
pub fn is_ssrf_blocked(&self) -> bool {
|
|
// Check if error data contains "code": "SSRF_BLOCKED"
|
|
if let Some(data) = &self.data {
|
|
if let Some(code) = data.get("code").and_then(|c| c.as_str()) {
|
|
if code == "SSRF_BLOCKED" {
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check if the error message itself contains SSRF_BLOCKED
|
|
if self.message.contains("SSRF_BLOCKED") {
|
|
return true;
|
|
}
|
|
|
|
false
|
|
}
|
|
}
|
|
|
|
/// A generic JSON-RPC 2.0 response structure.
|
|
///
|
|
/// A response has either a result field (success) or an error field (failure),
|
|
/// never both. The id field must match the request id.
|
|
///
|
|
/// # Type Parameters
|
|
///
|
|
/// * `T` - The type of the result field for successful responses. Must implement
|
|
/// `serde::Deserialize` to parse from JSON.
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust
|
|
/// use mcp_helpers::{JsonRpcResponse, JsonRpcError};
|
|
///
|
|
/// // A successful response with a custom result type
|
|
/// let success_json = r#"{"jsonrpc":"2.0","result":{"status":"ok"},"id":1}"#;
|
|
/// let response: JsonRpcResponse<serde_json::Value> = serde_json::from_str(success_json).unwrap();
|
|
/// assert!(response.is_success());
|
|
/// assert_eq!(response.jsonrpc, "2.0");
|
|
///
|
|
/// // An error response
|
|
/// let error_json = r#"{"jsonrpc":"2.0","error":{"code":-32001,"message":"Error"},"id":1}"#;
|
|
/// let response: JsonRpcResponse<serde_json::Value> = serde_json::from_str(error_json).unwrap();
|
|
/// assert!(response.is_error());
|
|
/// ```
|
|
#[derive(Debug, Clone, serde::Deserialize)]
|
|
pub struct JsonRpcResponse<T> {
|
|
/// Must be exactly "2.0"
|
|
pub jsonrpc: String,
|
|
/// The successful result value (present only on success)
|
|
pub result: Option<T>,
|
|
/// The error object (present only on failure)
|
|
pub error: Option<JsonRpcError>,
|
|
/// Request identifier
|
|
pub id: ResponseId,
|
|
}
|
|
|
|
impl<T> JsonRpcResponse<T> {
|
|
/// Check if this is a successful response (has result field).
|
|
pub fn is_success(&self) -> bool {
|
|
self.result.is_some()
|
|
}
|
|
|
|
/// Check if this is an error response (has error field).
|
|
pub fn is_error(&self) -> bool {
|
|
self.error.is_some()
|
|
}
|
|
|
|
/// Get the error object if present.
|
|
pub fn get_error(&self) -> Option<&JsonRpcError> {
|
|
self.error.as_ref()
|
|
}
|
|
|
|
/// Get the result value if present.
|
|
pub fn get_result(&self) -> Option<&T> {
|
|
self.result.as_ref()
|
|
}
|
|
}
|
|
|
|
/// Check if a JSON-RPC error response contains the SSRF_BLOCKED error code.
|
|
///
|
|
/// This helper parses the error data field and checks for the SSRF_BLOCKED
|
|
/// code that indicates a URL was rejected due to SSRF protection.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `response_json` - The JSON string of the response to parse
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok(true)` if the response is an error with SSRF_BLOCKED code
|
|
/// * `Ok(false)` if the response is not an SSRF_BLOCKED error
|
|
/// * `Err(_)` if the JSON is invalid or not a valid JSON-RPC response
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust
|
|
/// use mcp_helpers::is_ssrf_blocked_error;
|
|
///
|
|
/// let error_json = r#"{"jsonrpc":"2.0","error":{"code":-32001,"message":"SSRF_BLOCKED","data":{"code":"SSRF_BLOCKED"}},"id":1}"#;
|
|
/// assert!(is_ssrf_blocked_error(error_json).unwrap());
|
|
/// ```
|
|
pub fn is_ssrf_blocked_error(response_json: &str) -> Result<bool, String> {
|
|
// Parse the JSON-RPC response
|
|
let parsed: JsonRpcResponse<serde_json::Value> =
|
|
serde_json::from_str(response_json).map_err(|e| format!("Invalid JSON: {}", e))?;
|
|
|
|
// Check if it has an error field
|
|
let error = parsed
|
|
.error
|
|
.ok_or_else(|| "Response has no error field".to_string())?;
|
|
|
|
// Check the error data for SSRF_BLOCKED code
|
|
if let Some(data) = &error.data {
|
|
if let Some(code) = data.get("code").and_then(|c| c.as_str()) {
|
|
return Ok(code == "SSRF_BLOCKED");
|
|
}
|
|
}
|
|
|
|
// Check if the error message itself contains SSRF_BLOCKED
|
|
if error.message.contains("SSRF_BLOCKED") {
|
|
return Ok(true);
|
|
}
|
|
|
|
Ok(false)
|
|
}
|
|
|
|
/// Extract error information from a JSON-RPC error response.
|
|
///
|
|
/// Returns a tuple of (code, message, optional_data) for easier error handling.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `response_json` - The JSON string of the error response
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok((code, message, data))` with the error details
|
|
/// * `Err(_)` if parsing fails or response is not an error
|
|
pub fn extract_error_info(
|
|
response_json: &str,
|
|
) -> Result<(i64, String, Option<serde_json::Value>), String> {
|
|
let parsed: JsonRpcResponse<serde_json::Value> =
|
|
serde_json::from_str(response_json).map_err(|e| format!("Failed to parse response: {}", e))?;
|
|
|
|
let error = parsed
|
|
.error
|
|
.ok_or_else(|| "Response is not an error".to_string())?;
|
|
|
|
Ok((error.code, error.message, error.data))
|
|
}
|
|
|
|
/// Parse a JSON-RPC response string into a JsonRpcResponse object.
|
|
///
|
|
/// This is a convenience wrapper around serde_json::from_str with better
|
|
/// error messages for JSON-RPC specific issues.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `response_json` - The JSON string of the response
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok(JsonRpcResponse)` if parsing succeeds
|
|
/// * `Err(String)` if parsing fails with descriptive error message
|
|
pub fn parse_response(response_json: &str) -> Result<JsonRpcResponse<serde_json::Value>, String> {
|
|
serde_json::from_str(response_json)
|
|
.map_err(|e| format!("Failed to parse JSON-RPC response: {}", e))
|
|
}
|
|
|
|
// ============================================================================
|
|
// JSON-RPC Framing Helpers
|
|
// ============================================================================
|
|
|
|
/// Read a framed JSON-RPC response from a reader.
|
|
///
|
|
/// Parses the Content-Length header, reads the exact number of bytes,
|
|
/// and returns the JSON body as a string.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `reader` - Buffered reader to read from (typically BufReader<ChildStdout>)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok(Some(json_body))` - Successfully read and parsed the frame
|
|
/// * `Ok(None)` - Reached EOF
|
|
/// * `Err(_)` - I/O error or invalid framing
|
|
fn read_framed_response<R: std::io::Read>(
|
|
reader: &mut BufReader<R>,
|
|
) -> std::io::Result<Option<String>> {
|
|
let mut content_length: Option<usize> = None;
|
|
|
|
// Read headers until empty line
|
|
loop {
|
|
let mut line = String::new();
|
|
let bytes_read = reader.read_line(&mut line)?;
|
|
if bytes_read == 0 {
|
|
return Ok(None); // EOF
|
|
}
|
|
|
|
let line = line.trim_end_matches(|c| c == '\r' || c == '\n');
|
|
if line.is_empty() {
|
|
break;
|
|
}
|
|
|
|
if let Some(value) = line.strip_prefix("Content-Length:") {
|
|
content_length = Some(
|
|
value
|
|
.trim()
|
|
.parse::<usize>()
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?,
|
|
);
|
|
}
|
|
}
|
|
|
|
let content_length = content_length.ok_or_else(|| {
|
|
std::io::Error::new(
|
|
std::io::ErrorKind::InvalidData,
|
|
"Missing Content-Length header",
|
|
)
|
|
})?;
|
|
|
|
let mut buffer = vec![0u8; content_length];
|
|
reader.read_exact(&mut buffer)?;
|
|
Ok(Some(String::from_utf8(buffer).map_err(|e| {
|
|
std::io::Error::new(std::io::ErrorKind::InvalidData, e)
|
|
})?))
|
|
}
|
|
|
|
/// Write a framed JSON-RPC message to a writer.
|
|
///
|
|
/// Uses LSP-style framing: Content-Length header, blank line, then JSON body.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `writer` - Writer to send the framed message to
|
|
/// * `json_body` - The JSON body string to send
|
|
fn write_framed_message<W: std::io::Write>(
|
|
writer: &mut W,
|
|
json_body: &str,
|
|
) -> std::io::Result<()> {
|
|
let header = format!("Content-Length: {}\r\n\r\n", json_body.len());
|
|
writer.write_all(header.as_bytes())?;
|
|
writer.write_all(json_body.as_bytes())?;
|
|
writer.flush()
|
|
}
|
|
|
|
// ============================================================================
|
|
// High-Level MCP Tool Call Helper
|
|
// ============================================================================
|
|
|
|
/// Result type for MCP tool call responses.
|
|
///
|
|
/// Provides a structured Result that distinguishes between successful
|
|
/// tool call results and JSON-RPC errors, making it easy to assert
|
|
/// expected behavior in tests.
|
|
#[derive(Debug, Clone)]
|
|
pub enum ToolCallResult {
|
|
/// Successful tool call with the result data
|
|
Success(serde_json::Value),
|
|
/// Error response with error code, message, and optional data
|
|
Error {
|
|
code: i64,
|
|
message: String,
|
|
data: Option<serde_json::Value>,
|
|
},
|
|
}
|
|
|
|
impl ToolCallResult {
|
|
/// Check if this is a successful result.
|
|
pub fn is_success(&self) -> bool {
|
|
matches!(self, ToolCallResult::Success(_))
|
|
}
|
|
|
|
/// Check if this is an error result.
|
|
pub fn is_error(&self) -> bool {
|
|
matches!(self, ToolCallResult::Error { .. })
|
|
}
|
|
|
|
/// Get the success value if present.
|
|
pub fn get_success(&self) -> Option<&serde_json::Value> {
|
|
match self {
|
|
ToolCallResult::Success(v) => Some(v),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/// Get the error components if present.
|
|
pub fn get_error(&self) -> Option<(i64, &str, Option<&serde_json::Value>)> {
|
|
match self {
|
|
ToolCallResult::Error { code, message, data } => {
|
|
Some((*code, message.as_str(), data.as_ref()))
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/// Check if this error has a specific error code in data.code field.
|
|
///
|
|
/// Used for checking specific diagnostic codes like "SSRF_BLOCKED".
|
|
pub fn has_error_code(&self, expected_code: &str) -> bool {
|
|
match self {
|
|
ToolCallResult::Error { data, .. } => {
|
|
if let Some(data_obj) = data {
|
|
if let Some(code) = data_obj.get("code").and_then(|c| c.as_str()) {
|
|
return code == expected_code;
|
|
}
|
|
}
|
|
false
|
|
}
|
|
_ => false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Send an MCP tool call request and parse the response.
|
|
///
|
|
/// This helper function handles the complete request/response cycle for
|
|
/// MCP tool calls:
|
|
/// 1. Constructs a JSON-RPC tools/call request
|
|
/// 2. Writes it to the server's stdin with proper framing
|
|
/// 3. Reads the framed response from stdout
|
|
/// 4. Parses and returns a structured ToolCallResult
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `server` - Mutable reference to the spawned MCP server guard
|
|
/// * `tool_name` - Name of the tool to call (e.g., "extract", "get_metadata")
|
|
/// * `arguments` - Tool arguments as a JSON value (typically an object)
|
|
/// * `timeout_ms` - Maximum time to wait for response (default 5000ms recommended)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok(ToolCallResult)` - Parsed response (success or error)
|
|
/// * `Err(String)` - I/O or parsing error with descriptive message
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```rust
|
|
/// use std::thread;
|
|
/// use std::time::Duration;
|
|
/// use serde_json::json;
|
|
///
|
|
/// let mut server = spawn_mcp_stdio();
|
|
/// thread::sleep(Duration::from_millis(50));
|
|
///
|
|
/// let result = send_tool_call(
|
|
/// &mut server,
|
|
/// "extract",
|
|
/// json!({"path": "https://example.com/doc.pdf"}),
|
|
/// 5000
|
|
/// ).expect("Tool call should succeed");
|
|
///
|
|
/// if result.is_error() && result.has_error_code("SSRF_BLOCKED") {
|
|
/// // URL was rejected as expected
|
|
/// }
|
|
/// ```
|
|
pub fn send_tool_call(
|
|
server: &mut std::process::Child,
|
|
tool_name: &str,
|
|
arguments: serde_json::Value,
|
|
timeout_ms: u64,
|
|
) -> Result<ToolCallResult, String> {
|
|
use std::io::{BufRead, BufReader, Write};
|
|
|
|
// Build the JSON-RPC request
|
|
let request = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": tool_name,
|
|
"arguments": arguments
|
|
}
|
|
});
|
|
|
|
let request_str = serde_json::to_string(&request)
|
|
.map_err(|e| format!("Failed to serialize request: {}", e))?;
|
|
|
|
// Write the framed request
|
|
{
|
|
let stdin = server.stdin.as_mut()
|
|
.ok_or_else(|| "Failed to get stdin handle".to_string())?;
|
|
|
|
let header = format!("Content-Length: {}\r\n\r\n", request_str.len());
|
|
stdin.write_all(header.as_bytes())
|
|
.map_err(|e| format!("Failed to write header: {}", e))?;
|
|
stdin.write_all(request_str.as_bytes())
|
|
.map_err(|e| format!("Failed to write request body: {}", e))?;
|
|
stdin.flush()
|
|
.map_err(|e| format!("Failed to flush stdin: {}", e))?;
|
|
}
|
|
|
|
// Read the framed response with timeout
|
|
let start = std::time::Instant::now();
|
|
let response = loop {
|
|
let stdout = server.stdout.as_mut()
|
|
.ok_or_else(|| "Failed to get stdout handle".to_string())?;
|
|
let mut reader = BufReader::new(stdout);
|
|
|
|
match read_framed_response(&mut reader) {
|
|
Ok(Some(resp)) => break resp,
|
|
Ok(None) => return Err("Unexpected EOF while reading response".to_string()),
|
|
Err(e) if start.elapsed() >= std::time::Duration::from_millis(timeout_ms) => {
|
|
return Err(format!("Timeout waiting for response: {}", e));
|
|
}
|
|
Err(_) => {
|
|
std::thread::sleep(std::time::Duration::from_millis(10));
|
|
if start.elapsed() >= std::time::Duration::from_millis(timeout_ms) {
|
|
return Err("Timeout waiting for response".to_string());
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
};
|
|
|
|
// Parse the JSON-RPC response
|
|
let parsed: JsonRpcResponse<serde_json::Value> = parse_response(&response)?;
|
|
|
|
// Convert to ToolCallResult
|
|
if let Some(result_value) = parsed.result {
|
|
Ok(ToolCallResult::Success(result_value))
|
|
} else if let Some(error_obj) = parsed.error {
|
|
Ok(ToolCallResult::Error {
|
|
code: error_obj.code,
|
|
message: error_obj.message,
|
|
data: error_obj.data,
|
|
})
|
|
} else {
|
|
Err("Response has neither result nor error field".to_string())
|
|
}
|
|
}
|
|
|
|
/// Convenience function to send an extract tool call with a URL parameter.
|
|
///
|
|
/// This is the most common pattern for SSRF testing.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `server` - Mutable reference to the spawned MCP server guard
|
|
/// * `url` - The URL to extract from (will be passed as the "path" argument)
|
|
/// * `timeout_ms` - Maximum time to wait for response (default 5000ms recommended)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok(ToolCallResult)` - Parsed response (success or error)
|
|
/// * `Err(String)` - I/O or parsing error with descriptive message
|
|
pub fn send_extract_call(
|
|
server: &mut std::process::Child,
|
|
url: &str,
|
|
timeout_ms: u64,
|
|
) -> Result<ToolCallResult, String> {
|
|
let arguments = json!({
|
|
"path": url
|
|
});
|
|
send_tool_call(server, "extract", arguments, timeout_ms)
|
|
}
|
|
|
|
/// Convenience function to send a get_metadata tool call with a URL parameter.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `server` - Mutable reference to the spawned MCP server guard
|
|
/// * `url` - The URL to fetch metadata from (will be passed as the "path" argument)
|
|
/// * `timeout_ms` - Maximum time to wait for response (default 5000ms recommended)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// * `Ok(ToolCallResult)` - Parsed response (success or error)
|
|
/// * `Err(String)` - I/O or parsing error with descriptive message
|
|
pub fn send_get_metadata_call(
|
|
server: &mut std::process::Child,
|
|
url: &str,
|
|
timeout_ms: u64,
|
|
) -> Result<ToolCallResult, String> {
|
|
let arguments = json!({
|
|
"path": url
|
|
});
|
|
send_tool_call(server, "get_metadata", arguments, timeout_ms)
|
|
}
|
|
|
|
mod mcp_helpers_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_is_ssrf_blocked_error_with_code_in_data() {
|
|
let error_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32001,
|
|
"message": "SSRF protection blocked this URL",
|
|
"data": {"code": "SSRF_BLOCKED"}
|
|
},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let result = is_ssrf_blocked_error(error_json);
|
|
assert!(result.is_ok());
|
|
assert!(result.unwrap(), "Should detect SSRF_BLOCKED in data.code");
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_ssrf_blocked_error_with_message() {
|
|
let error_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32001,
|
|
"message": "SSRF_BLOCKED: URL targets private network"
|
|
},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let result = is_ssrf_blocked_error(error_json);
|
|
assert!(result.is_ok());
|
|
assert!(result.unwrap(), "Should detect SSRF_BLOCKED in message");
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_ssrf_blocked_error_not_blocked() {
|
|
let error_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32601,
|
|
"message": "Method not found"
|
|
},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let result = is_ssrf_blocked_error(error_json);
|
|
assert!(result.is_ok());
|
|
assert!(!result.unwrap(), "Should not detect SSRF_BLOCKED");
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_ssrf_blocked_error_success_response() {
|
|
let success_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"result": {"status": "ok"},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let result = is_ssrf_blocked_error(success_json);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().contains("no error field"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_error_info_success() {
|
|
let error_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32001,
|
|
"message": "SSRF_BLOCKED",
|
|
"data": {"url": "http://127.0.0.1/"}
|
|
},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let (code, message, data) = extract_error_info(error_json).unwrap();
|
|
assert_eq!(code, -32001);
|
|
assert_eq!(message, "SSRF_BLOCKED");
|
|
assert!(data.is_some());
|
|
assert_eq!(data.unwrap()["url"], "http://127.0.0.1/");
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_error_info_not_an_error() {
|
|
let success_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"result": {"status": "ok"},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let result = extract_error_info(success_json);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().contains("not an error"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_response_success() {
|
|
let success_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"result": {"tools": []},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let response = parse_response(success_json).unwrap();
|
|
assert!(response.is_success());
|
|
assert!(!response.is_error());
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_response_error() {
|
|
let error_json = r#"{
|
|
"jsonrpc": "2.0",
|
|
"error": {
|
|
"code": -32601,
|
|
"message": "Method not found"
|
|
},
|
|
"id": 1
|
|
}"#;
|
|
|
|
let response = parse_response(error_json).unwrap();
|
|
assert!(!response.is_success());
|
|
assert!(response.is_error());
|
|
assert_eq!(response.get_error().unwrap().code, -32601);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_response_invalid_json() {
|
|
let invalid_json = r#"{not valid json}"#;
|
|
|
|
let result = parse_response(invalid_json);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().contains("Failed to parse JSON-RPC response"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_response_missing_jsonrpc_field() {
|
|
let invalid_json = r#"{"result": {}, "id": 1}"#;
|
|
|
|
let result = parse_response(invalid_json);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_read_framed_response_simple() {
|
|
let input = b"Content-Length: 25\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":1}";
|
|
let mut reader = BufReader::new(&input[..]);
|
|
|
|
let result = read_framed_response(&mut reader);
|
|
assert!(result.is_ok());
|
|
let json_body = result.unwrap().unwrap();
|
|
assert_eq!(json_body, "{\"jsonrpc\":\"2.0\",\"id\":1}");
|
|
}
|
|
|
|
#[test]
|
|
fn test_read_framed_response_with_extra_whitespace() {
|
|
let input = b"Content-Length: 25\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":1}";
|
|
let mut reader = BufReader::new(&input[..]);
|
|
|
|
let result = read_framed_response(&mut reader);
|
|
assert!(result.is_ok());
|
|
let json_body = result.unwrap().unwrap();
|
|
assert_eq!(json_body, "{\"jsonrpc\":\"2.0\",\"id\":1}");
|
|
}
|
|
|
|
#[test]
|
|
fn test_read_framed_response_eof() {
|
|
let input = b"";
|
|
let mut reader = BufReader::new(&input[..]);
|
|
|
|
let result = read_framed_response(&mut reader);
|
|
assert!(result.is_ok());
|
|
assert!(result.unwrap().is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_read_framed_response_missing_content_length() {
|
|
let input = b"\r\n\r\n{\"jsonrpc\":\"2.0\"}";
|
|
let mut reader = BufReader::new(&input[..]);
|
|
|
|
let result = read_framed_response(&mut reader);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().to_string().contains("Missing Content-Length"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_write_framed_message_simple() {
|
|
let mut buffer = Vec::new();
|
|
let json_body = r#"{"test":"data"}"#;
|
|
|
|
let result = write_framed_message(&mut buffer, json_body);
|
|
assert!(result.is_ok());
|
|
|
|
let output = String::from_utf8(buffer).unwrap();
|
|
assert!(output.contains("Content-Length: 15"));
|
|
assert!(output.contains("\r\n\r\n"));
|
|
assert!(output.contains(json_body));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_result_success() {
|
|
let result_value = json!({"status": "ok", "data": {"items": []}});
|
|
let result = ToolCallResult::Success(result_value);
|
|
|
|
assert!(result.is_success());
|
|
assert!(!result.is_error());
|
|
assert!(result.get_success().is_some());
|
|
assert!(result.get_error().is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_result_error() {
|
|
let data = json!({"code": "SSRF_BLOCKED"});
|
|
let result = ToolCallResult::Error {
|
|
code: -32001,
|
|
message: "URL blocked".to_string(),
|
|
data: Some(data),
|
|
};
|
|
|
|
assert!(!result.is_success());
|
|
assert!(result.is_error());
|
|
assert!(result.get_success().is_none());
|
|
assert!(result.get_error().is_some());
|
|
|
|
let (code, message, data_opt) = result.get_error().unwrap();
|
|
assert_eq!(code, -32001);
|
|
assert_eq!(message, "URL blocked");
|
|
assert!(data_opt.is_some());
|
|
assert_eq!(data_opt.unwrap().get("code").unwrap().as_str().unwrap(), "SSRF_BLOCKED");
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_result_has_error_code() {
|
|
let data = json!({"code": "SSRF_BLOCKED"});
|
|
let result = ToolCallResult::Error {
|
|
code: -32001,
|
|
message: "URL blocked".to_string(),
|
|
data: Some(data),
|
|
};
|
|
|
|
assert!(result.has_error_code("SSRF_BLOCKED"));
|
|
assert!(!result.has_error_code("OTHER_ERROR"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_result_has_error_code_no_data() {
|
|
let result = ToolCallResult::Error {
|
|
code: -32001,
|
|
message: "URL blocked".to_string(),
|
|
data: None,
|
|
};
|
|
|
|
assert!(!result.has_error_code("SSRF_BLOCKED"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_result_has_error_code_malformed_data() {
|
|
let data = json!({"message": "something"});
|
|
let result = ToolCallResult::Error {
|
|
code: -32001,
|
|
message: "URL blocked".to_string(),
|
|
data: Some(data),
|
|
};
|
|
|
|
assert!(!result.has_error_code("SSRF_BLOCKED"));
|
|
}
|
|
|
|
// Tests for JsonRpcError::is_ssrf_blocked()
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_with_code_in_data() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "SSRF protection blocked this URL".to_string(),
|
|
data: Some(json!({"code": "SSRF_BLOCKED"})),
|
|
};
|
|
assert!(error.is_ssrf_blocked(), "Should detect SSRF_BLOCKED in data.code");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_with_message() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "SSRF_BLOCKED: URL targets private network".to_string(),
|
|
data: None,
|
|
};
|
|
assert!(error.is_ssrf_blocked(), "Should detect SSRF_BLOCKED in message");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_not_blocked() {
|
|
let error = JsonRpcError {
|
|
code: -32601,
|
|
message: "Method not found".to_string(),
|
|
data: None,
|
|
};
|
|
assert!(!error.is_ssrf_blocked(), "Should not detect SSRF_BLOCKED");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_empty_data() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "Some error".to_string(),
|
|
data: Some(json!({})),
|
|
};
|
|
assert!(!error.is_ssrf_blocked(), "Should not detect SSRF_BLOCKED with empty data");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_different_code_in_data() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "Some error".to_string(),
|
|
data: Some(json!({"code": "OTHER_ERROR"})),
|
|
};
|
|
assert!(!error.is_ssrf_blocked(), "Should not detect OTHER_ERROR as SSRF_BLOCKED");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_case_sensitive_in_message() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "ssrf_blocked: lowercase".to_string(),
|
|
data: None,
|
|
};
|
|
assert!(!error.is_ssrf_blocked(), "Should be case-sensitive in message");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_case_sensitive_in_data() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "Some error".to_string(),
|
|
data: Some(json!({"code": "ssrf_blocked"})),
|
|
};
|
|
assert!(!error.is_ssrf_blocked(), "Should be case-sensitive in data");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_partial_match_in_message() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "SSRF_BLOCKED detected in request".to_string(),
|
|
data: None,
|
|
};
|
|
assert!(error.is_ssrf_blocked(), "Should detect partial match in message");
|
|
}
|
|
|
|
#[test]
|
|
fn test_json_rpc_error_is_ssrf_blocked_both_data_and_message() {
|
|
let error = JsonRpcError {
|
|
code: -32001,
|
|
message: "SSRF_BLOCKED: URL rejected".to_string(),
|
|
data: Some(json!({"code": "SSRF_BLOCKED"})),
|
|
};
|
|
assert!(error.is_ssrf_blocked(), "Should detect SSRF_BLOCKED in both data and message");
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_tool_call_builder_extract_basic() {
|
|
let request = ToolCallBuilder::extract()
|
|
.with_url("https://example.com/")
|
|
.build();
|
|
|
|
assert_eq!(request.method, "tools/call");
|
|
assert!(request.params.is_some());
|
|
|
|
let params = request.params.unwrap();
|
|
assert_eq!(params["name"], "extract");
|
|
assert_eq!(params["arguments"]["path"], "https://example.com/");
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_builder_extract_with_id() {
|
|
let request = ToolCallBuilder::extract()
|
|
.with_url("https://example.com/")
|
|
.with_id(Id::String("test-123".to_string()))
|
|
.build();
|
|
|
|
assert_eq!(request.request_id(), Id::String("test-123".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_call_builder_with_custom_argument() {
|
|
let request = ToolCallBuilder::extract()
|
|
.with_url("https://example.com/")
|
|
.with_argument("ocr", serde_json::Value::Bool(true))
|
|
.build();
|
|
|
|
let params = request.params.unwrap();
|
|
assert_eq!(params["arguments"]["ocr"], true);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extract_call_quick_helper() {
|
|
let request = extract_call("https://example.com/doc.pdf");
|
|
|
|
assert_eq!(request.method, "tools/call");
|
|
let params = request.params.unwrap();
|
|
assert_eq!(params["name"], "extract");
|
|
assert_eq!(params["arguments"]["path"], "https://example.com/doc.pdf");
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_metadata_call_quick_helper() {
|
|
let request = get_metadata_call("https://example.com/doc.pdf");
|
|
|
|
assert_eq!(request.method, "tools/call");
|
|
let params = request.params.unwrap();
|
|
assert_eq!(params["name"], "get_metadata");
|
|
assert_eq!(
|
|
params["arguments"]["path"],
|
|
"https://example.com/doc.pdf"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_serialization_format() {
|
|
let request = ToolCallBuilder::extract()
|
|
.with_url("https://example.com/")
|
|
.build();
|
|
|
|
let json = serde_json::to_string(&request).unwrap();
|
|
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
|
|
|
|
assert_eq!(parsed["jsonrpc"], "2.0");
|
|
assert_eq!(parsed["method"], "tools/call");
|
|
assert_eq!(parsed["params"]["name"], "extract");
|
|
assert_eq!(parsed["params"]["arguments"]["path"], "https://example.com/");
|
|
assert_eq!(parsed["id"], 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_arguments() {
|
|
let request = ToolCallBuilder::extract()
|
|
.with_url("https://example.com/")
|
|
.with_argument("password", serde_json::Value::String("secret123".to_string()))
|
|
.with_argument("ocr", serde_json::Value::Bool(true))
|
|
.with_argument("pages", serde_json::Value::String("1-5".to_string()))
|
|
.build();
|
|
|
|
let params = request.params.unwrap();
|
|
let args = ¶ms["arguments"];
|
|
|
|
assert_eq!(args["path"], "https://example.com/");
|
|
assert_eq!(args["password"], "secret123");
|
|
assert_eq!(args["ocr"], true);
|
|
assert_eq!(args["pages"], "1-5");
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// MCP Server Integration Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(feature = "remote")]
|
|
#[cfg(test)]
|
|
mod mcp_ssrf_tests {
|
|
//! MCP server integration tests for SSRF protection.
|
|
//!
|
|
//! These tests verify that when the MCP server receives URL parameters
|
|
//! through JSON-RPC tools/call, SSRF-prone URLs are properly rejected.
|
|
//!
|
|
//! Currently, the MCP server returns stub responses for all URLs
|
|
//! (remote extraction requires Phase 1.8). Once Phase 1.8 is implemented,
|
|
//! these tests should verify that:
|
|
//! 1. SSRF-prone URLs are rejected with SSRF_BLOCKED error
|
|
//! 2. No actual network connections are attempted
|
|
//! 3. Public URLs are accepted
|
|
|
|
use std::io::{BufRead, BufReader, Read, Write};
|
|
use std::process::{Child, Command, Stdio};
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
|
|
/// RAII guard for spawned child processes.
|
|
///
|
|
/// Ensures deterministic cleanup on Drop (even on panic) by:
|
|
/// 1. Killing the child process
|
|
/// 2. Waiting with a timeout to reap the zombie
|
|
///
|
|
/// Per CLAUDE.md test hygiene rules, this prevents orphaned processes
|
|
/// and hung tests.
|
|
struct ProcessGuard {
|
|
child: Option<Child>,
|
|
}
|
|
|
|
impl ProcessGuard {
|
|
/// Create a new guard from a spawned child process.
|
|
fn new(child: Child) -> Self {
|
|
Self { child: Some(child) }
|
|
}
|
|
|
|
/// Get a mutable reference to the child process.
|
|
fn child_mut(&mut self) -> &mut Child {
|
|
self.child.as_mut().expect("ProcessGuard: child taken without replacement")
|
|
}
|
|
|
|
/// Take ownership of the child, dropping the guard without cleanup.
|
|
///
|
|
/// Use this when you want manual control over cleanup.
|
|
fn take(mut self) -> Child {
|
|
self.child.take().expect("ProcessGuard: child already taken")
|
|
}
|
|
}
|
|
|
|
impl Drop for ProcessGuard {
|
|
fn drop(&mut self) {
|
|
if let Some(mut child) = self.child.take() {
|
|
// Kill the process (signal-based, may fail if already dead)
|
|
let _ = child.kill();
|
|
// Wait with timeout to reap zombie, don't block forever
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// SSRF test payloads for MCP server testing.
|
|
///
|
|
/// These are the critical payloads that MUST be rejected by the MCP server
|
|
/// when remote extraction is implemented (Phase 1.8).
|
|
const MCP_SSRF_PAYLOADS: &[&str] = &[
|
|
"http://127.0.0.1:9999/", // Loopback with non-standard port
|
|
"http://0.0.0.0/", // All interfaces
|
|
"http://169.254.169.254/latest/meta-data/", // AWS metadata endpoint
|
|
"http://10.0.0.1/internal", // RFC 1918 private network
|
|
"http://[::1]/", // IPv6 loopback
|
|
];
|
|
|
|
/// Helper to spawn the pdftract MCP server in stdio mode.
|
|
fn spawn_mcp_stdio() -> std::process::Child {
|
|
Command::new("pdftract")
|
|
.arg("mcp")
|
|
.arg("--stdio")
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::null())
|
|
.spawn()
|
|
.expect("Failed to spawn pdftract mcp --stdio")
|
|
}
|
|
|
|
/// Helper to write a framed JSON-RPC message to stdin.
|
|
fn write_framed_message(
|
|
stdin: &mut std::process::ChildStdin,
|
|
json_body: &str,
|
|
) -> std::io::Result<()> {
|
|
let header = format!("Content-Length: {}\r\n\r\n", json_body.len());
|
|
stdin.write_all(header.as_bytes())?;
|
|
stdin.write_all(json_body.as_bytes())?;
|
|
stdin.flush()
|
|
}
|
|
|
|
/// Helper to read a framed JSON-RPC response from stdout.
|
|
fn read_framed_response<R: std::io::Read>(
|
|
reader: &mut BufReader<R>,
|
|
) -> std::io::Result<Option<String>> {
|
|
let mut content_length: Option<usize> = None;
|
|
|
|
// Read headers until empty line
|
|
loop {
|
|
let mut line = String::new();
|
|
let bytes_read = reader.read_line(&mut line)?;
|
|
if bytes_read == 0 {
|
|
return Ok(None); // EOF
|
|
}
|
|
|
|
let line = line.trim_end_matches(|c| c == '\r' || c == '\n');
|
|
if line.is_empty() {
|
|
break;
|
|
}
|
|
|
|
if let Some(value) = line.strip_prefix("Content-Length:") {
|
|
content_length = Some(
|
|
value
|
|
.trim()
|
|
.parse::<usize>()
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?,
|
|
);
|
|
}
|
|
}
|
|
|
|
let content_length = content_length.ok_or_else(|| {
|
|
std::io::Error::new(
|
|
std::io::ErrorKind::InvalidData,
|
|
"Missing Content-Length header",
|
|
)
|
|
})?;
|
|
|
|
let mut buffer = vec![0u8; content_length];
|
|
reader.read_exact(&mut buffer)?;
|
|
Ok(Some(String::from_utf8(buffer).map_err(|e| {
|
|
std::io::Error::new(std::io::ErrorKind::InvalidData, e)
|
|
})?))
|
|
}
|
|
|
|
/// Wait for a process to complete with a timeout.
|
|
fn wait_with_timeout(
|
|
child: &mut std::process::Child,
|
|
timeout_ms: u64,
|
|
) -> std::io::Result<Option<i32>> {
|
|
let deadline = std::time::Instant::now() + Duration::from_millis(timeout_ms);
|
|
|
|
loop {
|
|
if let Some(status) = child.try_wait()? {
|
|
return Ok(status.code());
|
|
}
|
|
|
|
if std::time::Instant::now() >= deadline {
|
|
// Timeout: kill the process
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
return Err(std::io::Error::new(
|
|
std::io::ErrorKind::TimedOut,
|
|
"Process timed out",
|
|
));
|
|
}
|
|
|
|
thread::sleep(Duration::from_millis(50));
|
|
}
|
|
}
|
|
|
|
/// Test that MCP server handles SSRF-prone URLs safely.
|
|
///
|
|
/// Current behavior: Returns stub response (remote extraction not implemented).
|
|
/// Expected behavior (Phase 1.8): Returns JSON-RPC error with SSRF_BLOCKED code.
|
|
#[test]
|
|
fn test_mcp_extract_tool_rejects_ssrf_urls() {
|
|
for url in MCP_SSRF_PAYLOADS {
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the helper to construct a proper JSON-RPC tools/call request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(*url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str)
|
|
.expect("Failed to write request");
|
|
}
|
|
|
|
// Read the response
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
// Verify the response
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current behavior: stub response with _note
|
|
// Future behavior (Phase 1.8): should return error with SSRF_BLOCKED
|
|
if parsed.get("result").is_some() {
|
|
// Current implementation returns stub result
|
|
let result = &parsed["result"];
|
|
assert!(
|
|
result.get("_note").is_some(),
|
|
"SSRF URL '{}' should return stub response or error, got: {}",
|
|
url,
|
|
response
|
|
);
|
|
|
|
// Verify no actual network activity occurred
|
|
// (The stub response doesn't fetch the URL)
|
|
} else if parsed.get("error").is_some() {
|
|
// Future implementation: should return SSRF_BLOCKED error
|
|
let error = &parsed["error"];
|
|
let error_data = error.get("data");
|
|
|
|
// Once Phase 1.8 is implemented, this should check for SSRF_BLOCKED
|
|
if let Some(data) = error_data {
|
|
let _code = data.get("code").and_then(|c| c.as_str());
|
|
// Future: assert_eq!(_code, Some("SSRF_BLOCKED"));
|
|
}
|
|
} else {
|
|
panic!(
|
|
"SSRF URL '{}' should return stub response or error, got: {}",
|
|
url, response
|
|
);
|
|
}
|
|
|
|
// Clean shutdown
|
|
let _ = child.stdin.take();
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
}
|
|
|
|
/// Test that MCP server doesn't attempt actual connections to SSRF URLs.
|
|
///
|
|
/// This verifies that even if a URL is passed, the server doesn't try to
|
|
/// fetch it (since remote extraction is not implemented yet).
|
|
#[test]
|
|
fn test_mcp_no_network_connections_to_ssrf_urls() {
|
|
// Use localhost:9999 which is unlikely to have a listener
|
|
let dangerous_url = "http://127.0.0.1:9999/test.pdf";
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(dangerous_url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str).expect("Failed to write request");
|
|
}
|
|
|
|
// Read response within timeout
|
|
let start = std::time::Instant::now();
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
let elapsed = start.elapsed();
|
|
|
|
// Should return quickly (no network timeout)
|
|
assert!(
|
|
elapsed < Duration::from_millis(500),
|
|
"Response should return quickly without network attempt, took {:?}",
|
|
elapsed
|
|
);
|
|
|
|
// Verify stub response (not an error from failed connection)
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
assert!(
|
|
parsed.get("result").is_some(),
|
|
"Should return stub result, not connection error"
|
|
);
|
|
|
|
// Clean shutdown
|
|
let _ = child.stdin.take();
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
|
|
/// Test that IPv6 loopback is handled safely.
|
|
#[test]
|
|
fn test_mcp_ipv6_loopback_rejected() {
|
|
let ipv6_loopback_urls = &[
|
|
"http://[::1]/",
|
|
"http://[::1]:8080/test.pdf",
|
|
"http://[0:0:0:0:0:0:0:1]/", // Full form of ::1
|
|
];
|
|
|
|
for url in ipv6_loopback_urls {
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(*url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str)
|
|
.expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current: stub response
|
|
// Future: SSRF_BLOCKED error
|
|
assert!(
|
|
parsed.get("result").is_some() || parsed.get("error").is_some(),
|
|
"IPv6 loopback URL '{}' should return valid response",
|
|
url
|
|
);
|
|
|
|
// Clean shutdown
|
|
let _ = child.stdin.take();
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
}
|
|
|
|
/// Test that cloud metadata endpoints are blocked.
|
|
#[test]
|
|
fn test_mcp_cloud_metadata_endpoints_blocked() {
|
|
let metadata_urls = &[
|
|
"http://169.254.169.254/latest/meta-data/",
|
|
"http://metadata.google.internal/computeMetadata/v1/",
|
|
"http://168.63.129.16/latest/",
|
|
];
|
|
|
|
for url in metadata_urls {
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the helper to construct the request for get_metadata tool
|
|
use crate::mcp_helpers::get_metadata_call;
|
|
let request = get_metadata_call(*url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str)
|
|
.expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Should never succeed in accessing metadata endpoints
|
|
assert!(
|
|
parsed.get("error").is_some()
|
|
|| parsed
|
|
.get("result")
|
|
.and_then(|r| r.get("_note"))
|
|
.is_some(),
|
|
"Metadata endpoint '{}' should be blocked or return stub",
|
|
url
|
|
);
|
|
|
|
// Clean shutdown
|
|
let _ = child.stdin.take();
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
}
|
|
|
|
/// Test cleanup: ensure no orphaned processes after test completes.
|
|
///
|
|
/// Per CLAUDE.md test hygiene rules, all spawned processes must be
|
|
/// cleaned up deterministically.
|
|
#[test]
|
|
fn test_mcp_process_cleanup_on_completion() {
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Send a simple request
|
|
let request = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#;
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, request).expect("Failed to write request");
|
|
}
|
|
|
|
// Read response
|
|
let _ = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
// Close stdin to trigger clean shutdown
|
|
drop(child.stdin.take());
|
|
|
|
// Wait for process to exit (should exit within 200ms)
|
|
let result = wait_with_timeout(&mut child, 200);
|
|
|
|
assert!(
|
|
result.is_ok(),
|
|
"Process should exit cleanly after stdin close"
|
|
);
|
|
|
|
// Verify exit code is 0 (success)
|
|
let exit_code = result.unwrap().unwrap();
|
|
assert_eq!(exit_code, 0, "Process should exit with code 0");
|
|
}
|
|
|
|
// ============================================================================
|
|
// Individual SSRF URL Pattern Tests
|
|
// ============================================================================
|
|
|
|
/// Test that IPv4 loopback (127.0.0.1:9999) is rejected.
|
|
///
|
|
/// Tests the extract tool with an IPv4 loopback address on a non-standard port.
|
|
/// This verifies that SSRF protection catches localhost variants even with
|
|
/// custom port numbers.
|
|
#[test]
|
|
fn test_mcp_ipv4_loopback_rejected() {
|
|
let url = "http://127.0.0.1:9999/";
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the JSON-RPC helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str).expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current implementation returns stub response; future should return SSRF_BLOCKED
|
|
assert!(
|
|
parsed.get("result").is_some() || parsed.get("error").is_some(),
|
|
"IPv4 loopback URL should return valid response"
|
|
);
|
|
|
|
// Clean shutdown
|
|
drop(child.stdin.take());
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
|
|
/// Test that IPv4 all interfaces (0.0.0.0) is rejected.
|
|
///
|
|
/// Tests the extract tool with 0.0.0.0, which binds to all interfaces.
|
|
/// This is a dangerous SSRF target as it can access services listening
|
|
/// on any local interface.
|
|
#[test]
|
|
fn test_mcp_ipv4_all_interfaces_rejected() {
|
|
let url = "http://0.0.0.0/";
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the JSON-RPC helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str).expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current implementation returns stub response; future should return SSRF_BLOCKED
|
|
assert!(
|
|
parsed.get("result").is_some() || parsed.get("error").is_some(),
|
|
"IPv4 all-interfaces URL should return valid response"
|
|
);
|
|
|
|
// Clean shutdown
|
|
drop(child.stdin.take());
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
|
|
/// Test that IPv6 loopback (::1) is rejected.
|
|
///
|
|
/// Tests the extract tool with the IPv6 loopback address. This verifies
|
|
/// that SSRF protection handles IPv6 notation correctly, including the
|
|
/// bracket format for URLs.
|
|
#[test]
|
|
fn test_mcp_ipv6_loopback_rejected() {
|
|
let url = "http://[::1]/";
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the JSON-RPC helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str).expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current implementation returns stub response; future should return SSRF_BLOCKED
|
|
assert!(
|
|
parsed.get("result").is_some() || parsed.get("error").is_some(),
|
|
"IPv6 loopback URL should return valid response"
|
|
);
|
|
|
|
// Clean shutdown
|
|
drop(child.stdin.take());
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
|
|
/// Test that AWS metadata endpoint (169.254.169.254) is rejected.
|
|
///
|
|
/// Tests the extract tool with the AWS IMDSv2 endpoint path. This is a
|
|
/// critical SSRF payload as cloud metadata endpoints can expose instance
|
|
/// credentials and sensitive metadata.
|
|
#[test]
|
|
fn test_mcp_aws_metadata_endpoint_rejected() {
|
|
let url = "http://169.254.169.254/latest/meta-data/";
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the JSON-RPC helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str).expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current implementation returns stub response; future should return SSRF_BLOCKED
|
|
assert!(
|
|
parsed.get("result").is_some() || parsed.get("error").is_some(),
|
|
"AWS metadata endpoint URL should return valid response"
|
|
);
|
|
|
|
// Clean shutdown
|
|
drop(child.stdin.take());
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
|
|
/// Test that RFC 1918 private network (10.0.0.1) is rejected.
|
|
///
|
|
/// Tests the extract tool with a private network IP from the 10.0.0.0/8 range.
|
|
/// This verifies that SSRF protection blocks access to internal network
|
|
/// services that are not publicly accessible.
|
|
#[test]
|
|
fn test_mcp_private_network_rejected() {
|
|
let url = "http://10.0.0.1/internal";
|
|
let mut child = spawn_mcp_stdio();
|
|
thread::sleep(Duration::from_millis(50));
|
|
|
|
// Use the JSON-RPC helper to construct the request
|
|
use crate::mcp_helpers::extract_call;
|
|
let request = extract_call(url);
|
|
let request_str = serde_json::to_string(&request).unwrap();
|
|
|
|
{
|
|
let stdin = child.stdin.as_mut().expect("Failed to open stdin");
|
|
write_framed_message(stdin, &request_str).expect("Failed to write request");
|
|
}
|
|
|
|
let response = {
|
|
let stdout = child.stdout.as_mut().expect("Failed to open stdout");
|
|
let mut reader = BufReader::new(stdout);
|
|
read_framed_response(&mut reader)
|
|
.expect("Failed to read response")
|
|
.expect("No response received")
|
|
};
|
|
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&response).expect("Response is not valid JSON");
|
|
|
|
// Current implementation returns stub response; future should return SSRF_BLOCKED
|
|
assert!(
|
|
parsed.get("result").is_some() || parsed.get("error").is_some(),
|
|
"Private network URL should return valid response"
|
|
);
|
|
|
|
// Clean shutdown
|
|
drop(child.stdin.take());
|
|
let _ = wait_with_timeout(&mut child, 1000);
|
|
}
|
|
}
|