fakecloud_stepfunctions/

io_processing.rs

use serde_json::Value;

use crate::jsonpath::{JsonPath, PathError, Step};

/// A Step Functions error: `(error, cause)`.
pub type StatesError = (String, String);

/// Apply InputPath to extract a subset of the raw input.
/// - `None` or `Some("$")` -> return input unchanged
/// - `Some("null")` handled at call site (pass `{}`)
/// - any other JSONPath -> the selected value; a definite path that matches
///   nothing fails the state with `States.Runtime`, as on AWS
pub fn apply_input_path(input: &Value, path: Option<&str>) -> Result<Value, StatesError> {
    match path {
        None | Some("$") => Ok(input.clone()),
        Some(p) => resolve_reference(input, p),
    }
}

/// Apply OutputPath to extract a subset of the effective output.
/// Same semantics as InputPath.
pub fn apply_output_path(output: &Value, path: Option<&str>) -> Result<Value, StatesError> {
    match path {
        None | Some("$") => Ok(output.clone()),
        Some(p) => resolve_reference(output, p),
    }
}

/// Apply ResultPath to merge a state's result into the input.
/// - `None` or `Some("$")` -> result replaces input entirely
/// - `Some("null")` -> discard result, return original input
/// - `Some("$.foo")` -> set result at that path within input
///
/// A path that walks into a value that can't hold its next step (`$.a.b` with
/// `a: 1`) fails with `States.ReferencePathConflict`, as recorded.
pub fn apply_result_path(
    input: &Value,
    result: &Value,
    path: Option<&str>,
) -> Result<Value, StatesError> {
    match path {
        None | Some("$") => Ok(result.clone()),
        Some("null") => Ok(input.clone()),
        Some(p) => set_at_path(input, p, result),
    }
}

/// The error a `ResultPath` that can't be applied fails with.
pub const REFERENCE_PATH_CONFLICT: &str = "States.ReferencePathConflict";

/// As recorded: `Unable to apply step "b" to input {"a":1}`, naming the step and
/// the whole input.
fn reference_path_conflict(step: &Step, input: &Value) -> StatesError {
    let step = match step {
        Step::Name(name) => name.clone(),
        Step::Index(index) => index.to_string(),
    };
    (
        REFERENCE_PATH_CONFLICT.to_string(),
        format!("Unable to apply step \"{step}\" to input {input}"),
    )
}

/// Evaluate a JSONPath against `root` (see [`crate::jsonpath`]).
pub fn resolve_path(root: &Value, path: &str) -> Result<Value, PathError> {
    crate::jsonpath::evaluate(root, path)
}

/// Resolve a path, mapping a miss to the `States.Runtime` error AWS raises
/// for InputPath / OutputPath / ItemsPath / SecondsPath and friends:
/// `Invalid path '$.x' : No results for path: $['x']`.
pub fn resolve_reference(root: &Value, path: &str) -> Result<Value, StatesError> {
    resolve_with_context(root, None, path).map_err(|e| runtime_error(path_error_cause(path, &e)))
}

/// [`resolve_reference`] for a path that may address the context object
/// (`$$...`), as `ItemsPath`, `SecondsPath` and friends may.
pub fn resolve_reference_with_context(
    input: &Value,
    context: Option<&Value>,
    path: &str,
) -> Result<Value, StatesError> {
    resolve_with_context(input, context, path)
        .map_err(|e| runtime_error(path_error_cause(path, &e)))
}

/// Resolve a path that may address the context object (`$$...`). Without a
/// context object, `$$` paths resolve to null.
pub fn resolve_with_context(
    input: &Value,
    context: Option<&Value>,
    path: &str,
) -> Result<Value, PathError> {
    if let Some((name, within)) = crate::variables::split_reference(path) {
        let value = crate::variables::lookup(name)
            .ok_or_else(|| PathError::Invalid(format!("the variable '{name}' is not defined")))?;
        return resolve_path(&value, &within);
    }
    match path.strip_prefix("$$") {
        Some(rest) => match context {
            Some(ctx) => resolve_path(ctx, &format!("${rest}")),
            None => Ok(Value::Null),
        },
        None => resolve_path(input, path),
    }
}

pub fn path_error_cause(path: &str, err: &PathError) -> String {
    match err {
        PathError::NoResults(norm) => {
            format!("Invalid path '{path}' : No results for path: {norm}")
        }
        // As recorded for `$.missing` on an array: `Invalid path '$.missing' :
        // Property ['missing'] not found in path $`.
        PathError::MissingProperty { property, parent } => {
            format!("Invalid path '{path}' : Property {property} not found in path {parent}")
        }
        PathError::Invalid(msg) => format!("Invalid path '{path}' : {msg}"),
    }
}

/// A value as error messages quote it: a string's text, anything else as JSON.
pub(crate) fn value_text(value: &Value) -> String {
    match value {
        Value::String(text) => text.clone(),
        other => other.to_string(),
    }
}

pub fn runtime_error(cause: String) -> StatesError {
    ("States.Runtime".to_string(), cause)
}

/// Set a value at a reference path (`$.a.b`, `$['a'][0]`) within a JSON
/// structure. An index step descends into the array stored at that point,
/// growing it with nulls as needed. A step into a value that can't hold it (a
/// property of a non-object, an index of a non-array) is a conflict.
fn set_at_path(root: &Value, path: &str, value: &Value) -> Result<Value, StatesError> {
    let mut result = root.clone();
    let Some(steps) = JsonPath::parse(path).ok().and_then(|p| p.reference_steps()) else {
        return Ok(result);
    };

    // The empty container a missing intermediate node is created as.
    fn container_for(next: &Step) -> Value {
        match next {
            Step::Name(_) => serde_json::json!({}),
            Step::Index(_) => Value::Array(Vec::new()),
        }
    }

    /// Assign `value` at `steps` below `current`, or return the step that
    /// can't be taken.
    fn assign<'s>(current: &mut Value, steps: &'s [Step], value: &Value) -> Result<(), &'s Step> {
        let Some((step, rest)) = steps.split_first() else {
            *current = value.clone();
            return Ok(());
        };
        match step {
            Step::Name(name) => {
                let Some(obj) = current.as_object_mut() else {
                    return Err(step);
                };
                match rest.first() {
                    None => {
                        obj.insert(name.clone(), value.clone());
                        Ok(())
                    }
                    Some(next) => {
                        let child = obj
                            .entry(name.clone())
                            .or_insert_with(|| container_for(next));
                        assign(child, rest, value)
                    }
                }
            }
            Step::Index(idx) => {
                let Some(arr) = current.as_array_mut() else {
                    return Err(step);
                };
                {
                    let idx = if *idx < 0 {
                        match usize::try_from(arr.len() as i64 + idx) {
                            Ok(i) => i,
                            Err(_) => return Err(step),
                        }
                    } else {
                        *idx as usize
                    };
                    if arr.len() <= idx {
                        arr.resize(idx + 1, Value::Null);
                    }
                    if let Some(next) = rest.first() {
                        if arr[idx].is_null() {
                            arr[idx] = container_for(next);
                        }
                    }
                    assign(&mut arr[idx], rest, value)
                }
            }
        }
    }

    assign(&mut result, &steps, value).map_err(|step| reference_path_conflict(step, root))?;
    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn test_resolve_path_root() {
        let input = json!({"a": 1});
        assert_eq!(resolve_path(&input, "$").unwrap(), input);
    }

    #[test]
    fn test_resolve_path_simple_field() {
        let input = json!({"name": "hello", "value": 42});
        assert_eq!(resolve_path(&input, "$.name").unwrap(), json!("hello"));
        assert_eq!(resolve_path(&input, "$.value").unwrap(), json!(42));
    }

    #[test]
    fn test_resolve_path_nested() {
        let input = json!({"a": {"b": {"c": 99}}});
        assert_eq!(resolve_path(&input, "$.a.b.c").unwrap(), json!(99));
    }

    #[test]
    fn test_resolve_path_missing() {
        let input = json!({"a": 1});
        // Both recorded.
        let error = resolve_path(&input, "$.missing").unwrap_err();
        assert_eq!(
            path_error_cause("$.missing", &error),
            "Invalid path '$.missing' : No results for path: $['missing']"
        );
        let error = resolve_path(&json!([{}]), "$.missing").unwrap_err();
        assert_eq!(
            path_error_cause("$.missing", &error),
            "Invalid path '$.missing' : Property ['missing'] not found in path $"
        );
    }

    #[test]
    fn test_resolve_path_array_index() {
        let input = json!({"items": [10, 20, 30]});
        assert_eq!(resolve_path(&input, "$.items[0]").unwrap(), json!(10));
        assert_eq!(resolve_path(&input, "$.items[2]").unwrap(), json!(30));
    }

    #[test]
    fn test_apply_input_path_default() {
        let input = json!({"x": 1});
        assert_eq!(apply_input_path(&input, None).unwrap(), input);
        assert_eq!(apply_input_path(&input, Some("$")).unwrap(), input);
    }

    #[test]
    fn test_apply_result_path_default() {
        let input = json!({"x": 1});
        let result = json!({"y": 2});
        // Default: result replaces input
        assert_eq!(apply_result_path(&input, &result, None).unwrap(), result);
        assert_eq!(
            apply_result_path(&input, &result, Some("$")).unwrap(),
            result
        );
    }

    #[test]
    fn test_apply_result_path_null() {
        let input = json!({"x": 1});
        let result = json!({"y": 2});
        // null: discard result, keep input
        assert_eq!(
            apply_result_path(&input, &result, Some("null")).unwrap(),
            input
        );
    }

    #[test]
    fn test_apply_result_path_nested() {
        let input = json!({"x": 1});
        let result = json!("hello");
        let output = apply_result_path(&input, &result, Some("$.result")).unwrap();
        assert_eq!(output, json!({"x": 1, "result": "hello"}));
    }

    #[test]
    fn test_set_at_path_non_object_intermediate() {
        // An intermediate non-object can't take the next step: a conflict, as
        // recorded, never a panic.
        let input = json!({"x": 42});
        let result = json!("hello");
        let (error, cause) = apply_result_path(&input, &result, Some("$.x.nested")).unwrap_err();
        assert_eq!(error, REFERENCE_PATH_CONFLICT);
        assert_eq!(cause, r#"Unable to apply step "nested" to input {"x":42}"#);
    }

    // L8: ResultPath with an array-index segment must descend into the array,
    // not insert a literal key like "a[0]".
    #[test]
    fn test_set_at_path_array_index_leaf() {
        let input = json!({"a": [1, 2, 3]});
        let result = json!(99);
        let output = apply_result_path(&input, &result, Some("$.a[1]")).unwrap();
        assert_eq!(output, json!({"a": [1, 99, 3]}));
        // No stray literal "a[1]" key was created.
        assert!(output.get("a[1]").is_none());
    }

    #[test]
    fn test_set_at_path_array_index_then_field() {
        let input = json!({"items": [{"v": 1}]});
        let result = json!("done");
        let output = apply_result_path(&input, &result, Some("$.items[0].status")).unwrap();
        assert_eq!(output, json!({"items": [{"v": 1, "status": "done"}]}));
    }

    #[test]
    fn test_set_at_path_array_index_grows_array() {
        let input = json!({});
        let result = json!("x");
        let output = apply_result_path(&input, &result, Some("$.a[2]")).unwrap();
        assert_eq!(output, json!({"a": [null, null, "x"]}));
    }

    #[test]
    fn test_apply_output_path() {
        let output = json!({"a": 1, "b": 2});
        assert_eq!(apply_output_path(&output, Some("$.a")).unwrap(), json!(1));
        assert_eq!(apply_output_path(&output, None).unwrap(), output);
    }

    #[test]
    fn test_resolve_path_unclosed_bracket_does_not_panic() {
        // Malformed JSONPath: `[` with no trailing `]`. Previously this sliced
        // `[bracket_pos + 1 .. part.len() - 1]` which underflows -> panic.
        let input = json!({"arr": [1, 2, 3]});
        // No field literally named "arr[" exists, so this resolves to Null,
        // but the key requirement is that it must NOT panic.
        assert!(matches!(
            resolve_path(&input, "$.arr["),
            Err(PathError::Invalid(_))
        ));
    }

    #[test]
    fn test_resolve_path_multibyte_after_bracket_does_not_panic() {
        // Multibyte char where the close bracket would be. `part.len() - 1`
        // previously landed mid-char -> "byte index is not a char boundary".
        let input = json!({"x": [1, 2, 3]});
        assert!(matches!(
            resolve_path(&input, "$.x[é"),
            Err(PathError::Invalid(_))
        ));
        // Also the closed-but-multibyte-inner case.
        assert!(matches!(
            resolve_path(&input, "$.x[é]"),
            Err(PathError::Invalid(_))
        ));
    }

    #[test]
    fn test_resolve_path_empty_brackets_do_not_panic() {
        let input = json!({"x": [1, 2, 3]});
        assert!(matches!(
            resolve_path(&input, "$.x[]"),
            Err(PathError::Invalid(_))
        ));
    }

    #[test]
    fn test_resolve_path_bracket_only_segment() {
        // A bare `[` as an entire segment.
        let input = json!({"a": 1});
        assert!(matches!(
            resolve_path(&input, "$.["),
            Err(PathError::Invalid(_))
        ));
        assert!(matches!(
            resolve_path(&input, "$.]"),
            Err(PathError::Invalid(_))
        ));
    }

    #[test]
    fn test_split_path_segments_well_formed_index_still_works() {
        // Ensure the hardening did not break the happy path.
        let input = json!({"items": [10, 20, 30]});
        assert_eq!(resolve_path(&input, "$.items[1]").unwrap(), json!(20));
        let nested = json!({"a": {"b": [{"c": 7}]}});
        assert_eq!(resolve_path(&nested, "$.a.b[0].c").unwrap(), json!(7));
    }

    #[test]
    fn test_apply_input_path_malformed_does_not_panic() {
        let input = json!({"arr": [1, 2, 3]});
        // Exercised via the public apply_* entrypoints used by the interpreter.
        let (err, _) = apply_input_path(&input, Some("$.arr[")).unwrap_err();
        assert_eq!(err, "States.Runtime");
        assert!(apply_output_path(&input, Some("$.x[é")).is_err());
    }
}

#[cfg(test)]
mod missing_path_tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn missing_input_path_is_states_runtime() {
        let (err, cause) = apply_input_path(&json!({"a": 1}), Some("$.x")).unwrap_err();
        assert_eq!(err, "States.Runtime");
        assert_eq!(cause, "Invalid path '$.x' : No results for path: $['x']");
        let (err, _) = apply_output_path(&json!({"a": 1}), Some("$.a.b")).unwrap_err();
        assert_eq!(err, "States.Runtime");
        // A present null is fine.
        assert_eq!(
            apply_input_path(&json!({"a": null}), Some("$.a")).unwrap(),
            Value::Null
        );
    }

    #[test]
    fn full_jsonpath_syntax_in_input_path() {
        let input =
            json!({"items": [{"id": 1, "ok": true}, {"id": 2, "ok": false}], "k": {"x y": 5}});
        assert_eq!(
            apply_input_path(&input, Some("$['k']['x y']")).unwrap(),
            json!(5)
        );
        assert_eq!(
            apply_input_path(&input, Some("$.items[-1].id")).unwrap(),
            json!(2)
        );
        assert_eq!(
            apply_input_path(&input, Some("$.items[*].id")).unwrap(),
            json!([1, 2])
        );
        assert_eq!(
            apply_input_path(&input, Some("$.items[?(@.ok == true)].id")).unwrap(),
            json!([1])
        );
        assert_eq!(
            apply_input_path(&input, Some("$.items[0:1]")).unwrap(),
            json!([{"id": 1, "ok": true}])
        );
    }

    #[test]
    fn result_path_bracket_notation() {
        let out = apply_result_path(&json!({"a": {}}), &json!(1), Some("$['a']['b c']")).unwrap();
        assert_eq!(out, json!({"a": {"b c": 1}}));
        let out = apply_result_path(&json!({"a": [1, 2]}), &json!(9), Some("$.a[-1]")).unwrap();
        assert_eq!(out, json!({"a": [1, 9]}));
    }

    #[test]
    fn context_paths() {
        let ctx = json!({"Execution": {"Id": "arn"}});
        assert_eq!(
            resolve_with_context(&json!({}), Some(&ctx), "$$.Execution.Id").unwrap(),
            json!("arn")
        );
        assert!(resolve_with_context(&json!({}), Some(&ctx), "$$.Nope").is_err());
        assert_eq!(
            resolve_with_context(&json!({"a": 1}), None, "$.a").unwrap(),
            json!(1)
        );
    }
}

#[cfg(test)]
mod reference_path_conflict_tests {
    use super::*;
    use serde_json::json;

    /// Recorded: `ResultPath: $.a.b` over `{"a":1}`.
    #[test]
    fn a_step_into_a_scalar_is_a_conflict() {
        assert_eq!(
            apply_result_path(&json!({"a": 1}), &json!([{}]), Some("$.a.b")).unwrap_err(),
            (
                "States.ReferencePathConflict".to_string(),
                r#"Unable to apply step "b" to input {"a":1}"#.to_string()
            )
        );
        // Inferred: an index into a non-array, and a property of a scalar input.
        assert_eq!(
            apply_result_path(&json!({"a": {}}), &json!(1), Some("$.a[0]"))
                .unwrap_err()
                .0,
            REFERENCE_PATH_CONFLICT
        );
        assert!(apply_result_path(&json!(5), &json!(1), Some("$.x")).is_err());
        // Missing nodes are still created.
        assert_eq!(
            apply_result_path(&json!({}), &json!(1), Some("$.a.b")).unwrap(),
            json!({"a": {"b": 1}})
        );
    }
}