Skip to main content

mcp_usage_tower/
classify.rs

1//! Allocation-free classification of mandatory MCP metadata headers.
2
3use std::borrow::Cow;
4use std::fmt;
5
6use mcp_usage_core::{Method, name, validate_header};
7
8/// Trusted protocol metadata lifted from request headers.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ProtocolHeaders<'a> {
11    /// Classified method. Known core methods do not allocate.
12    pub method: Method,
13    /// Decoded `Mcp-Name`, borrowed for ordinary ASCII values.
14    pub name: Option<Cow<'a, str>>,
15}
16
17/// Why mandatory MCP request metadata could not be trusted.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum ClassificationError {
20    /// Protocol version is missing, legacy, or malformed.
21    UntrustedVersion(String),
22    /// `Mcp-Method` was absent.
23    MissingMethod,
24    /// A name-bearing method omitted `Mcp-Name`.
25    MissingName,
26    /// The sentinel-encoded name was malformed.
27    InvalidName(String),
28}
29
30impl fmt::Display for ClassificationError {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        match self {
33            Self::UntrustedVersion(reason) => write!(f, "untrusted MCP headers: {reason}"),
34            Self::MissingMethod => f.write_str("missing Mcp-Method header"),
35            Self::MissingName => f.write_str("missing Mcp-Name header for name-bearing method"),
36            Self::InvalidName(reason) => write!(f, "invalid Mcp-Name header: {reason}"),
37        }
38    }
39}
40
41impl std::error::Error for ClassificationError {}
42
43/// Validate and classify raw MCP header values.
44///
45/// The common path (known method and plain ASCII name) performs no allocation.
46///
47/// # Errors
48///
49/// Returns [`ClassificationError`] for untrusted protocol versions, missing
50/// mandatory metadata, or malformed Base64 sentinel values.
51pub fn classify_protocol_headers<'a>(
52    version_header: Option<&str>,
53    method_header: Option<&str>,
54    name_header: Option<&'a str>,
55) -> Result<ProtocolHeaders<'a>, ClassificationError> {
56    if let Err(reason) = validate_header(version_header) {
57        return Err(ClassificationError::UntrustedVersion(reason.to_string()));
58    }
59    let method = Method::parse(method_header.ok_or(ClassificationError::MissingMethod)?);
60    if method.carries_name() && name_header.is_none() {
61        return Err(ClassificationError::MissingName);
62    }
63    let name = name_header
64        .map(name::decode)
65        .transpose()
66        .map_err(|error| ClassificationError::InvalidName(error.to_string()))?;
67    Ok(ProtocolHeaders { method, name })
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn common_path_borrows_the_name() {
76        let classified =
77            classify_protocol_headers(Some("2026-07-28"), Some("tools/call"), Some("weather"))
78                .unwrap();
79        assert_eq!(classified.method, Method::ToolsCall);
80        assert!(matches!(classified.name, Some(Cow::Borrowed("weather"))));
81    }
82
83    #[test]
84    fn legacy_requests_are_rejected_before_headers_are_trusted() {
85        let error =
86            classify_protocol_headers(Some("2025-11-25"), Some("tools/call"), Some("cheap_tool"))
87                .unwrap_err();
88        assert!(matches!(error, ClassificationError::UntrustedVersion(_)));
89    }
90
91    #[test]
92    fn name_is_required_only_where_the_spec_requires_it() {
93        assert!(matches!(
94            classify_protocol_headers(Some("2026-07-28"), Some("tools/call"), None),
95            Err(ClassificationError::MissingName)
96        ));
97        assert!(classify_protocol_headers(Some("2026-07-28"), Some("tools/list"), None).is_ok());
98    }
99
100    #[test]
101    fn header_hot_path_has_no_body_dependency() {
102        // Invalid JSON is used here: classification still succeeds because the
103        // common request path never deserializes a body.
104        let invalid_body = b"this is not JSON";
105        let classified =
106            classify_protocol_headers(Some("2026-07-28"), Some("tools/call"), Some("weather"))
107                .unwrap();
108        assert_eq!(classified.method, Method::ToolsCall);
109        assert_eq!(invalid_body, b"this is not JSON");
110    }
111}