magi-code 0.63.1

use crate::mcp::{McpError, McpResult};
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use std::collections::BTreeMap;

pub(crate) const REDACTED_HEADER_VALUE: &str = "[REDACTED]";

pub(crate) fn resolve_http_headers(headers: &BTreeMap<String, String>) -> McpResult<HeaderMap> {
    let mut resolved = HeaderMap::new();
    for (name, value) in headers {
        let header_name = HeaderName::from_bytes(name.as_bytes())
            .map_err(|_| McpError::Config(format!("invalid MCP HTTP header name '{name}'")))?;
        let value = match parse_env_header_ref(value) {
            Some(var_name) => std::env::var(var_name).map_err(|_| {
                McpError::Config(format!(
                    "MCP HTTP header '{name}' references unset environment variable '{var_name}'"
                ))
            })?,
            None => {
                if value.starts_with("{env:") {
                    return Err(McpError::Config(format!(
                        "MCP HTTP header '{name}' has invalid environment reference syntax"
                    )));
                }
                if is_sensitive_header(name) {
                    return Err(McpError::Config(format!(
                        "MCP HTTP header '{name}' must reference an environment variable"
                    )));
                }
                value.clone()
            }
        };
        let header_value = HeaderValue::from_str(&value)
            .map_err(|_| McpError::Config(format!("invalid MCP HTTP header value for '{name}'")))?;
        resolved.insert(header_name, header_value);
    }
    Ok(resolved)
}

pub(crate) fn parse_env_header_ref(value: &str) -> Option<&str> {
    value
        .strip_prefix("{env:")
        .and_then(|rest| rest.strip_suffix('}'))
        .filter(|name| is_valid_env_name(name))
}

pub(crate) fn is_env_header_ref_syntax(value: &str) -> bool {
    value.starts_with("{env:") && value.ends_with('}')
}

pub(crate) fn is_valid_env_name(name: &str) -> bool {
    let mut bytes = name.bytes();
    let Some(first) = bytes.next() else {
        return false;
    };
    if !(first.is_ascii_alphabetic() || first == b'_') {
        return false;
    }
    bytes.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
}

pub(crate) fn is_sensitive_header(name: &str) -> bool {
    let lower = name.to_ascii_lowercase();
    lower == "authorization"
        || lower == "proxy-authorization"
        || lower.contains("token")
        || lower.contains("secret")
        || lower.contains("api-key")
}

pub(crate) fn redact_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
    headers
        .keys()
        .map(|name| (name.clone(), REDACTED_HEADER_VALUE.to_string()))
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn resolves_literal_and_env_headers() {
        let env = crate::test_support::env::env_lock();
        env.set_var("MCP_HEADER_TEST_TOKEN", "secret-token");
        let mut headers = BTreeMap::new();
        headers.insert(
            "Authorization".to_string(),
            "{env:MCP_HEADER_TEST_TOKEN}".to_string(),
        );
        headers.insert("X-Team".to_string(), "platform".to_string());

        let resolved = resolve_http_headers(&headers).unwrap();

        assert_eq!(resolved.get("authorization").unwrap(), "secret-token");
        assert_eq!(resolved.get("x-team").unwrap(), "platform");
        env.remove_var("MCP_HEADER_TEST_TOKEN");
    }

    #[test]
    fn missing_env_error_does_not_echo_value() {
        let env = crate::test_support::env::env_lock();
        env.remove_var("MCP_HEADER_TEST_MISSING");
        let mut headers = BTreeMap::new();
        headers.insert(
            "Authorization".to_string(),
            "{env:MCP_HEADER_TEST_MISSING}".to_string(),
        );

        let error = resolve_http_headers(&headers).unwrap_err().to_string();

        assert!(error.contains("MCP_HEADER_TEST_MISSING"), "{error}");
        assert!(!error.contains("secret-token"), "{error}");
    }

    #[test]
    fn rejects_literal_sensitive_header_values() {
        let mut headers = BTreeMap::new();
        headers.insert("Authorization".to_string(), "literal-token".to_string());

        let error = resolve_http_headers(&headers).unwrap_err().to_string();

        assert!(
            error.contains("must reference an environment variable"),
            "{error}"
        );
        assert!(!error.contains("literal-token"), "{error}");
    }

    #[test]
    fn rejects_malformed_env_header_references() {
        let mut headers = BTreeMap::new();
        headers.insert("X-Team".to_string(), "{env:NOT-A-VAR}".to_string());

        let error = resolve_http_headers(&headers).unwrap_err().to_string();

        assert!(
            error.contains("invalid environment reference syntax"),
            "{error}"
        );
        assert!(!error.contains("NOT-A-VAR"), "{error}");
    }

    #[test]
    fn detects_sensitive_headers_case_insensitively() {
        for name in [
            "Authorization",
            "proxy-authorization",
            "X-Access-Token",
            "x-client-secret",
            "x-api-key",
        ] {
            assert!(is_sensitive_header(name), "{name}");
        }
        assert!(!is_sensitive_header("X-Team"));
    }

    #[test]
    fn redacts_all_header_values() {
        let mut headers = BTreeMap::new();
        headers.insert("X-Team".to_string(), "platform".to_string());
        headers.insert("Authorization".to_string(), "{env:TOKEN}".to_string());

        let redacted = redact_headers(&headers);

        assert_eq!(redacted["X-Team"], REDACTED_HEADER_VALUE);
        assert_eq!(redacted["Authorization"], REDACTED_HEADER_VALUE);
        assert!(!format!("{redacted:?}").contains("platform"));
    }
}