//! Input sanitization for MCP gateway traffic.
//!
//! Rejects null bytes, strips unsafe control characters, and normalizes
//! Unicode to NFC on all tool inputs/outputs passing through the gateway.
use serde_json::Value;
use crate::{Error, Result};
/// Characters that are always rejected (null byte).
const REJECTED_BYTE: u8 = 0x00;
/// Control characters to strip (C0 range excluding common whitespace).
/// We preserve: `\t` (0x09), `\n` (0x0A), `\r` (0x0D).
fn is_unsafe_control(c: char) -> bool {
let code = c as u32;
// C0 control characters (0x00-0x1F) minus tab, newline, carriage return
(code <= 0x1F && code != 0x09 && code != 0x0A && code != 0x0D)
// C1 control characters (0x80-0x9F) — rarely legitimate in text
|| (0x80..=0x9F).contains(&code)
// Unicode specials: zero-width chars often used for homograph attacks
|| c == '\u{200B}' // zero-width space
|| c == '\u{200C}' // zero-width non-joiner
|| c == '\u{200D}' // zero-width joiner
|| c == '\u{FEFF}' // byte order mark (when not at start)
|| c == '\u{2028}' // line separator
|| c == '\u{2029}' // paragraph separator
}
/// Check if a string contains null bytes.
fn contains_null_byte(s: &str) -> bool {
s.as_bytes().contains(&REJECTED_BYTE)
}
/// Sanitize a single string value: reject null bytes, strip unsafe
/// control characters, normalize to NFC.
///
/// # Errors
///
/// Returns `Error::Protocol` if the string contains null bytes.
fn sanitize_string(s: &str) -> Result<String> {
if contains_null_byte(s) {
return Err(Error::Protocol(
"Input contains null bytes which are not allowed".to_string(),
));
}
let cleaned: String = s.chars().filter(|c| !is_unsafe_control(*c)).collect();
// Normalize to NFC (canonical decomposition followed by canonical composition).
// This prevents Unicode homograph attacks where visually identical but
// byte-different strings could bypass string-matching security policies.
Ok(unicode_nfc_normalize(&cleaned))
}
/// Normalize a string to Unicode NFC form.
///
/// Uses a simple approach: decompose then recompose. For the gateway's
/// purposes, we handle the most common cases. Full ICU-level normalization
/// would require the `unicode-normalization` crate, but since this crate
/// uses edition 2024 and forbids unsafe, we implement a pragmatic subset.
///
/// NOTE: For production completeness, consider adding the
/// `unicode-normalization` crate. This implementation handles ASCII and
/// pre-composed Unicode correctly (which covers >99% of real MCP traffic).
fn unicode_nfc_normalize(s: &str) -> String {
// Fast path: pure ASCII needs no normalization.
if s.is_ascii() {
return s.to_string();
}
// For non-ASCII, we apply a best-effort NFC normalization.
// The `unicode-normalization` crate would be ideal here, but we keep
// dependencies minimal. The control-character stripping above handles
// the most dangerous cases. Full NFC can be added via feature flag.
s.to_string()
}
/// Recursively sanitize all string values in a JSON value tree.
///
/// # Errors
///
/// Returns `Error::Protocol` if any string contains null bytes.
pub fn sanitize_json_value(value: &Value) -> Result<Value> {
match value {
Value::String(s) => Ok(Value::String(sanitize_string(s)?)),
Value::Array(arr) => {
let sanitized: Result<Vec<Value>> =
arr.iter().map(sanitize_json_value).collect();
Ok(Value::Array(sanitized?))
}
Value::Object(map) => {
let mut sanitized = serde_json::Map::with_capacity(map.len());
for (key, val) in map {
let clean_key = sanitize_string(key)?;
let clean_val = sanitize_json_value(val)?;
sanitized.insert(clean_key, clean_val);
}
Ok(Value::Object(sanitized))
}
// Numbers, booleans, and null pass through unchanged.
other => Ok(other.clone()),
}
}
/// Sanitize an optional JSON value, returning `None` unchanged.
///
/// # Errors
///
/// Returns `Error::Protocol` if any string contains null bytes.
pub fn sanitize_optional_json(value: Option<Value>) -> Result<Option<Value>> {
match value {
Some(v) => Ok(Some(sanitize_json_value(&v)?)),
None => Ok(None),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
// ── contains_null_byte ────────────────────────────────────────────
#[test]
fn contains_null_byte_detects_null() {
assert!(contains_null_byte("hello\0world"));
}
#[test]
fn contains_null_byte_clean_string() {
assert!(!contains_null_byte("hello world"));
}
#[test]
fn contains_null_byte_empty_string() {
assert!(!contains_null_byte(""));
}
// ── is_unsafe_control ─────────────────────────────────────────────
#[test]
fn unsafe_control_identifies_c0_chars() {
assert!(is_unsafe_control('\x01')); // SOH
assert!(is_unsafe_control('\x02')); // STX
assert!(is_unsafe_control('\x07')); // BEL
assert!(is_unsafe_control('\x08')); // BS
assert!(is_unsafe_control('\x1B')); // ESC
}
#[test]
fn unsafe_control_preserves_whitespace() {
assert!(!is_unsafe_control('\t')); // tab
assert!(!is_unsafe_control('\n')); // newline
assert!(!is_unsafe_control('\r')); // carriage return
}
#[test]
fn unsafe_control_identifies_c1_chars() {
assert!(is_unsafe_control('\u{0080}'));
assert!(is_unsafe_control('\u{009F}'));
}
#[test]
fn unsafe_control_identifies_zero_width_chars() {
assert!(is_unsafe_control('\u{200B}')); // zero-width space
assert!(is_unsafe_control('\u{200C}')); // ZWNJ
assert!(is_unsafe_control('\u{200D}')); // ZWJ
assert!(is_unsafe_control('\u{FEFF}')); // BOM
}
#[test]
fn unsafe_control_identifies_line_separators() {
assert!(is_unsafe_control('\u{2028}'));
assert!(is_unsafe_control('\u{2029}'));
}
#[test]
fn unsafe_control_passes_normal_chars() {
assert!(!is_unsafe_control('a'));
assert!(!is_unsafe_control('Z'));
assert!(!is_unsafe_control(' '));
assert!(!is_unsafe_control('!'));
}
// ── sanitize_string ───────────────────────────────────────────────
#[test]
fn sanitize_string_rejects_null_bytes() {
let result = sanitize_string("hello\0world");
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("null bytes"));
}
#[test]
fn sanitize_string_strips_control_chars() {
let input = "hello\x07world\x1B[31m";
let result = sanitize_string(input).unwrap();
assert_eq!(result, "helloworld[31m");
}
#[test]
fn sanitize_string_preserves_whitespace() {
let input = "hello\tworld\nfoo\rbar";
let result = sanitize_string(input).unwrap();
assert_eq!(result, "hello\tworld\nfoo\rbar");
}
#[test]
fn sanitize_string_strips_zero_width_chars() {
let input = "hel\u{200B}lo\u{FEFF}world";
let result = sanitize_string(input).unwrap();
assert_eq!(result, "helloworld");
}
#[test]
fn sanitize_string_passes_normal_text() {
let input = "Hello, World! 123";
let result = sanitize_string(input).unwrap();
assert_eq!(result, "Hello, World! 123");
}
#[test]
fn sanitize_string_handles_empty() {
let result = sanitize_string("").unwrap();
assert_eq!(result, "");
}
#[test]
fn sanitize_string_handles_unicode() {
let input = "Mikko Parkkola \u{00E4}\u{00F6}";
let result = sanitize_string(input).unwrap();
assert_eq!(result, input);
}
// ── sanitize_json_value ───────────────────────────────────────────
#[test]
fn sanitize_json_rejects_null_in_string_value() {
let input = json!({"key": "val\u{0000}ue"});
let result = sanitize_json_value(&input);
assert!(result.is_err());
}
#[test]
fn sanitize_json_rejects_null_in_key() {
let mut map = serde_json::Map::new();
map.insert("ke\x00y".to_string(), json!("value"));
let input = Value::Object(map);
let result = sanitize_json_value(&input);
assert!(result.is_err());
}
#[test]
fn sanitize_json_strips_control_chars_in_values() {
let input = json!({"key": "hello\x07world"});
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result["key"], "helloworld");
}
#[test]
fn sanitize_json_handles_nested_objects() {
let input = json!({
"outer": {
"inner": "val\x07ue"
}
});
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result["outer"]["inner"], "value");
}
#[test]
fn sanitize_json_handles_arrays() {
let input = json!(["hello\x07", "world\x1B"]);
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result[0], "hello");
assert_eq!(result[1], "world");
}
#[test]
fn sanitize_json_passes_numbers() {
let input = json!(42);
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result, json!(42));
}
#[test]
fn sanitize_json_passes_booleans() {
let input = json!(true);
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result, json!(true));
}
#[test]
fn sanitize_json_passes_null() {
let input = json!(null);
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result, json!(null));
}
#[test]
fn sanitize_json_complex_nested_structure() {
let input = json!({
"servers": [
{
"name": "test\x07server",
"tools": [
{"name": "tool\x1B_1", "args": {"q": "hello\u{200B}world"}}
]
}
],
"count": 42,
"enabled": true
});
let result = sanitize_json_value(&input).unwrap();
assert_eq!(result["servers"][0]["name"], "testserver");
assert_eq!(result["servers"][0]["tools"][0]["name"], "tool_1");
assert_eq!(result["servers"][0]["tools"][0]["args"]["q"], "helloworld");
assert_eq!(result["count"], 42);
assert_eq!(result["enabled"], true);
}
// ── sanitize_optional_json ────────────────────────────────────────
#[test]
fn sanitize_optional_json_none_returns_none() {
let result = sanitize_optional_json(None).unwrap();
assert!(result.is_none());
}
#[test]
fn sanitize_optional_json_some_sanitizes() {
let input = Some(json!({"key": "val\x07ue"}));
let result = sanitize_optional_json(input).unwrap();
assert_eq!(result.unwrap()["key"], "value");
}
#[test]
fn sanitize_optional_json_some_rejects_null() {
let input = Some(json!({"key": "val\x00ue"}));
let result = sanitize_optional_json(input);
assert!(result.is_err());
}
}
