northroot-proof-engine 0.3.2

//! Canonicalization for v0.2 proof types
//!
//! This module provides deterministic JSON canonicalization using BTreeMap
//! for sorted keys. All canonicalization functions produce identical output
//! for identical inputs, ensuring hash determinism across platforms.

use chrono::{DateTime, Utc};
use std::collections::BTreeMap;
use thiserror::Error;

use crate::types::{
    IsoTime, MethodRef, MethodShapeManifest, SpanCommitment, SpanMeta, SpanShape, SpanStatus,
    Trace, TraceRoots,
};

/// Error returned when canonicalization fails
#[derive(Debug, Error, PartialEq, Eq)]
pub enum CanonicalError {
    /// Failed to serialize to canonical form
    #[error("serialization error: {0}")]
    Serialization(String),

    /// Invalid timestamp format
    #[error("invalid timestamp: {0}")]
    InvalidTimestamp(String),
}

/// Normalizes a timestamp to RFC3339 format with millisecond precision (UTC).
///
/// Converts any RFC3339 timestamp to UTC and normalizes to millisecond precision.
/// Format: `YYYY-MM-DDTHH:MM:SS.sssZ` (always Z suffix, no timezone offset).
///
/// # Examples
///
/// ```
/// use northroot_proof_engine::canonical::normalize_timestamp;
///
/// // UTC conversion
/// assert_eq!(
///     normalize_timestamp("2025-11-04T05:00:00-05:00").unwrap(),
///     "2025-11-04T10:00:00.000Z"
/// );
///
/// // Millisecond precision
/// assert_eq!(
///     normalize_timestamp("2025-11-04T10:00:00.123456Z").unwrap(),
///     "2025-11-04T10:00:00.123Z"
/// );
/// ```
pub fn normalize_timestamp(timestamp: &str) -> Result<IsoTime, CanonicalError> {
    // Try parsing as RFC3339 (handles both with and without timezone offset)
    match DateTime::parse_from_rfc3339(timestamp) {
        Ok(dt) => {
            // Convert to UTC
            let utc_dt = dt.with_timezone(&Utc);
            Ok(format_timestamp_millis(utc_dt))
        }
        Err(_) => {
            // If parsing fails, try parsing as UTC directly
            match timestamp.parse::<DateTime<Utc>>() {
                Ok(dt) => Ok(format_timestamp_millis(dt)),
                Err(_) => Err(CanonicalError::InvalidTimestamp(timestamp.to_string())),
            }
        }
    }
}

/// Formats a UTC DateTime to RFC3339 with millisecond precision.
fn format_timestamp_millis(dt: DateTime<Utc>) -> String {
    let nanos = dt.timestamp_subsec_nanos();
    let millis = nanos / 1_000_000; // Truncate to milliseconds
    format!("{}.{:03}Z", dt.format("%Y-%m-%dT%H:%M:%S"), millis)
}

/// Canonicalizes a span shape for hashing.
///
/// Creates a deterministic JSON representation of the span shape.
/// Fields are sorted lexicographically by key.
pub fn canonicalize_span_shape(shape: &SpanShape) -> Result<String, CanonicalError> {
    let mut map = BTreeMap::new();
    map.insert("tool_id", serde_json::Value::String(shape.tool_id.clone()));
    map.insert(
        "tool_version",
        serde_json::Value::String(shape.tool_version.clone()),
    );
    map.insert("param_schema", shape.param_schema.clone());
    if let Some(ref checker_ref) = shape.checker_ref {
        map.insert(
            "checker_ref",
            serde_json::Value::String(checker_ref.clone()),
        );
    }

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

/// Canonicalizes a method shape manifest for hashing.
///
/// Note: This excludes the publisher_sig field (signatures are not hashed).
pub fn canonicalize_method_shape(manifest: &MethodShapeManifest) -> Result<String, CanonicalError> {
    let mut map = BTreeMap::new();
    map.insert(
        "schema_version",
        serde_json::Value::String(manifest.schema_version.clone()),
    );
    map.insert(
        "method_id",
        serde_json::Value::String(manifest.method_id.clone()),
    );
    map.insert(
        "version",
        serde_json::Value::String(manifest.version.clone()),
    );
    map.insert(
        "method_shape_root",
        serde_json::Value::String(manifest.method_shape_root.clone()),
    );
    if let Some(ref dag_hash) = manifest.method_dag_hash {
        map.insert(
            "method_dag_hash",
            serde_json::Value::String(dag_hash.clone()),
        );
    }

    // Runtime pins
    let mut runtime = BTreeMap::new();
    runtime.insert(
        "image_digest",
        serde_json::Value::String(manifest.runtime.image_digest.clone()),
    );
    runtime.insert(
        "method_fingerprint",
        serde_json::Value::String(manifest.runtime.method_fingerprint.clone()),
    );
    map.insert(
        "runtime",
        serde_json::Value::Object(
            runtime
                .into_iter()
                .map(|(k, v)| (k.to_string(), v))
                .collect(),
        ),
    );

    if !manifest.checker_refs.is_empty() {
        map.insert(
            "checker_refs",
            serde_json::Value::Array(
                manifest
                    .checker_refs
                    .iter()
                    .map(|s| serde_json::Value::String(s.clone()))
                    .collect(),
            ),
        );
    }
    if let Some(ref policy_id) = manifest.default_policy_id {
        map.insert(
            "default_policy_id",
            serde_json::Value::String(policy_id.clone()),
        );
    }

    // Note: publisher_sig is excluded (signatures are not hashed)

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

/// Canonicalizes a span commitment for hashing.
///
/// This is the core commitment that gets hashed and signed.
/// All timestamps are normalized to UTC with millisecond precision.
pub fn canonicalize_span_commitment(commitment: &SpanCommitment) -> Result<String, CanonicalError> {
    let normalized_start = normalize_timestamp(&commitment.start_time)?;
    let normalized_end = normalize_timestamp(&commitment.end_time)?;

    let mut map = BTreeMap::new();
    map.insert(
        "schema_version",
        serde_json::Value::String(commitment.schema_version.clone()),
    );
    map.insert(
        "trace_id",
        serde_json::Value::String(commitment.trace_id.clone()),
    );
    map.insert(
        "span_id",
        serde_json::Value::String(commitment.span_id.clone()),
    );

    // Method ref
    let method_ref_json = canonicalize_method_ref(&commitment.method_ref)?;
    let method_ref_value: serde_json::Value =
        serde_json::from_str(&method_ref_json).map_err(|e| {
            CanonicalError::Serialization(format!("method_ref canonicalization: {}", e))
        })?;
    map.insert("method_ref", method_ref_value);

    map.insert(
        "data_shape_hash",
        serde_json::Value::String(commitment.data_shape_hash.clone()),
    );
    map.insert(
        "span_shape_hash",
        serde_json::Value::String(commitment.span_shape_hash.clone()),
    );

    // Inputs
    let inputs: Vec<serde_json::Value> = commitment
        .inputs
        .iter()
        .map(|input| {
            let mut input_map = BTreeMap::new();
            input_map.insert("name", serde_json::Value::String(input.name.clone()));
            input_map.insert(
                "content_hash",
                serde_json::Value::String(input.content_hash.clone()),
            );
            if let Some(ref uri) = input.uri {
                input_map.insert("uri", serde_json::Value::String(uri.clone()));
            }
            serde_json::Value::Object(
                input_map
                    .into_iter()
                    .map(|(k, v)| (k.to_string(), v))
                    .collect(),
            )
        })
        .collect();
    map.insert("inputs", serde_json::Value::Array(inputs));

    // Outputs
    let outputs: Vec<serde_json::Value> = commitment
        .outputs
        .iter()
        .map(|output| {
            let mut output_map = BTreeMap::new();
            output_map.insert("name", serde_json::Value::String(output.name.clone()));
            output_map.insert(
                "content_hash",
                serde_json::Value::String(output.content_hash.clone()),
            );
            if let Some(ref uri) = output.uri {
                output_map.insert("uri", serde_json::Value::String(uri.clone()));
            }
            serde_json::Value::Object(
                output_map
                    .into_iter()
                    .map(|(k, v)| (k.to_string(), v))
                    .collect(),
            )
        })
        .collect();
    map.insert("outputs", serde_json::Value::Array(outputs));

    map.insert("start_time", serde_json::Value::String(normalized_start));
    map.insert("end_time", serde_json::Value::String(normalized_end));
    map.insert(
        "status",
        serde_json::Value::String(match commitment.status {
            SpanStatus::OK => "OK".to_string(),
            SpanStatus::ERROR => "ERROR".to_string(),
        }),
    );

    if let Some(ref meta) = commitment.meta {
        let meta_json = canonicalize_span_meta(meta)?;
        let meta_value: serde_json::Value = serde_json::from_str(&meta_json).map_err(|e| {
            CanonicalError::Serialization(format!("span_meta canonicalization: {}", e))
        })?;
        map.insert("meta", meta_value);
    }

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

/// Canonicalizes a method reference.
fn canonicalize_method_ref(method_ref: &MethodRef) -> Result<String, CanonicalError> {
    let mut map = BTreeMap::new();
    map.insert(
        "method_id",
        serde_json::Value::String(method_ref.method_id.clone()),
    );
    map.insert(
        "version",
        serde_json::Value::String(method_ref.version.clone()),
    );
    map.insert(
        "method_shape_root",
        serde_json::Value::String(method_ref.method_shape_root.clone()),
    );
    map.insert(
        "image_digest",
        serde_json::Value::String(method_ref.image_digest.clone()),
    );
    map.insert(
        "method_fingerprint",
        serde_json::Value::String(method_ref.method_fingerprint.clone()),
    );

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

/// Canonicalizes span metadata.
fn canonicalize_span_meta(meta: &SpanMeta) -> Result<String, CanonicalError> {
    let mut map = BTreeMap::new();
    if let Some(lamport) = meta.lamport {
        map.insert("lamport", serde_json::Value::Number(lamport.into()));
    }
    if let Some(ref parents) = meta.parents {
        map.insert(
            "parents",
            serde_json::Value::Array(
                parents
                    .iter()
                    .map(|s| serde_json::Value::String(s.clone()))
                    .collect(),
            ),
        );
    }
    if let Some(order_index) = meta.order_index {
        map.insert("order_index", serde_json::Value::Number(order_index.into()));
    }

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

/// Canonicalizes a trace for hashing.
pub fn canonicalize_trace(trace: &Trace) -> Result<String, CanonicalError> {
    let mut map = BTreeMap::new();
    map.insert(
        "schema_version",
        serde_json::Value::String(trace.schema_version.clone()),
    );
    map.insert(
        "trace_id",
        serde_json::Value::String(trace.trace_id.clone()),
    );

    // Method ref
    let method_ref_json = canonicalize_method_ref(&trace.method_ref)?;
    let method_ref_value: serde_json::Value =
        serde_json::from_str(&method_ref_json).map_err(|e| {
            CanonicalError::Serialization(format!("method_ref canonicalization: {}", e))
        })?;
    map.insert("method_ref", method_ref_value);

    map.insert(
        "data_shape_hash",
        serde_json::Value::String(trace.data_shape_hash.clone()),
    );
    map.insert(
        "policy_id",
        serde_json::Value::String(trace.policy_id.clone()),
    );
    map.insert(
        "identity_route",
        serde_json::Value::String(trace.identity_route.clone()),
    );
    map.insert(
        "content_route",
        serde_json::Value::String(trace.content_route.clone()),
    );

    if let Some(ref exec_route) = trace.execution_route {
        let mut exec_map = BTreeMap::new();
        exec_map.insert(
            "provider",
            serde_json::Value::String(exec_route.provider.clone()),
        );
        if let Some(ref account) = exec_route.account {
            exec_map.insert("account", serde_json::Value::String(account.clone()));
        }
        exec_map.insert(
            "region",
            serde_json::Value::String(exec_route.region.clone()),
        );
        exec_map.insert(
            "runtime_digest",
            serde_json::Value::String(exec_route.runtime_digest.clone()),
        );
        map.insert(
            "execution_route",
            serde_json::Value::Object(
                exec_map
                    .into_iter()
                    .map(|(k, v)| (k.to_string(), v))
                    .collect(),
            ),
        );
    }

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

/// Canonicalizes trace roots for hashing.
///
/// Note: This excludes server_sig (signatures are not hashed).
pub fn canonicalize_trace_roots(roots: &TraceRoots) -> Result<String, CanonicalError> {
    let mut map = BTreeMap::new();
    map.insert(
        "trace_id",
        serde_json::Value::String(roots.trace_id.clone()),
    );
    map.insert(
        "span_commitments",
        serde_json::Value::Array(
            roots
                .span_commitments
                .iter()
                .map(|s| serde_json::Value::String(s.clone()))
                .collect(),
        ),
    );
    map.insert(
        "trace_set_root",
        serde_json::Value::String(roots.trace_set_root.clone()),
    );
    if let Some(ref identity_root) = roots.identity_root {
        map.insert(
            "identity_root",
            serde_json::Value::String(identity_root.clone()),
        );
    }
    if let Some(ref trace_seq_root) = roots.trace_seq_root {
        map.insert(
            "trace_seq_root",
            serde_json::Value::String(trace_seq_root.clone()),
        );
    }
    map.insert(
        "method_shape_root",
        serde_json::Value::String(roots.method_shape_root.clone()),
    );
    if let Some(ref tier01_status) = roots.tier01_status {
        map.insert(
            "tier01_status",
            serde_json::Value::String(tier01_status.clone()),
        );
    }

    // Note: server_sig is excluded (signatures are not hashed)

    serde_json::to_string(&serde_json::Value::Object(
        map.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
    ))
    .map_err(|e| CanonicalError::Serialization(e.to_string()))
}

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

    #[test]
    fn test_normalize_timestamp_utc_conversion() {
        let t1 = normalize_timestamp("2025-11-04T05:00:00-05:00").unwrap();
        let t2 = normalize_timestamp("2025-11-04T10:00:00Z").unwrap();
        assert_eq!(t1, t2);
        assert_eq!(t1, "2025-11-04T10:00:00.000Z");
    }

    #[test]
    fn test_normalize_timestamp_millisecond_precision() {
        let normalized = normalize_timestamp("2025-11-04T10:00:00.123456Z").unwrap();
        assert_eq!(normalized, "2025-11-04T10:00:00.123Z");
    }

    #[test]
    fn test_normalize_timestamp_sub_millisecond_truncation() {
        let normalized = normalize_timestamp("2025-11-04T10:00:00.000999Z").unwrap();
        assert_eq!(normalized, "2025-11-04T10:00:00.000Z");
    }

    #[test]
    fn test_canonicalize_span_shape_determinism() {
        let shape = SpanShape {
            tool_id: "io.read_csv".to_string(),
            tool_version: "1".to_string(),
            param_schema: serde_json::json!({"delimiter": {"type": "string", "default": ","}}),
            checker_ref: Some("chk:csv-valid".to_string()),
        };

        let c1 = canonicalize_span_shape(&shape).unwrap();
        let c2 = canonicalize_span_shape(&shape).unwrap();
        assert_eq!(c1, c2);
    }
}