newton-core 0.4.16

use serde::{Deserialize, Serialize};
use thiserror::Error;

/// Protocol-level directives extracted from the `_newton` reserved key in `wasmArgs`.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct NewtonDirectives {
    /// IPFS CID of a TLSNotary presentation to fetch and verify.
    #[serde(alias = "proofCid", skip_serializing_if = "Option::is_none")]
    pub proof_cid: Option<String>,
    /// Proof type identifier (for example, `"tlsn"`).
    #[serde(alias = "proofType", skip_serializing_if = "Option::is_none")]
    pub proof_type: Option<String>,
    /// Multiple proof CIDs (vectorized). Alias: proofCids.
    /// If both proof_cid and proof_cids are set, proof_cid is prepended.
    #[serde(alias = "proofCids", default, skip_serializing_if = "Option::is_none")]
    pub proof_cids: Option<Vec<String>>,
    /// Inline ephemeral privacy data: base64-encoded SecureEnvelope JSON strings.
    /// Operators decrypt locally (direct HPKE or threshold partial DH).
    /// Max 10 envelopes per request. Decrypted data injected into Rego `data.privacy.*`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub privacy: Option<Vec<String>>,
    /// Additional forward-compatible `_newton` directive fields.
    #[serde(default, flatten)]
    pub extra: serde_json::Map<String, serde_json::Value>,
}

/// Maximum number of inline privacy envelopes per task request.
pub const MAX_INLINE_PRIVACY_ENVELOPES: usize = 10;

impl NewtonDirectives {
    /// Returns merged, deduplicated list of all proof CIDs.
    /// If both `proof_cid` and `proof_cids` are set, `proof_cid` is prepended.
    pub fn all_proof_cids(&self) -> Vec<&str> {
        let mut cids = Vec::new();
        if let Some(ref cid) = self.proof_cid {
            cids.push(cid.as_str());
        }
        if let Some(ref cids_vec) = self.proof_cids {
            for cid in cids_vec {
                if !cids.contains(&cid.as_str()) {
                    cids.push(cid.as_str());
                }
            }
        }
        cids
    }

    /// Validates inline privacy envelope count is within bounds.
    pub fn validate_privacy_count(&self) -> Result<(), WasmArgsError> {
        if let Some(ref envelopes) = self.privacy {
            if envelopes.len() > MAX_INLINE_PRIVACY_ENVELOPES {
                return Err(WasmArgsError::TooManyPrivacyEnvelopes {
                    count: envelopes.len(),
                    max: MAX_INLINE_PRIVACY_ENVELOPES,
                });
            }
        }
        Ok(())
    }
}

/// Errors while parsing or normalizing `_newton` directives in `wasmArgs`.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum WasmArgsError {
    /// The request provided both alias and canonical keys with different values.
    #[error(
        "conflicting _newton aliases for `{canonical}`: `{canonical}`={canonical_value} but `{alias}`={alias_value}"
    )]
    ConflictingAlias {
        /// The alias key (for example `proofCid`).
        alias: &'static str,
        /// The canonical key (for example `proof_cid`).
        canonical: &'static str,
        /// Serialized alias value.
        alias_value: String,
        /// Serialized canonical value.
        canonical_value: String,
    },
    /// The request provided a top-level `proof_cid` that conflicts with `_newton.proof_cid`.
    #[error("conflicting proof CID values: request proof_cid={request_value} but _newton.proof_cid={existing_value}")]
    ConflictingProofCid {
        /// Serialized value already present under `_newton.proof_cid`.
        existing_value: String,
        /// Serialized request-level proof CID value.
        request_value: String,
    },
    /// Too many inline privacy envelopes in `_newton.privacy`.
    #[error("too many inline privacy envelopes: {count} exceeds max {max}")]
    TooManyPrivacyEnvelopes {
        /// Actual number of envelopes provided.
        count: usize,
        /// Maximum allowed number of envelopes.
        max: usize,
    },
}

fn normalize_newton_aliases(value: &mut serde_json::Value) -> Result<(), WasmArgsError> {
    let Some(obj) = value.as_object_mut() else {
        return Ok(());
    };

    normalize_newton_alias(obj, "proofCid", "proof_cid")?;
    normalize_newton_alias(obj, "proofType", "proof_type")?;
    normalize_newton_alias(obj, "proofCids", "proof_cids")?;
    Ok(())
}

fn normalize_newton_alias(
    obj: &mut serde_json::Map<String, serde_json::Value>,
    alias: &'static str,
    canonical: &'static str,
) -> Result<(), WasmArgsError> {
    if let (Some(alias_value), Some(canonical_value)) = (obj.get(alias), obj.get(canonical)) {
        if alias_value != canonical_value {
            return Err(WasmArgsError::ConflictingAlias {
                alias,
                canonical,
                alias_value: alias_value.to_string(),
                canonical_value: canonical_value.to_string(),
            });
        }
    }

    if let Some(alias_value) = obj.remove(alias) {
        obj.entry(canonical.to_string()).or_insert(alias_value);
    }

    Ok(())
}

fn normalize_embedded_newton_namespace(
    json: &mut serde_json::Value,
) -> Result<Option<serde_json::Value>, WasmArgsError> {
    let Some(obj) = json.as_object_mut() else {
        return Ok(None);
    };

    let Some(mut value) = obj.remove("_newton") else {
        return Ok(None);
    };

    normalize_newton_aliases(&mut value)?;
    Ok(Some(value))
}

/// Parses raw `wasmArgs` bytes into Newton directives plus WASM passthrough bytes.
///
/// The `_newton` reserved object is removed from the returned passthrough payload.
/// If `wasm_args` is not valid JSON or does not contain `_newton`, this returns
/// default directives and the original bytes unchanged.
pub fn parse_wasm_args(wasm_args: &[u8]) -> Result<(NewtonDirectives, Vec<u8>), WasmArgsError> {
    let Ok(mut json) = serde_json::from_slice::<serde_json::Value>(wasm_args) else {
        return Ok((NewtonDirectives::default(), wasm_args.to_vec()));
    };

    let directives = normalize_embedded_newton_namespace(&mut json)?
        .and_then(|value| serde_json::from_value::<NewtonDirectives>(value).ok())
        .unwrap_or_default();

    directives.validate_privacy_count()?;

    let passthrough = serde_json::to_vec(&json).unwrap_or_else(|_| wasm_args.to_vec());
    Ok((directives, passthrough))
}

/// Injects `proof_cid` into the `_newton` namespace of `wasm_args`.
///
/// This validates any existing `_newton` aliases first. If both `proofCid` and
/// `proof_cid` are present and differ, or if an existing `_newton.proof_cid`
/// conflicts with the request-level `proof_cid`, this returns an error instead
/// of silently choosing one value.
pub fn inject_proof_cid_into_wasm_args(
    wasm_args: Option<alloy::primitives::Bytes>,
    proof_cid: Option<&str>,
) -> Result<Option<alloy::primitives::Bytes>, WasmArgsError> {
    let mut json: serde_json::Value = wasm_args
        .as_ref()
        .and_then(|b| serde_json::from_slice(b).ok())
        .unwrap_or_else(|| serde_json::json!({}));

    // If existing wasm_args isn't a JSON object, wrap it.
    if !json.is_object() {
        json = serde_json::json!({ "_original": json });
    }

    let Some(obj) = json.as_object_mut() else {
        return Ok(wasm_args);
    };

    if let Some(newton) = obj.get_mut("_newton") {
        normalize_newton_aliases(newton)?;

        if let (Some(existing_value), Some(cid)) = (
            newton
                .as_object()
                .and_then(|newton_obj| newton_obj.get("proof_cid"))
                .cloned(),
            proof_cid,
        ) {
            let requested_value = serde_json::json!(cid);
            if existing_value != requested_value {
                return Err(WasmArgsError::ConflictingProofCid {
                    existing_value: existing_value.to_string(),
                    request_value: requested_value.to_string(),
                });
            }
        }
    }

    let Some(cid) = proof_cid else {
        // Validation-only path: normalize to detect alias conflicts, but return
        // original bytes since there is no proof_cid to inject.
        return Ok(wasm_args);
    };

    let newton = obj.entry("_newton").or_insert_with(|| serde_json::json!({}));

    if let Some(newton_obj) = newton.as_object_mut() {
        newton_obj.remove("proofCid");
        newton_obj.insert("proof_cid".to_string(), serde_json::json!(cid));
    } else {
        // _newton exists but isn't an object — overwrite.
        *newton = serde_json::json!({ "proof_cid": cid });
    }

    match serde_json::to_vec(&json) {
        Ok(bytes) => Ok(Some(alloy::primitives::Bytes::from(bytes))),
        Err(_) => Ok(wasm_args),
    }
}

#[cfg(test)]
mod tests {
    use super::{
        inject_proof_cid_into_wasm_args, parse_wasm_args, NewtonDirectives, WasmArgsError, MAX_INLINE_PRIVACY_ENVELOPES,
    };
    use alloy::primitives::Bytes;
    use serde_json::json;

    #[test]
    fn parse_wasm_args_empty_bytes_returns_defaults_and_original() {
        let input = b"";
        let (directives, passthrough) = parse_wasm_args(input).expect("empty bytes should parse");
        assert_eq!(directives, NewtonDirectives::default());
        assert_eq!(passthrough, input);
    }

    #[test]
    fn parse_wasm_args_without_newton_returns_defaults_and_original_json() {
        let input = br#"{"foo":"bar"}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("request should parse");
        assert_eq!(directives, NewtonDirectives::default());
        assert_eq!(passthrough, input);
    }

    #[test]
    fn parse_wasm_args_with_newton_extracts_directives_and_strips_namespace() {
        let input = br#"{"foo":"bar","_newton":{"proof_cid":"bafyproof"}}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("request should parse");
        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: None,
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(passthrough, br#"{"foo":"bar"}"#);
    }

    #[test]
    fn parse_wasm_args_non_json_returns_defaults_and_original() {
        let input = b"\xff\xfe\xfd";
        let (directives, passthrough) = parse_wasm_args(input).expect("non-json bytes should pass through");
        assert_eq!(directives, NewtonDirectives::default());
        assert_eq!(passthrough, input);
    }

    #[test]
    fn parse_wasm_args_extracts_all_supported_directives() {
        let input = br#"{"_newton":{"proof_cid":"bafyproof","proof_type":"tlsn"},"foo":"bar"}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("request should parse");

        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: Some("tlsn".to_string()),
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&passthrough).ok(),
            Some(json!({"foo": "bar"}))
        );
    }

    #[test]
    fn parse_wasm_args_accepts_camel_case_directive_aliases() {
        let input = br#"{"_newton":{"proofCid":"bafyproof","proofType":"tlsn"},"foo":"bar"}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("request should parse");

        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: Some("tlsn".to_string()),
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&passthrough).ok(),
            Some(json!({"foo": "bar"}))
        );
    }

    #[test]
    fn parse_wasm_args_rejects_conflicting_proof_cid_aliases() {
        let input = br#"{"_newton":{"proof_cid":"bafycanonical","proofCid":"bafycamel"}}"#;
        let err = parse_wasm_args(input).expect_err("conflicting aliases should be rejected");

        assert_eq!(
            err,
            WasmArgsError::ConflictingAlias {
                alias: "proofCid",
                canonical: "proof_cid",
                alias_value: "\"bafycamel\"".to_string(),
                canonical_value: "\"bafycanonical\"".to_string(),
            }
        );
    }

    #[test]
    fn inject_proof_cid_roundtrip_preserves_passthrough_fields() {
        let input = Some(Bytes::from(
            br#"{"foo":"bar","count":1,"_newton":{"proof_type":"tlsn"}}"#.to_vec(),
        ));

        let injected = inject_proof_cid_into_wasm_args(input, Some("bafyproof"))
            .expect("request should inject")
            .expect("expected injected wasm args");

        let (directives, passthrough) = parse_wasm_args(injected.as_ref()).expect("roundtrip should parse");

        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: Some("tlsn".to_string()),
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&passthrough).ok(),
            Some(json!({"foo": "bar", "count": 1}))
        );
    }

    #[test]
    fn inject_proof_cid_wraps_non_object_json_before_roundtrip() {
        let input = Some(Bytes::from(br#"[1,2,3]"#.to_vec()));

        let injected = inject_proof_cid_into_wasm_args(input, Some("bafyproof"))
            .expect("request should inject")
            .expect("expected injected wasm args");

        let injected_json = serde_json::from_slice::<serde_json::Value>(injected.as_ref()).ok();
        assert_eq!(
            injected_json,
            Some(json!({"_original": [1, 2, 3], "_newton": {"proof_cid": "bafyproof"}}))
        );

        let (directives, passthrough) = parse_wasm_args(injected.as_ref()).expect("roundtrip should parse");
        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: None,
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&passthrough).ok(),
            Some(json!({"_original": [1, 2, 3]}))
        );
    }

    #[test]
    fn parse_wasm_args_roundtrips_injected_proof_cid_and_preserves_passthrough() {
        let original = Bytes::from_static(br#"{"foo":"bar","count":1}"#);
        let injected = inject_proof_cid_into_wasm_args(Some(original.clone()), Some("bafyproof"))
            .expect("request should inject")
            .expect("injected wasm args should exist");

        let (directives, passthrough) = parse_wasm_args(injected.as_ref()).expect("roundtrip should parse");

        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: None,
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(
            serde_json::from_slice::<serde_json::Value>(&passthrough).ok(),
            serde_json::from_slice::<serde_json::Value>(original.as_ref()).ok()
        );
    }

    #[test]
    fn inject_proof_cid_roundtrip_overwrites_existing_newton_namespace() {
        let original = Bytes::from_static(br#"{"foo":"bar","_newton":"legacy"}"#);
        let injected = inject_proof_cid_into_wasm_args(Some(original), Some("bafyproof"))
            .expect("request should inject")
            .expect("injected wasm args should exist");

        let (directives, passthrough) = parse_wasm_args(injected.as_ref()).expect("roundtrip should parse");

        assert_eq!(directives.proof_cid.as_deref(), Some("bafyproof"));
        assert_eq!(passthrough, br#"{"foo":"bar"}"#);
    }

    #[test]
    fn inject_proof_cid_normalizes_existing_camel_case_alias() {
        let original = Bytes::from_static(br#"{"foo":"bar","_newton":{"proofCid":"bafyproof"}}"#);
        let injected = inject_proof_cid_into_wasm_args(Some(original), Some("bafyproof"))
            .expect("request should inject")
            .expect("injected wasm args should exist");

        let injected_json =
            serde_json::from_slice::<serde_json::Value>(injected.as_ref()).expect("injected wasm args should be json");
        assert_eq!(injected_json, json!({"foo":"bar","_newton":{"proof_cid":"bafyproof"}}));

        let (directives, passthrough) = parse_wasm_args(injected.as_ref()).expect("roundtrip should parse");
        assert_eq!(directives.proof_cid.as_deref(), Some("bafyproof"));
        assert_eq!(passthrough, br#"{"foo":"bar"}"#);
    }

    #[test]
    fn inject_proof_cid_rejects_conflict_with_existing_newton_proof_cid() {
        let original = Bytes::from_static(br#"{"foo":"bar","_newton":{"proof_cid":"legacy"}}"#);
        let err = inject_proof_cid_into_wasm_args(Some(original), Some("bafyproof"))
            .expect_err("conflicting proof_cid values should be rejected");

        assert_eq!(
            err,
            WasmArgsError::ConflictingProofCid {
                existing_value: "\"legacy\"".to_string(),
                request_value: "\"bafyproof\"".to_string(),
            }
        );
    }

    #[test]
    fn parse_wasm_args_flattens_unknown_newton_fields_into_extra() {
        let input = br#"{"_newton":{"proof_cid":"bafyproof","proof_type":"tlsn","network":"testnet","retries":2}}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("request should parse");

        assert_eq!(
            directives,
            NewtonDirectives {
                proof_cid: Some("bafyproof".to_string()),
                proof_type: Some("tlsn".to_string()),
                extra: serde_json::Map::from_iter([
                    ("network".to_string(), json!("testnet")),
                    ("retries".to_string(), json!(2)),
                ]),
                ..NewtonDirectives::default()
            }
        );
        assert_eq!(passthrough, br#"{}"#);
    }

    #[test]
    fn parse_wasm_args_extracts_privacy_envelopes() {
        let input = br#"{"_newton":{"privacy":["YmFzZTY0ZW52ZWxvcGUx","YmFzZTY0ZW52ZWxvcGUy"]}}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("should parse");
        assert_eq!(
            directives.privacy,
            Some(vec![
                "YmFzZTY0ZW52ZWxvcGUx".to_string(),
                "YmFzZTY0ZW52ZWxvcGUy".to_string()
            ])
        );
        assert_eq!(passthrough, br#"{}"#);
    }

    #[test]
    fn parse_wasm_args_rejects_too_many_privacy_envelopes() {
        let envelopes: Vec<String> = (0..11).map(|i| format!("envelope_{i}")).collect();
        let input = serde_json::json!({"_newton": {"privacy": envelopes}});
        let input_bytes = serde_json::to_vec(&input).unwrap();
        let err = parse_wasm_args(&input_bytes).expect_err("should reject >10 envelopes");
        assert!(matches!(
            err,
            WasmArgsError::TooManyPrivacyEnvelopes {
                count: 11,
                max: MAX_INLINE_PRIVACY_ENVELOPES
            }
        ));
    }

    #[test]
    fn parse_wasm_args_extracts_proof_cids_plural() {
        let input = br#"{"_newton":{"proof_cids":["bafyCID1","bafyCID2"]}}"#;
        let (directives, _) = parse_wasm_args(input).expect("should parse");
        assert_eq!(
            directives.proof_cids,
            Some(vec!["bafyCID1".to_string(), "bafyCID2".to_string()])
        );
    }

    #[test]
    fn parse_wasm_args_accepts_camel_case_proof_cids_alias() {
        let input = br#"{"_newton":{"proofCids":["bafyCID1","bafyCID2"]}}"#;
        let (directives, _) = parse_wasm_args(input).expect("should parse");
        assert_eq!(
            directives.proof_cids,
            Some(vec!["bafyCID1".to_string(), "bafyCID2".to_string()])
        );
    }

    #[test]
    fn all_proof_cids_merges_singular_and_plural() {
        let d = NewtonDirectives {
            proof_cid: Some("bafySINGLE".to_string()),
            proof_cids: Some(vec!["bafyA".to_string(), "bafySINGLE".to_string(), "bafyB".to_string()]),
            ..Default::default()
        };
        assert_eq!(d.all_proof_cids(), vec!["bafySINGLE", "bafyA", "bafyB"]);
    }

    #[test]
    fn all_proof_cids_returns_empty_when_none() {
        let d = NewtonDirectives::default();
        assert!(d.all_proof_cids().is_empty());
    }

    #[test]
    fn parse_wasm_args_extracts_privacy_and_proofs_together() {
        let input = br#"{"foo":"bar","_newton":{"proof_cid":"bafyPROOF","privacy":["ZW52ZWxvcGU="]}}"#;
        let (directives, passthrough) = parse_wasm_args(input).expect("should parse");
        assert_eq!(directives.proof_cid.as_deref(), Some("bafyPROOF"));
        assert_eq!(directives.privacy, Some(vec!["ZW52ZWxvcGU=".to_string()]));
        assert_eq!(passthrough, br#"{"foo":"bar"}"#);
    }
}