xdoc-rs 0.1.1

use sha2::{Digest, Sha256};

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

use super::xmldsig::{
    element_children, find_signature, required_child, required_child_text, XMLDSIG_NAMESPACE_URI,
};
use super::{
    canonicalize_node, decode_standard_base64, digest_bytes, encode_standard_base64,
    CanonicalizationConfig, DigestAlgorithm, XADES_NAMESPACE_URI,
};

/// Request sent to a timestamp authority abstraction.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TimestampRequest {
    pub digest_algorithm: DigestAlgorithm,
    pub message_imprint: Vec<u8>,
}

/// Opaque timestamp token bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TimestampToken {
    pub encoded: Vec<u8>,
}

impl TimestampToken {
    pub fn new(encoded: impl Into<Vec<u8>>) -> Self {
        Self {
            encoded: encoded.into(),
        }
    }
}

/// Supplies timestamp tokens for XAdES unsigned properties.
///
/// Implementations may call a real TSA, HSM, KMS, or local service, but the
/// engine does not include any network client by default.
pub trait TimestampAuthorityClient {
    fn timestamp(&self, request: &TimestampRequest) -> XmlResult<TimestampToken>;

    fn verify(&self, request: &TimestampRequest, token: &TimestampToken) -> XmlResult<bool> {
        Ok(self.timestamp(request)?.encoded == token.encoded)
    }
}

/// Deterministic timestamp authority for tests and fixtures.
///
/// This is not an RFC 3161 implementation. It creates stable opaque bytes so
/// the XAdES-T structure and message-imprint flow can be tested without
/// enabling network access.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeterministicTimestampAuthority {
    secret: Vec<u8>,
}

impl DeterministicTimestampAuthority {
    pub fn new(secret: impl Into<Vec<u8>>) -> Self {
        Self {
            secret: secret.into(),
        }
    }
}

impl TimestampAuthorityClient for DeterministicTimestampAuthority {
    fn timestamp(&self, request: &TimestampRequest) -> XmlResult<TimestampToken> {
        let mut hasher = Sha256::new();
        hasher.update(b"xdoc-deterministic-timestamp");
        hasher.update([0]);
        hasher.update(request.digest_algorithm.uri().as_bytes());
        hasher.update([0]);
        hasher.update(&self.secret);
        hasher.update([0]);
        hasher.update(&request.message_imprint);
        Ok(TimestampToken::new(hasher.finalize().to_vec()))
    }
}

/// Configuration for adding and validating XAdES signature timestamps.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct XadesTimestampConfig {
    digest_algorithm: DigestAlgorithm,
    canonicalization: CanonicalizationConfig,
}

impl XadesTimestampConfig {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_digest_algorithm(mut self, digest_algorithm: DigestAlgorithm) -> Self {
        self.digest_algorithm = digest_algorithm;
        self
    }

    pub fn with_canonicalization(mut self, canonicalization: CanonicalizationConfig) -> Self {
        self.canonicalization = canonicalization;
        self
    }

    pub fn digest_algorithm(&self) -> DigestAlgorithm {
        self.digest_algorithm
    }

    pub fn canonicalization(&self) -> &CanonicalizationConfig {
        &self.canonicalization
    }
}

impl Default for XadesTimestampConfig {
    fn default() -> Self {
        Self {
            digest_algorithm: DigestAlgorithm::Sha256,
            canonicalization: CanonicalizationConfig::default(),
        }
    }
}

/// Result of validating a XAdES SignatureTimeStamp.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TimestampValidationReport {
    pub timestamp_present: bool,
    pub structure_valid: bool,
    pub token_valid: bool,
    pub message_imprint: Vec<u8>,
}

/// Adds `xades:SignatureTimeStamp` as an unsigned property to an existing XAdES signature.
pub fn add_signature_timestamp<C>(
    document: &Document,
    client: &C,
    config: &XadesTimestampConfig,
) -> XmlResult<Document>
where
    C: TimestampAuthorityClient + ?Sized,
{
    config.digest_algorithm.ensure_allowed_for_generation()?;

    let mut timestamped = document.clone();
    let signature = find_signature(&timestamped)?;
    let signature_value = required_child(&timestamped, signature, "SignatureValue")?;
    let message_imprint = signature_value_message_imprint(&timestamped, signature_value, config)?;
    let request = TimestampRequest {
        digest_algorithm: config.digest_algorithm,
        message_imprint,
    };
    let token = client.timestamp(&request)?;
    let qualifying_properties = find_qualifying_properties(&timestamped, signature)?;
    let unsigned_signature_properties =
        ensure_unsigned_signature_properties(&mut timestamped, qualifying_properties)?;

    if optional_xades_child(
        &timestamped,
        unsigned_signature_properties,
        "SignatureTimeStamp",
    )?
    .is_some()
    {
        return Err(XmlError::new(
            ErrorKind::Signature,
            "XAdES SignatureTimeStamp already exists",
        ));
    }

    let signature_timestamp = timestamped.add_element(
        unsigned_signature_properties,
        QName::qualified("xades", "SignatureTimeStamp", XADES_NAMESPACE_URI)?,
    )?;
    let canonicalization = timestamped.add_element(
        signature_timestamp,
        QName::qualified("ds", "CanonicalizationMethod", XMLDSIG_NAMESPACE_URI)?,
    )?;
    timestamped.add_attribute(
        canonicalization,
        Attribute::new(
            QName::new("Algorithm")?,
            config.canonicalization.algorithm().uri(),
        ),
    )?;
    let encoded = timestamped.add_element(
        signature_timestamp,
        QName::qualified("xades", "EncapsulatedTimeStamp", XADES_NAMESPACE_URI)?,
    )?;
    timestamped.add_text(encoded, encode_standard_base64(&token.encoded))?;

    Ok(timestamped)
}

/// Validates the presence, structure, and opaque token of a XAdES SignatureTimeStamp.
pub fn verify_signature_timestamp<C>(
    document: &Document,
    client: &C,
    config: &XadesTimestampConfig,
) -> XmlResult<TimestampValidationReport>
where
    C: TimestampAuthorityClient + ?Sized,
{
    config.digest_algorithm.ensure_allowed_for_generation()?;

    let signature = find_signature(document)?;
    let signature_value = required_child(document, signature, "SignatureValue")?;
    let message_imprint = signature_value_message_imprint(document, signature_value, config)?;
    let Some(signature_timestamp) = find_signature_timestamp(document, signature)? else {
        return Ok(TimestampValidationReport {
            timestamp_present: false,
            structure_valid: false,
            token_valid: false,
            message_imprint,
        });
    };

    let canonicalization_method =
        required_child(document, signature_timestamp, "CanonicalizationMethod")?;
    let canonicalization_algorithm =
        optional_attribute(document, canonicalization_method, "Algorithm")?.ok_or_else(|| {
            XmlError::new(
                ErrorKind::Signature,
                "XAdES SignatureTimeStamp CanonicalizationMethod requires Algorithm",
            )
        })?;
    let structure_valid = canonicalization_algorithm == config.canonicalization.algorithm().uri();
    let token_text = required_child_text(document, signature_timestamp, "EncapsulatedTimeStamp")?;
    let token = TimestampToken::new(decode_standard_base64(&token_text)?);
    let request = TimestampRequest {
        digest_algorithm: config.digest_algorithm,
        message_imprint: message_imprint.clone(),
    };
    let token_valid = structure_valid && client.verify(&request, &token)?;

    Ok(TimestampValidationReport {
        timestamp_present: true,
        structure_valid,
        token_valid,
        message_imprint,
    })
}

fn signature_value_message_imprint(
    document: &Document,
    signature_value: NodeId,
    config: &XadesTimestampConfig,
) -> XmlResult<Vec<u8>> {
    let canonicalized = canonicalize_node(document, signature_value, &config.canonicalization)?;
    digest_bytes(config.digest_algorithm, canonicalized)
}

fn find_signature_timestamp(document: &Document, signature: NodeId) -> XmlResult<Option<NodeId>> {
    let qualifying_properties = find_qualifying_properties(document, signature)?;
    let Some(unsigned_properties) =
        optional_xades_child(document, qualifying_properties, "UnsignedProperties")?
    else {
        return Ok(None);
    };
    let Some(unsigned_signature_properties) =
        optional_xades_child(document, unsigned_properties, "UnsignedSignatureProperties")?
    else {
        return Ok(None);
    };
    optional_xades_child(
        document,
        unsigned_signature_properties,
        "SignatureTimeStamp",
    )
}

fn ensure_unsigned_signature_properties(
    document: &mut Document,
    qualifying_properties: NodeId,
) -> XmlResult<NodeId> {
    let unsigned_properties =
        match optional_xades_child(document, qualifying_properties, "UnsignedProperties")? {
            Some(node) => node,
            None => document.add_element(
                qualifying_properties,
                QName::qualified("xades", "UnsignedProperties", XADES_NAMESPACE_URI)?,
            )?,
        };

    match optional_xades_child(document, unsigned_properties, "UnsignedSignatureProperties")? {
        Some(node) => Ok(node),
        None => document.add_element(
            unsigned_properties,
            QName::qualified("xades", "UnsignedSignatureProperties", XADES_NAMESPACE_URI)?,
        ),
    }
}

fn find_qualifying_properties(document: &Document, signature: NodeId) -> XmlResult<NodeId> {
    let object = required_child(document, signature, "Object")?;
    element_children(document, object)?
        .into_iter()
        .find(|child| is_xades_element(document, *child, "QualifyingProperties"))
        .ok_or_else(|| {
            XmlError::new(
                ErrorKind::Signature,
                "missing required XAdES QualifyingProperties",
            )
        })
}

fn optional_xades_child(
    document: &Document,
    parent: NodeId,
    local: &str,
) -> XmlResult<Option<NodeId>> {
    Ok(element_children(document, parent)?
        .into_iter()
        .find(|child| is_xades_element(document, *child, local)))
}

fn optional_attribute(document: &Document, node: NodeId, local: &str) -> XmlResult<Option<String>> {
    let NodeKind::Element(element) = document.node(node)?.kind() else {
        return Ok(None);
    };
    Ok(element
        .attributes()
        .iter()
        .find(|attribute| attribute.name().local() == local)
        .map(|attribute| attribute.value().to_owned()))
}

fn is_xades_element(document: &Document, node: NodeId, local: &str) -> bool {
    matches!(
        document.node(node).map(|node| node.kind()),
        Ok(NodeKind::Element(element))
            if element.name().namespace_uri().map(|uri| uri.as_str()) == Some(XADES_NAMESPACE_URI)
                && element.name().local() == local
    )
}

#[cfg(test)]
mod tests {
    use crate::parser::parse_str;
    use crate::signature::{
        sign_xades_bes_enveloped, verify_xades_bes_enveloped, DeterministicSigningProvider,
        XadesConfig,
    };
    use crate::writer::to_string_compact;

    use super::*;

    fn signing_provider() -> DeterministicSigningProvider {
        DeterministicSigningProvider::new(b"test-cert".to_vec(), b"test-secret".to_vec())
    }

    fn timestamp_client() -> DeterministicTimestampAuthority {
        DeterministicTimestampAuthority::new(b"timestamp-secret".to_vec())
    }

    fn unsigned_document() -> XmlResult<Document> {
        parse_str(r#"<Root Id="doc-1"><Item>value</Item></Root>"#)
    }

    fn signed_document() -> XmlResult<Document> {
        sign_xades_bes_enveloped(
            &unsigned_document()?,
            &signing_provider(),
            &XadesConfig::new().with_signing_time("2026-06-11T12:00:00Z"),
        )
    }

    #[test]
    fn xades_timestamp_adds_unsigned_signature_timestamp() -> XmlResult<()> {
        let timestamped = add_signature_timestamp(
            &signed_document()?,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )?;
        let xml = to_string_compact(&timestamped)?;
        let report = verify_signature_timestamp(
            &timestamped,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )?;

        assert!(report.timestamp_present);
        assert!(report.structure_valid);
        assert!(report.token_valid);
        assert!(xml.contains("<xades:UnsignedProperties>"));
        assert!(xml.contains("<xades:UnsignedSignatureProperties>"));
        assert!(xml.contains("<xades:SignatureTimeStamp>"));
        assert!(xml.contains("<xades:EncapsulatedTimeStamp>"));
        assert!(
            verify_xades_bes_enveloped(&timestamped, &signing_provider(), &XadesConfig::new())?
                .valid
        );
        Ok(())
    }

    #[test]
    fn xades_timestamp_report_marks_missing_timestamp() -> XmlResult<()> {
        let report = verify_signature_timestamp(
            &signed_document()?,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )?;

        assert!(!report.timestamp_present);
        assert!(!report.structure_valid);
        assert!(!report.token_valid);
        assert!(!report.message_imprint.is_empty());
        Ok(())
    }

    #[test]
    fn xades_timestamp_rejects_duplicate_timestamp() -> XmlResult<()> {
        let timestamped = add_signature_timestamp(
            &signed_document()?,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )?;
        let error = add_signature_timestamp(
            &timestamped,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )
        .expect_err("duplicate timestamp must fail");

        assert_eq!(error.kind(), &ErrorKind::Signature);
        assert!(error.message().contains("already exists"));
        Ok(())
    }

    #[test]
    fn xades_timestamp_detects_tampered_signature_value() -> XmlResult<()> {
        let timestamped = add_signature_timestamp(
            &signed_document()?,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )?;
        let xml = to_string_compact(&timestamped)?
            .replace("<ds:SignatureValue>", "<ds:SignatureValue>tampered");
        let tampered = parse_str(&xml)?;
        let report = verify_signature_timestamp(
            &tampered,
            &timestamp_client(),
            &XadesTimestampConfig::new(),
        )?;

        assert!(report.timestamp_present);
        assert!(report.structure_valid);
        assert!(!report.token_valid);
        Ok(())
    }
}