headless_lms_mcp/protocol/
headers.rs

1//! Standard request headers (`MCP-Protocol-Version`, `Mcp-Method`, `Mcp-Name`, `Mcp-Param-*`),
2//! which let rate limiting and logging route without parsing the body. They must agree with it.
3
4use serde_json::{Map, Value};
5
6const BASE64_PREFIX: &str = "=?base64?";
7const BASE64_SUFFIX: &str = "?=";
8
9/// Header naming the protocol revision of the request.
10pub(crate) const PROTOCOL_VERSION_HEADER: &str = "mcp-protocol-version";
11/// Header repeating the JSON-RPC method of the body.
12pub(crate) const METHOD_HEADER: &str = "mcp-method";
13/// Header repeating the tool or resource name of the body.
14pub(crate) const NAME_HEADER: &str = "mcp-name";
15/// Prefix of the headers that repeat the arguments a tool promotes to headers.
16pub(crate) const PARAM_HEADER_PREFIX: &str = "mcp-param-";
17
18/// Request headers with lowercase names.
19#[derive(Debug, Clone, Default)]
20pub struct RequestHeaders {
21    entries: Vec<(String, String)>,
22}
23
24impl RequestHeaders {
25    /// Headers from name-value pairs; names match case-insensitively.
26    pub fn new<N: AsRef<str>, V: Into<String>>(entries: impl IntoIterator<Item = (N, V)>) -> Self {
27        Self {
28            entries: entries
29                .into_iter()
30                .map(|(n, v)| (n.as_ref().to_ascii_lowercase(), v.into()))
31                .collect(),
32        }
33    }
34
35    /// The single value of `name`; an error when it is repeated.
36    pub(crate) fn single(&self, name: &str) -> Result<Option<&str>, String> {
37        let name = name.to_ascii_lowercase();
38        let mut values = self
39            .entries
40            .iter()
41            .filter(|(n, _)| *n == name)
42            .map(|(_, v)| v.as_str());
43        let first = values.next();
44        if values.next().is_some() {
45            return Err(format!("duplicate {name} header"));
46        }
47        Ok(first)
48    }
49}
50
51fn decode(value: &str) -> Option<String> {
52    use base64::Engine;
53    match value
54        .strip_prefix(BASE64_PREFIX)
55        .and_then(|v| v.strip_suffix(BASE64_SUFFIX))
56    {
57        Some(encoded) => base64::engine::general_purpose::STANDARD
58            .decode(encoded)
59            .ok()
60            .and_then(|bytes| String::from_utf8(bytes).ok()),
61        None => Some(value.to_string()),
62    }
63}
64
65fn primitive(value: &Value) -> Option<String> {
66    match value {
67        Value::String(s) => Some(s.clone()),
68        Value::Bool(b) => Some(b.to_string()),
69        Value::Number(n) => Some(n.to_string()),
70        _ => None,
71    }
72}
73
74/// Checks the routing headers against the body. `header_params` are the called tool's
75/// `(argument, header suffix)` pairs.
76pub(crate) fn validate(
77    headers: &RequestHeaders,
78    method: &str,
79    params: &Map<String, Value>,
80    header_params: &[(String, String)],
81) -> Result<(), String> {
82    match headers.single(METHOD_HEADER)? {
83        None => return Err("missing Mcp-Method header".into()),
84        Some(value) if value != method => {
85            return Err(format!(
86                "Mcp-Method header \"{value}\" doesn't match the body method \"{method}\""
87            ));
88        }
89        Some(_) => {}
90    }
91    if method == "tools/call" {
92        let name = params
93            .get("name")
94            .and_then(Value::as_str)
95            .unwrap_or_default();
96        match headers.single(NAME_HEADER)? {
97            None => return Err("missing Mcp-Name header for tools/call".into()),
98            Some(raw) => {
99                let decoded = decode(raw).ok_or("Mcp-Name header isn't valid base64")?;
100                if decoded != name {
101                    return Err(format!(
102                        "Mcp-Name header \"{decoded}\" doesn't match the tool \"{name}\""
103                    ));
104                }
105            }
106        }
107        let arguments = params.get("arguments");
108        for (argument, suffix) in header_params {
109            let header = format!("{PARAM_HEADER_PREFIX}{}", suffix.to_ascii_lowercase());
110            let body = arguments.and_then(|a| a.get(argument)).and_then(primitive);
111            match (headers.single(&header)?, body) {
112                (None, None) => {}
113                (Some(_), None) => {
114                    return Err(format!(
115                        "Mcp-Param-{suffix} header sent without a {argument} argument"
116                    ));
117                }
118                (None, Some(_)) => {
119                    return Err(format!("missing Mcp-Param-{suffix} header for {argument}"));
120                }
121                (Some(raw), Some(expected)) => {
122                    let decoded = decode(raw)
123                        .ok_or_else(|| format!("Mcp-Param-{suffix} header isn't valid base64"))?;
124                    if decoded != expected {
125                        return Err(format!(
126                            "Mcp-Param-{suffix} header \"{decoded}\" doesn't match {argument} \"{expected}\""
127                        ));
128                    }
129                }
130            }
131        }
132    }
133    Ok(())
134}
135
136#[cfg(test)]
137mod tests {
138    use serde_json::json;
139
140    use super::*;
141
142    fn headers(entries: &[(&str, &str)]) -> RequestHeaders {
143        RequestHeaders::new(entries.iter().map(|(n, v)| (*n, v.to_string())))
144    }
145
146    #[test]
147    fn method_name_and_params_must_match() {
148        let params =
149            json!({"name": "get_page_content", "arguments": {"course": "c1", "page": "p1"}});
150        let params = params.as_object().unwrap();
151        let course = vec![("course".to_string(), "Course".to_string())];
152        let ok = headers(&[
153            ("Mcp-Method", "tools/call"),
154            ("Mcp-Name", "get_page_content"),
155            ("Mcp-Param-Course", "c1"),
156        ]);
157        assert!(validate(&ok, "tools/call", params, &course).is_ok());
158        let wrong_course = headers(&[
159            ("Mcp-Method", "tools/call"),
160            ("Mcp-Name", "get_page_content"),
161            ("Mcp-Param-Course", "c2"),
162        ]);
163        assert!(validate(&wrong_course, "tools/call", params, &course).is_err());
164        let missing_name = headers(&[("Mcp-Method", "tools/call")]);
165        assert!(validate(&missing_name, "tools/call", params, &course).is_err());
166        let wrong_method = headers(&[("Mcp-Method", "tools/list")]);
167        assert!(validate(&wrong_method, "tools/call", params, &course).is_err());
168        let duplicated = headers(&[("Mcp-Method", "tools/list"), ("mcp-method", "tools/list")]);
169        assert!(validate(&duplicated, "tools/list", &Map::new(), &[]).is_err());
170    }
171
172    #[test]
173    fn base64_wrapped_values() {
174        assert_eq!(decode("=?base64?w6RpdGk=?="), Some("äiti".to_string()));
175        assert_eq!(decode("plain"), Some("plain".to_string()));
176    }
177}
headless_lms_mcp/protocol/headers.rs

headless_lms_mcp/protocol/
headers.rs
