perimeterx-fastly-enforcer 2.2.2

use crate::modules::pxconstants::{APPLICATION_JSON, MCP_BODY_MAX_LENGTH, MCP_SESSION_ID_HEADER};
use crate::pxconfig::PXConfig;
use crate::pxcontext::PXContext;
use fastly::http::{Method, StatusCode};
use fastly::{Request, Response};
use serde_json::Value;

/// MCP metadata extracted from a matching Agentic Trust request.
#[derive(Default, Debug, Clone)]
pub struct AgenticTrustData {
    pub mcp_method: Option<String>,
    pub mcp_tool_name: Option<String>,
    pub mcp_tool_argument_keys: Option<String>,
    pub mcp_session_id: Option<String>,
    pub mcp_http_method: Option<String>,
}

fn first_json_rpc_message(value: &Value) -> Option<&Value> {
    if let Some(arr) = value.as_array() {
        arr.first()
    } else {
        Some(value)
    }
}

fn argument_keys_joined(arguments: &Value) -> Option<String> {
    let obj = arguments.as_object()?;
    if obj.is_empty() {
        return None;
    }
    let keys: Vec<&str> = obj.keys().map(String::as_str).collect();
    Some(keys.join(","))
}

/// Parse a JSON-RPC body and populate MCP fields on `data`. Parse failures are ignored.
pub(crate) fn extract_from_body(body: &[u8], data: &mut AgenticTrustData) {
    if body.is_empty() {
        return;
    }

    let value: Value = match serde_json::from_slice(body) {
        Ok(v) => v,
        Err(_) => return,
    };

    if value.is_null() {
        return;
    }

    let message = match first_json_rpc_message(&value) {
        Some(m) if m.is_object() => m,
        _ => return,
    };

    let params = match message.get("params") {
        Some(p) => p,
        None => return,
    };

    let jsonrpc = match message.get("jsonrpc") {
        Some(j) => j,
        None => return,
    };
    // only support jsonrpc 2.0
    if jsonrpc.as_str() != Some("2.0") {
        return;
    }

    if let Some(method) = message.get("method").and_then(Value::as_str) {
        data.mcp_method = Some(method.to_owned());
    }
    if let Some(name) = params.get("name").and_then(Value::as_str) {
        data.mcp_tool_name = Some(name.to_owned());
    }

    if let Some(arguments) = params.get("arguments") {
        if let Some(keys) = argument_keys_joined(arguments) {
            data.mcp_tool_argument_keys = Some(keys);
        }
    }
}

/// Enrich `ctx` with MCP metadata when the request matches the configured MCP endpoint.
pub fn enrich_context_from_request(req: &mut Request, conf: &PXConfig, ctx: &mut PXContext) {
    if !conf.agentic_trust_enabled {
        return;
    }

    if req.get_path() != conf.agentic_trust_mcp_endpoint_path {
        return;
    }

    let mut data = AgenticTrustData {
        ..AgenticTrustData::default()
    };

    if let Some(session_id) = req.get_header_str_lossy(MCP_SESSION_ID_HEADER) {
        let session_id = session_id.into_owned();
        if !session_id.is_empty() {
            data.mcp_session_id = Some(session_id);
        }
    }

    let method = req.get_method_str().to_string();

    if req.get_method() == Method::POST {
        let is_json = req
            .get_header_str_lossy("content-type")
            .map(|v| v.into_owned())
            .unwrap_or_default()
            .contains(APPLICATION_JSON);
        if is_json {
            let should_extract_from_body = match req.get_content_length() {
                Some(content_length) => content_length <= MCP_BODY_MAX_LENGTH,
                None => true,
            };

            if should_extract_from_body {
                let body = req.get_body_prefix_mut(MCP_BODY_MAX_LENGTH);
                extract_from_body(body.as_slice(), &mut data);
                if data.mcp_method.is_some() {
                    data.mcp_http_method = Some(method);
                }
            }
        }
    }

    ctx.agentic_trust_data = Some(data);
}

// Determine if activities should be postponed based on Agentic Trust data
// postpone for "method":"initialize"
pub fn should_postpone_activities(ctx: &PXContext) -> bool {
    ctx.agentic_trust_data.as_ref().is_some_and(|data| {
        data.mcp_method
            .as_ref()
            .is_some_and(|method| method == "initialize")
    })
}

// extract mcp-session-id from response header
pub fn enrich_context_from_response(resp: &Response, conf: &PXConfig, ctx: &mut PXContext) {
    if !conf.agentic_trust_enabled {
        return;
    }

    if resp.get_status() != StatusCode::OK {
        return;
    }

    if ctx.agentic_trust_data.is_none() {
        return;
    }

    let content_type = resp
        .get_header_str_lossy("content-type")
        .map(|v| v.into_owned())
        .unwrap_or_default();
    if content_type.is_empty() {
        return;
    }

    if !content_type.contains(APPLICATION_JSON) {
        return;
    }

    let mcp_session_id = resp
        .get_header_str_lossy(MCP_SESSION_ID_HEADER)
        .map(|v| v.into_owned())
        .unwrap_or_default();
    if mcp_session_id.is_empty() {
        return;
    }

    if let Some(data) = ctx.agentic_trust_data.as_mut() {
        data.mcp_session_id = Some(mcp_session_id);
    }
}