xdoc-rs 0.1.1

use std::collections::BTreeMap;

use crate::core::{
    Document, ElementData, ErrorKind, NamespaceDeclaration, NodeId, NodeKind, QName, XmlError,
    XmlResult,
};

use super::CanonicalizationAlgorithm;

type NamespaceMap = BTreeMap<Option<String>, String>;

/// Configuration for XML canonicalization used by signatures.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CanonicalizationConfig {
    algorithm: CanonicalizationAlgorithm,
    with_comments: bool,
}

impl CanonicalizationConfig {
    pub fn new(algorithm: CanonicalizationAlgorithm) -> Self {
        Self {
            algorithm,
            with_comments: false,
        }
    }

    pub fn canonical_xml_11() -> Self {
        Self::new(CanonicalizationAlgorithm::CanonicalXml11)
    }

    pub fn canonical_xml_10() -> Self {
        Self::new(CanonicalizationAlgorithm::CanonicalXml10)
    }

    pub fn with_comments(mut self, with_comments: bool) -> Self {
        self.with_comments = with_comments;
        self
    }

    pub fn algorithm(&self) -> CanonicalizationAlgorithm {
        self.algorithm
    }

    pub fn comments_enabled(&self) -> bool {
        self.with_comments
    }
}

impl Default for CanonicalizationConfig {
    fn default() -> Self {
        Self::canonical_xml_11()
    }
}

/// Canonicalizes a whole XML document using the configured algorithm.
pub fn canonicalize_document(
    document: &Document,
    config: &CanonicalizationConfig,
) -> XmlResult<Vec<u8>> {
    let root = document.root().ok_or_else(|| {
        XmlError::new(
            ErrorKind::Signature,
            "cannot canonicalize a document without a root element",
        )
    })?;

    canonicalize_node(document, root, config)
}

/// Canonicalizes a node and its descendants.
pub fn canonicalize_node(
    document: &Document,
    node: NodeId,
    config: &CanonicalizationConfig,
) -> XmlResult<Vec<u8>> {
    validate_config(config)?;
    let mut output = String::new();
    let inherited_namespaces = inherited_namespace_context(document, node)?;
    write_node(
        document,
        node,
        config,
        &[],
        &NamespaceMap::new(),
        &inherited_namespaces,
        &mut output,
    )?;
    Ok(output.into_bytes())
}

pub(crate) fn canonicalize_node_excluding(
    document: &Document,
    node: NodeId,
    excluded: &[NodeId],
    config: &CanonicalizationConfig,
) -> XmlResult<Vec<u8>> {
    validate_config(config)?;
    let mut output = String::new();
    let inherited_namespaces = inherited_namespace_context(document, node)?;
    write_node(
        document,
        node,
        config,
        excluded,
        &NamespaceMap::new(),
        &inherited_namespaces,
        &mut output,
    )?;
    Ok(output.into_bytes())
}

fn validate_config(config: &CanonicalizationConfig) -> XmlResult<()> {
    match config.algorithm {
        CanonicalizationAlgorithm::CanonicalXml10 | CanonicalizationAlgorithm::CanonicalXml11 => {}
        CanonicalizationAlgorithm::ExclusiveXml10 => {
            return Err(XmlError::new(
                ErrorKind::Signature,
                format!(
                    "canonicalization algorithm `{}` is not implemented yet",
                    config.algorithm.uri()
                ),
            ));
        }
    }
    if config.with_comments {
        return Err(XmlError::new(
            ErrorKind::Signature,
            "canonicalization with comments is not implemented yet",
        ));
    }
    Ok(())
}

fn write_node(
    document: &Document,
    node: NodeId,
    config: &CanonicalizationConfig,
    excluded: &[NodeId],
    rendered_namespaces: &NamespaceMap,
    in_scope_namespaces: &NamespaceMap,
    output: &mut String,
) -> XmlResult<()> {
    if excluded.contains(&node) {
        return Ok(());
    }

    match document.node(node)?.kind() {
        NodeKind::Element(element) => write_element(
            document,
            element,
            config,
            excluded,
            rendered_namespaces,
            in_scope_namespaces,
            output,
        )?,
        NodeKind::Text(text) | NodeKind::CData(text) => write_canonical_text(text, output),
        NodeKind::Comment(_) => {}
        NodeKind::ProcessingInstruction { target, data } => {
            output.push_str("<?");
            output.push_str(target);
            if let Some(data) = data {
                output.push(' ');
                output.push_str(data);
            }
            output.push_str("?>");
        }
    }
    Ok(())
}

fn write_element(
    document: &Document,
    element: &ElementData,
    config: &CanonicalizationConfig,
    excluded: &[NodeId],
    rendered_namespaces: &NamespaceMap,
    in_scope_namespaces: &NamespaceMap,
    output: &mut String,
) -> XmlResult<()> {
    let current_namespaces =
        namespace_context_with_declarations(in_scope_namespaces, element.namespace_declarations());

    output.push('<');
    output.push_str(&element.name().lexical_name());
    write_namespace_declarations(&current_namespaces, rendered_namespaces, output);
    write_attributes(element.attributes(), output);
    output.push('>');

    for child in element.children() {
        write_node(
            document,
            *child,
            config,
            excluded,
            &current_namespaces,
            &current_namespaces,
            output,
        )?;
    }

    output.push_str("</");
    output.push_str(&element.name().lexical_name());
    output.push('>');
    Ok(())
}

fn inherited_namespace_context(document: &Document, node: NodeId) -> XmlResult<NamespaceMap> {
    let mut ancestors = Vec::new();
    let mut current = document.parent(node)?;
    while let Some(parent) = current {
        ancestors.push(parent);
        current = document.parent(parent)?;
    }
    ancestors.reverse();

    let mut namespaces = NamespaceMap::new();
    for ancestor in ancestors {
        if let NodeKind::Element(element) = document.node(ancestor)?.kind() {
            merge_namespace_declarations(&mut namespaces, element.namespace_declarations());
        }
    }
    Ok(namespaces)
}

fn namespace_context_with_declarations(
    in_scope: &NamespaceMap,
    declarations: &[NamespaceDeclaration],
) -> NamespaceMap {
    let mut namespaces = in_scope.clone();
    merge_namespace_declarations(&mut namespaces, declarations);
    namespaces
}

fn merge_namespace_declarations(
    namespaces: &mut NamespaceMap,
    declarations: &[NamespaceDeclaration],
) {
    for declaration in declarations {
        namespaces.insert(
            namespace_prefix_key(declaration),
            declaration.uri().as_str().to_owned(),
        );
    }
}

fn write_namespace_declarations(
    current: &NamespaceMap,
    rendered: &NamespaceMap,
    output: &mut String,
) {
    for (prefix, uri) in current {
        if rendered
            .get(prefix)
            .is_some_and(|rendered_uri| rendered_uri == uri)
        {
            continue;
        }

        output.push(' ');
        match prefix {
            Some(prefix) => {
                output.push_str("xmlns:");
                output.push_str(prefix);
            }
            None => output.push_str("xmlns"),
        }
        output.push_str("=\"");
        write_canonical_attribute_value(uri, output);
        output.push('"');
    }
}

fn write_attributes(attributes: &[crate::core::Attribute], output: &mut String) {
    let mut attributes = attributes.iter().collect::<Vec<_>>();
    attributes.sort_by(|left, right| {
        attribute_sort_key(left.name()).cmp(&attribute_sort_key(right.name()))
    });

    for attribute in attributes {
        output.push(' ');
        output.push_str(&attribute.name().lexical_name());
        output.push_str("=\"");
        write_canonical_attribute_value(attribute.value(), output);
        output.push('"');
    }
}

fn namespace_prefix_key(declaration: &NamespaceDeclaration) -> Option<String> {
    declaration
        .prefix()
        .map(|prefix| prefix.as_str().to_owned())
}

fn attribute_sort_key(name: &QName) -> (String, String) {
    (
        name.namespace_uri()
            .map(|uri| uri.as_str().to_owned())
            .unwrap_or_default(),
        name.local().to_owned(),
    )
}

fn write_canonical_text(value: &str, output: &mut String) {
    for ch in value.chars() {
        match ch {
            '&' => output.push_str("&amp;"),
            '<' => output.push_str("&lt;"),
            '>' => output.push_str("&gt;"),
            '\r' => output.push_str("&#xD;"),
            _ => output.push(ch),
        }
    }
}

fn write_canonical_attribute_value(value: &str, output: &mut String) {
    for ch in value.chars() {
        match ch {
            '&' => output.push_str("&amp;"),
            '<' => output.push_str("&lt;"),
            '"' => output.push_str("&quot;"),
            '\t' => output.push_str("&#x9;"),
            '\n' => output.push_str("&#xA;"),
            '\r' => output.push_str("&#xD;"),
            _ => output.push(ch),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::core::{Attribute, NamespaceDeclaration, XML_NAMESPACE_URI};
    use crate::parser::parse_str;

    use super::*;

    #[test]
    fn c14n_canonicalizes_document_without_xml_declaration_or_comments() -> XmlResult<()> {
        let document = parse_str(
            r#"<?xml version="1.0"?>
<doc:Root xmlns:z="urn:z" xmlns:doc="urn:doc" z:flag="yes" Id="r1">
  <!--ignored-->
  <Empty/>
  <Text a="1 &amp; 2">A &amp; B</Text>
</doc:Root>"#,
        )?;

        let actual = canonicalize_document(&document, &CanonicalizationConfig::default())?;
        let expected = include_str!("../../tests/golden/signature/c14n_document.xml")
            .strip_suffix('\n')
            .expect("golden fixture ends with newline")
            .as_bytes();

        assert_eq!(actual, expected);
        Ok(())
    }

    #[test]
    fn c14n10_canonicalizes_document_with_supported_inclusive_surface() -> XmlResult<()> {
        let document = parse_str(
            r#"<?xml version="1.0"?>
<doc:Root xmlns:z="urn:z" xmlns:doc="urn:doc" z:flag="yes" Id="r1"><Empty/><Text a="1 &amp; 2">A &amp; B</Text></doc:Root>"#,
        )?;

        let actual = canonicalize_document(&document, &CanonicalizationConfig::canonical_xml_10())?;
        let expected = include_str!("../../tests/golden/signature/c14n10_document.xml")
            .strip_suffix('\n')
            .expect("golden fixture ends with newline")
            .as_bytes();

        assert_eq!(actual, expected);
        Ok(())
    }

    #[test]
    fn c14n_canonicalizes_node_subset() -> XmlResult<()> {
        let document = parse_str(r#"<Root><Item Id="A">one</Item><Item Id="B">two</Item></Root>"#)?;
        let item = crate::signature::find_element_by_id(
            &document,
            "B",
            &crate::signature::IdAttributePolicy::Standard,
        )?;

        let actual = canonicalize_node(&document, item, &CanonicalizationConfig::default())?;

        assert_eq!(actual, b"<Item Id=\"B\">two</Item>");
        Ok(())
    }

    #[test]
    fn c14n_orders_namespaces_and_attributes() -> XmlResult<()> {
        let mut document = Document::new();
        let root = document.add_root_element(QName::qualified("b", "Root", "urn:b")?)?;
        document.add_namespace_declaration(root, NamespaceDeclaration::prefixed("z", "urn:z")?)?;
        document.add_namespace_declaration(root, NamespaceDeclaration::default("urn:default")?)?;
        document.add_namespace_declaration(root, NamespaceDeclaration::prefixed("b", "urn:b")?)?;
        document.add_attribute(root, Attribute::new(QName::new("plain")?, "1"))?;
        document.add_attribute(
            root,
            Attribute::new(QName::qualified("z", "attr", "urn:z")?, "2"),
        )?;
        document.add_attribute(
            root,
            Attribute::new(QName::qualified("xml", "lang", XML_NAMESPACE_URI)?, "en"),
        )?;

        let actual = canonicalize_document(&document, &CanonicalizationConfig::default())?;

        assert_eq!(
            actual,
            br#"<b:Root xmlns="urn:default" xmlns:b="urn:b" xmlns:z="urn:z" plain="1" xml:lang="en" z:attr="2"></b:Root>"#
        );
        Ok(())
    }

    #[test]
    fn c14n_omits_redundant_child_namespace_declarations() -> XmlResult<()> {
        let document = parse_str(r#"<Root xmlns:p="urn:p"><p:Child xmlns:p="urn:p"/></Root>"#)?;

        let actual = canonicalize_document(&document, &CanonicalizationConfig::default())?;

        assert_eq!(
            actual,
            br#"<Root xmlns:p="urn:p"><p:Child></p:Child></Root>"#
        );
        Ok(())
    }

    #[test]
    fn c14n_node_subset_includes_in_scope_ancestor_namespaces() -> XmlResult<()> {
        let document = parse_str(
            r#"<Root xmlns="urn:root" xmlns:a="urn:a"><Container xmlns:ds="http://www.w3.org/2000/09/xmldsig#"><ds:KeyInfo Id="key-1"><ds:X509Data/></ds:KeyInfo></Container></Root>"#,
        )?;
        let key_info = crate::signature::find_element_by_id(
            &document,
            "key-1",
            &crate::signature::IdAttributePolicy::Standard,
        )?;

        let actual = canonicalize_node(&document, key_info, &CanonicalizationConfig::default())?;

        assert_eq!(
            actual,
            br#"<ds:KeyInfo xmlns="urn:root" xmlns:a="urn:a" xmlns:ds="http://www.w3.org/2000/09/xmldsig#" Id="key-1"><ds:X509Data></ds:X509Data></ds:KeyInfo>"#
        );
        Ok(())
    }

    #[test]
    fn c14n_escapes_text_and_attribute_values() -> XmlResult<()> {
        let mut document = Document::new();
        let root = document.add_root_element(QName::new("Root")?)?;
        document.add_attribute(root, Attribute::new(QName::new("value")?, "\"\t\n\r<&"))?;
        document.add_text(root, "A&B < C > D\r")?;

        let actual = canonicalize_document(&document, &CanonicalizationConfig::default())?;

        assert_eq!(
            actual,
            b"<Root value=\"&quot;&#x9;&#xA;&#xD;&lt;&amp;\">A&amp;B &lt; C &gt; D&#xD;</Root>"
        );
        Ok(())
    }

    #[test]
    fn c14n10_and_c14n11_are_selected_by_explicit_config() -> XmlResult<()> {
        let document = parse_str(r#"<Root xml:id="r1"><Item value="1"/></Root>"#)?;
        let c14n10 = canonicalize_document(&document, &CanonicalizationConfig::canonical_xml_10())?;
        let c14n11 = canonicalize_document(&document, &CanonicalizationConfig::canonical_xml_11())?;

        assert_eq!(
            c14n10,
            br#"<Root xml:id="r1"><Item value="1"></Item></Root>"#
        );
        assert_eq!(c14n11, c14n10);
        Ok(())
    }

    #[test]
    fn c14n_rejects_unimplemented_modes() {
        let document = parse_str("<Root/>").expect("valid document");
        let config = CanonicalizationConfig::new(CanonicalizationAlgorithm::ExclusiveXml10);

        let error =
            canonicalize_document(&document, &config).expect_err("exclusive c14n is future work");

        assert_eq!(error.kind(), &ErrorKind::Signature);
    }
}