crank/crates/crank-mapping/src/jsonpath.rs

use serde_json::Value;

use crate::MappingError;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum JsonPathRoot {
    Mcp,
    RequestPath,
    RequestQuery,
    RequestHeaders,
    RequestBody,
    ResponseBody,
    ResponseData,
    Output,
}

impl JsonPathRoot {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Mcp => "mcp",
            Self::RequestPath => "request.path",
            Self::RequestQuery => "request.query",
            Self::RequestHeaders => "request.headers",
            Self::RequestBody => "request.body",
            Self::ResponseBody => "response.body",
            Self::ResponseData => "response.data",
            Self::Output => "output",
        }
    }

    pub fn root_fields(&self) -> &'static [&'static str] {
        match self {
            Self::Mcp => &["mcp"],
            Self::RequestPath => &["request", "path"],
            Self::RequestQuery => &["request", "query"],
            Self::RequestHeaders => &["request", "headers"],
            Self::RequestBody => &["request", "body"],
            Self::ResponseBody => &["response", "body"],
            Self::ResponseData => &["response", "data"],
            Self::Output => &["output"],
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum JsonPathSegment {
    Field(String),
    Index(usize),
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct JsonPath {
    pub root: JsonPathRoot,
    pub segments: Vec<JsonPathSegment>,
}

impl JsonPath {
    pub fn parse(path: &str) -> Result<Self, MappingError> {
        let body = path
            .strip_prefix("$.")
            .ok_or_else(|| MappingError::InvalidJsonPath {
                path: path.to_owned(),
            })?;

        let (root, rest) = match_root(body).ok_or_else(|| MappingError::InvalidJsonPath {
            path: path.to_owned(),
        })?;

        if !(rest.is_empty() || rest.starts_with('.') || rest.starts_with('[')) {
            return Err(MappingError::InvalidJsonPath {
                path: path.to_owned(),
            });
        }

        let segments = parse_segments(rest, path)?;

        Ok(Self { root, segments })
    }

    pub fn to_string_path(&self) -> String {
        let mut path = format!("$.{}", self.root.as_str());

        for segment in &self.segments {
            match segment {
                JsonPathSegment::Field(field) => {
                    path.push('.');
                    path.push_str(field);
                }
                JsonPathSegment::Index(index) => {
                    path.push('[');
                    path.push_str(&index.to_string());
                    path.push(']');
                }
            }
        }

        path
    }

    pub fn read<'a>(&self, value: &'a Value) -> Option<&'a Value> {
        let mut current = value;

        for field in self.root.root_fields() {
            current = current.get(*field)?;
        }

        for segment in &self.segments {
            current = match segment {
                JsonPathSegment::Field(field) => current.get(field)?,
                JsonPathSegment::Index(index) => current.get(*index)?,
            };
        }

        Some(current)
    }
}

fn match_root(input: &str) -> Option<(JsonPathRoot, &str)> {
    const ROOTS: [(&str, JsonPathRoot); 8] = [
        ("request.headers", JsonPathRoot::RequestHeaders),
        ("response.body", JsonPathRoot::ResponseBody),
        ("response.data", JsonPathRoot::ResponseData),
        ("request.path", JsonPathRoot::RequestPath),
        ("request.query", JsonPathRoot::RequestQuery),
        ("request.body", JsonPathRoot::RequestBody),
        ("output", JsonPathRoot::Output),
        ("mcp", JsonPathRoot::Mcp),
    ];

    for (prefix, root) in ROOTS {
        if let Some(rest) = input.strip_prefix(prefix) {
            return Some((root, rest));
        }
    }

    None
}

fn parse_segments(mut input: &str, original: &str) -> Result<Vec<JsonPathSegment>, MappingError> {
    let mut segments = Vec::new();

    while !input.is_empty() {
        if let Some(rest) = input.strip_prefix('.') {
            let end = rest.find(['.', '[']).unwrap_or(rest.len());
            let field = &rest[..end];

            if field.is_empty() || !is_valid_field(field) {
                return Err(MappingError::InvalidJsonPath {
                    path: original.to_owned(),
                });
            }

            segments.push(JsonPathSegment::Field(field.to_owned()));
            input = &rest[end..];
            continue;
        }

        if let Some(rest) = input.strip_prefix('[') {
            let Some(end) = rest.find(']') else {
                return Err(MappingError::InvalidJsonPath {
                    path: original.to_owned(),
                });
            };

            let index =
                rest[..end]
                    .parse::<usize>()
                    .map_err(|_| MappingError::InvalidJsonPath {
                        path: original.to_owned(),
                    })?;

            segments.push(JsonPathSegment::Index(index));
            input = &rest[end + 1..];
            continue;
        }

        return Err(MappingError::InvalidJsonPath {
            path: original.to_owned(),
        });
    }

    Ok(segments)
}

fn is_valid_field(value: &str) -> bool {
    value
        .chars()
        .all(|char| char.is_ascii_alphanumeric() || char == '_' || char == '-')
}

#[cfg(test)]
mod tests {
    use super::{JsonPath, JsonPathRoot, JsonPathSegment};

    #[test]
    fn parses_nested_jsonpath_with_array_index() {
        let path = JsonPath::parse("$.request.body.contacts[0].email").unwrap();

        assert_eq!(path.root, JsonPathRoot::RequestBody);
        assert_eq!(
            path.segments,
            vec![
                JsonPathSegment::Field("contacts".to_owned()),
                JsonPathSegment::Index(0),
                JsonPathSegment::Field("email".to_owned())
            ]
        );
    }

    #[test]
    fn rejects_unknown_root() {
        let error = JsonPath::parse("$.unknown.field").unwrap_err();

        assert_eq!(
            error,
            crate::MappingError::InvalidJsonPath {
                path: "$.unknown.field".to_owned(),
            }
        );
    }
}