1use std::io::Write;
7
8use quick_xml::events::{BytesText, Event};
9use quick_xml::name::{Namespace, ResolveResult};
10use quick_xml::reader::NsReader;
11use quick_xml::{Reader, Writer};
12
13use super::parse::XMLDSIG_NS;
14
15#[derive(Debug, thiserror::Error)]
17pub enum XmlMutationError {
18 #[error("XML parsing error: {0}")]
20 XmlParse(#[from] roxmltree::Error),
21 #[error("XML read error: {0}")]
23 Read(#[from] quick_xml::Error),
24 #[error("XML write error: {0}")]
26 Write(#[from] std::io::Error),
27 #[error("XML writer emitted invalid UTF-8: {0}")]
29 InvalidUtf8(#[from] std::string::FromUtf8Error),
30 #[error("signature template root must be one XMLDSig Signature element")]
32 InvalidSignatureTemplate,
33 #[error("expected {expected} XMLDSig {element} values, got {actual}")]
35 ValueCountMismatch {
36 element: &'static str,
38 expected: usize,
40 actual: usize,
42 },
43 #[error("source XML must contain a root element")]
45 MissingRootElement,
46}
47
48pub fn append_signature_to_root(
51 xml: &str,
52 signature_template: &str,
53) -> Result<String, XmlMutationError> {
54 validate_signature_template(signature_template)?;
55 let source = roxmltree::Document::parse(xml)?;
56 if !source.root().children().any(|node| node.is_element()) {
57 return Err(XmlMutationError::MissingRootElement);
58 }
59
60 let mut reader = Reader::from_str(xml);
61 let mut writer = Writer::new(Vec::new());
62 let mut root_depth = 0usize;
63 let mut saw_root = false;
64 let mut buf = Vec::new();
65
66 loop {
67 match reader.read_event_into(&mut buf)? {
68 Event::Start(element) if root_depth == 0 => {
69 saw_root = true;
70 root_depth = 1;
71 writer.write_event(Event::Start(element))?;
72 }
73 Event::Start(element) => {
74 root_depth += 1;
75 writer.write_event(Event::Start(element))?;
76 }
77 Event::Empty(element) if root_depth == 0 => {
78 saw_root = true;
79 writer.write_event(Event::Start(element.borrow()))?;
80 writer.get_mut().write_all(signature_template.as_bytes())?;
81 writer.write_event(Event::End(element.to_end()))?;
82 }
83 Event::End(element) if root_depth == 1 => {
84 writer.get_mut().write_all(signature_template.as_bytes())?;
85 writer.write_event(Event::End(element))?;
86 root_depth = 0;
87 }
88 Event::End(element) => {
89 root_depth = root_depth.saturating_sub(1);
90 writer.write_event(Event::End(element))?;
91 }
92 Event::Eof => break,
93 event => writer.write_event(event)?,
94 }
95 buf.clear();
96 }
97
98 if !saw_root {
99 return Err(XmlMutationError::MissingRootElement);
100 }
101
102 let output = String::from_utf8(writer.into_inner())?;
103 roxmltree::Document::parse(&output)?;
104 Ok(output)
105}
106
107pub fn fill_digest_values<I, S>(xml: &str, values: I) -> Result<String, XmlMutationError>
109where
110 I: IntoIterator<Item = S>,
111 S: AsRef<str>,
112{
113 fill_dsig_values(xml, "DigestValue", values)
114}
115
116pub fn fill_signed_info_digest_values<I, S>(
118 xml: &str,
119 values: I,
120) -> Result<String, XmlMutationError>
121where
122 I: IntoIterator<Item = S>,
123 S: AsRef<str>,
124{
125 let values: Vec<String> = values
126 .into_iter()
127 .map(|value| value.as_ref().to_owned())
128 .collect();
129 let expected = count_signed_info_digest_values(xml)?;
130 if expected != values.len() {
131 return Err(XmlMutationError::ValueCountMismatch {
132 element: "DigestValue",
133 expected,
134 actual: values.len(),
135 });
136 }
137
138 fill_dsig_values_matching(xml, "DigestValue", values, is_signed_info_reference_context)
139}
140
141pub fn fill_signature_values<I, S>(xml: &str, values: I) -> Result<String, XmlMutationError>
143where
144 I: IntoIterator<Item = S>,
145 S: AsRef<str>,
146{
147 fill_dsig_values(xml, "SignatureValue", values)
148}
149
150pub fn fill_signature_value(xml: &str, value: &str) -> Result<String, XmlMutationError> {
152 let expected = count_direct_signature_values(xml)?;
153 if expected != 1 {
154 return Err(XmlMutationError::ValueCountMismatch {
155 element: "SignatureValue",
156 expected,
157 actual: 1,
158 });
159 }
160
161 fill_dsig_values_matching(
162 xml,
163 "SignatureValue",
164 vec![value.to_owned()],
165 is_direct_signature_context,
166 )
167}
168
169fn fill_dsig_values<I, S>(
170 xml: &str,
171 local_name: &'static str,
172 values: I,
173) -> Result<String, XmlMutationError>
174where
175 I: IntoIterator<Item = S>,
176 S: AsRef<str>,
177{
178 let values: Vec<String> = values
179 .into_iter()
180 .map(|value| value.as_ref().to_owned())
181 .collect();
182 let expected = count_dsig_elements(xml, local_name)?;
183 if expected != values.len() {
184 return Err(XmlMutationError::ValueCountMismatch {
185 element: local_name,
186 expected,
187 actual: values.len(),
188 });
189 }
190
191 fill_dsig_values_matching(xml, local_name, values, |_, _| true)
192}
193
194fn fill_dsig_values_matching(
195 xml: &str,
196 local_name: &'static str,
197 values: Vec<String>,
198 mut should_replace: impl FnMut(&[(bool, Vec<u8>)], &ResolveResult<'_>) -> bool,
199) -> Result<String, XmlMutationError> {
200 let mut reader = NsReader::from_str(xml);
201 let mut writer = Writer::new(Vec::new());
202 let mut buf = Vec::new();
203 let mut value_index = 0usize;
204 let mut replacing_depth: Option<usize> = None;
205 let mut element_stack: Vec<(bool, Vec<u8>)> = Vec::new();
206
207 loop {
208 let (namespace, event) = reader.read_resolved_event_into(&mut buf)?;
209 if let Some(depth) = replacing_depth.as_mut() {
210 match event {
211 Event::Start(_) => *depth += 1,
212 Event::End(end) if *depth == 0 => {
213 writer.write_event(Event::End(end))?;
214 replacing_depth = None;
215 element_stack.pop();
216 }
217 Event::End(_) => *depth -= 1,
218 Event::Eof => break,
219 _ => {}
220 }
221 buf.clear();
222 continue;
223 }
224
225 match event {
226 Event::Start(element)
227 if is_dsig_element(&namespace, element.local_name().as_ref(), local_name)
228 && should_replace(&element_stack, &namespace) =>
229 {
230 element_stack.push((
231 is_dsig_namespace(&namespace),
232 element.local_name().as_ref().to_vec(),
233 ));
234 writer.write_event(Event::Start(element))?;
235 writer.write_event(Event::Text(BytesText::new(&values[value_index])))?;
236 value_index += 1;
237 replacing_depth = Some(0);
238 }
239 Event::Empty(element)
240 if is_dsig_element(&namespace, element.local_name().as_ref(), local_name)
241 && should_replace(&element_stack, &namespace) =>
242 {
243 writer.write_event(Event::Start(element.borrow()))?;
244 writer.write_event(Event::Text(BytesText::new(&values[value_index])))?;
245 value_index += 1;
246 writer.write_event(Event::End(element.to_end()))?;
247 }
248 Event::Start(element) => {
249 element_stack.push((
250 is_dsig_namespace(&namespace),
251 element.local_name().as_ref().to_vec(),
252 ));
253 writer.write_event(Event::Start(element))?;
254 }
255 Event::Empty(element) => writer.write_event(Event::Empty(element))?,
256 Event::End(element) => {
257 element_stack.pop();
258 writer.write_event(Event::End(element))?;
259 }
260 Event::Eof => break,
261 event => writer.write_event(event)?,
262 }
263 buf.clear();
264 }
265
266 if value_index != values.len() {
267 return Err(XmlMutationError::ValueCountMismatch {
268 element: local_name,
269 expected: values.len(),
270 actual: value_index,
271 });
272 }
273
274 let output = String::from_utf8(writer.into_inner())?;
275 roxmltree::Document::parse(&output)?;
276 Ok(output)
277}
278
279fn validate_signature_template(signature_template: &str) -> Result<(), XmlMutationError> {
280 let document = roxmltree::Document::parse(signature_template)?;
281 let root = document.root_element();
282 if root.tag_name().namespace() == Some(XMLDSIG_NS) && root.tag_name().name() == "Signature" {
283 Ok(())
284 } else {
285 Err(XmlMutationError::InvalidSignatureTemplate)
286 }
287}
288
289fn count_dsig_elements(xml: &str, local_name: &str) -> Result<usize, XmlMutationError> {
290 let document = roxmltree::Document::parse(xml)?;
291 Ok(document
292 .descendants()
293 .filter(|node| {
294 node.is_element()
295 && node.tag_name().namespace() == Some(XMLDSIG_NS)
296 && node.tag_name().name() == local_name
297 })
298 .count())
299}
300
301fn count_signed_info_digest_values(xml: &str) -> Result<usize, XmlMutationError> {
302 let document = roxmltree::Document::parse(xml)?;
303 Ok(document
304 .descendants()
305 .filter(|node| is_direct_signed_info_reference_digest(*node))
306 .count())
307}
308
309fn count_direct_signature_values(xml: &str) -> Result<usize, XmlMutationError> {
310 let document = roxmltree::Document::parse(xml)?;
311 Ok(document
312 .descendants()
313 .filter(|node| {
314 node.is_element()
315 && node.tag_name().namespace() == Some(XMLDSIG_NS)
316 && node.tag_name().name() == "SignatureValue"
317 && node
318 .parent()
319 .is_some_and(|parent| is_dsig_node(parent, "Signature"))
320 })
321 .count())
322}
323
324fn is_direct_signed_info_reference_digest(node: roxmltree::Node<'_, '_>) -> bool {
325 node.is_element()
326 && node.tag_name().namespace() == Some(XMLDSIG_NS)
327 && node.tag_name().name() == "DigestValue"
328 && node
329 .parent()
330 .is_some_and(|parent| is_dsig_node(parent, "Reference"))
331 && node
332 .parent()
333 .and_then(|parent| parent.parent())
334 .is_some_and(|grandparent| is_dsig_node(grandparent, "SignedInfo"))
335}
336
337fn is_dsig_node(node: roxmltree::Node<'_, '_>, expected_local: &str) -> bool {
338 node.is_element()
339 && node.tag_name().namespace() == Some(XMLDSIG_NS)
340 && node.tag_name().name() == expected_local
341}
342
343fn is_signed_info_reference_context(
344 element_stack: &[(bool, Vec<u8>)],
345 namespace: &ResolveResult<'_>,
346) -> bool {
347 is_dsig_namespace(namespace)
348 && matches!(
349 element_stack,
350 [.., (true, signed_info), (true, reference)]
351 if signed_info.as_slice() == b"SignedInfo"
352 && reference.as_slice() == b"Reference"
353 )
354}
355
356fn is_direct_signature_context(
357 element_stack: &[(bool, Vec<u8>)],
358 namespace: &ResolveResult<'_>,
359) -> bool {
360 is_dsig_namespace(namespace)
361 && matches!(
362 element_stack,
363 [.., (true, signature)] if signature.as_slice() == b"Signature"
364 )
365}
366
367fn is_dsig_element(namespace: &ResolveResult<'_>, local: &[u8], expected_local: &str) -> bool {
368 is_dsig_namespace(namespace) && local == expected_local.as_bytes()
369}
370
371fn is_dsig_namespace(namespace: &ResolveResult<'_>) -> bool {
372 matches!(namespace, ResolveResult::Bound(Namespace(ns)) if *ns == XMLDSIG_NS.as_bytes())
373}
374
375#[cfg(test)]
376mod tests {
377 use crate::c14n::{C14nAlgorithm, C14nMode};
378 use crate::xmldsig::{
379 DigestAlgorithm, ReferenceBuilder, SignatureAlgorithm, SignatureBuilder, Transform,
380 };
381
382 use super::*;
383
384 fn template(reference_count: usize) -> String {
385 let mut builder = SignatureBuilder::new(
386 C14nAlgorithm::new(C14nMode::Exclusive1_0, false),
387 SignatureAlgorithm::RsaSha256,
388 )
389 .ns_prefix("ds");
390 for index in 0..reference_count {
391 builder = builder.add_reference(
392 ReferenceBuilder::new(DigestAlgorithm::Sha256)
393 .uri(format!("#ref-{index}"))
394 .transform(Transform::Enveloped),
395 );
396 }
397 builder.build_template().expect("valid template")
398 }
399
400 #[test]
401 fn appends_signature_template_to_non_empty_root() {
402 let signed = append_signature_to_root("<root><payload ID=\"ref-0\"/></root>", &template(1))
403 .expect("append signature");
404 let document = roxmltree::Document::parse(&signed).expect("parse output");
405 let root = document.root_element();
406 let children: Vec<_> = root
407 .children()
408 .filter(roxmltree::Node::is_element)
409 .map(|node| node.tag_name().name())
410 .collect();
411 assert_eq!(children, ["payload", "Signature"]);
412 assert_eq!(
413 root.last_element_child()
414 .expect("signature")
415 .tag_name()
416 .namespace(),
417 Some(XMLDSIG_NS)
418 );
419 }
420
421 #[test]
422 fn appends_signature_template_to_empty_root() {
423 let signed = append_signature_to_root("<root/>", &template(1)).expect("append signature");
424 let document = roxmltree::Document::parse(&signed).expect("parse output");
425 let root = document.root_element();
426 assert_eq!(
427 root.first_element_child()
428 .expect("signature")
429 .tag_name()
430 .name(),
431 "Signature"
432 );
433 }
434
435 #[test]
436 fn rejects_non_signature_template() {
437 let err = append_signature_to_root("<root/>", "<NotSignature/>")
438 .expect_err("template must be a Signature");
439 assert!(matches!(err, XmlMutationError::InvalidSignatureTemplate));
440 }
441
442 #[test]
443 fn fills_digest_values_in_xml_dsig_document_order() {
444 let signed = append_signature_to_root("<root/>", &template(2)).expect("append signature");
445 let filled =
446 fill_digest_values(&signed, ["digest-one", "digest-two"]).expect("fill digest values");
447 let document = roxmltree::Document::parse(&filled).expect("parse output");
448 let values: Vec<_> = document
449 .descendants()
450 .filter(|node| node.has_tag_name((XMLDSIG_NS, "DigestValue")))
451 .map(|node| node.text())
452 .collect();
453 assert_eq!(values, [Some("digest-one"), Some("digest-two")]);
454 }
455
456 #[test]
457 fn fills_signature_value_without_touching_digest_values() {
458 let signed = append_signature_to_root("<root/>", &template(1)).expect("append signature");
459 let filled =
460 fill_signature_values(&signed, ["signature&bytes"]).expect("fill signature value");
461 let document = roxmltree::Document::parse(&filled).expect("parse output");
462 let signature_value = document
463 .descendants()
464 .find(|node| node.has_tag_name((XMLDSIG_NS, "SignatureValue")))
465 .expect("SignatureValue");
466 assert_eq!(signature_value.text(), Some("signature&bytes"));
467 let digest_value = document
468 .descendants()
469 .find(|node| node.has_tag_name((XMLDSIG_NS, "DigestValue")))
470 .expect("DigestValue");
471 assert_eq!(digest_value.text(), None);
472 }
473
474 #[test]
475 fn replacement_count_must_match_dsig_elements() {
476 let signed = append_signature_to_root("<root/>", &template(2)).expect("append signature");
477 let err = fill_digest_values(&signed, ["only-one"]).expect_err("mismatch");
478 assert!(matches!(
479 err,
480 XmlMutationError::ValueCountMismatch {
481 element: "DigestValue",
482 expected: 2,
483 actual: 1
484 }
485 ));
486 }
487
488 #[test]
489 fn does_not_replace_foreign_same_local_name_elements() {
490 let source = r#"<root xmlns:foreign="urn:test"><foreign:DigestValue>keep</foreign:DigestValue></root>"#;
491 let signed = append_signature_to_root(source, &template(1)).expect("append signature");
492 let filled = fill_digest_values(&signed, ["digest"]).expect("fill digest");
493 let document = roxmltree::Document::parse(&filled).expect("parse output");
494 let foreign = document
495 .descendants()
496 .find(|node| node.has_tag_name(("urn:test", "DigestValue")))
497 .expect("foreign DigestValue");
498 assert_eq!(foreign.text(), Some("keep"));
499 let dsig = document
500 .descendants()
501 .find(|node| node.has_tag_name((XMLDSIG_NS, "DigestValue")))
502 .expect("dsig DigestValue");
503 assert_eq!(dsig.text(), Some("digest"));
504 }
505
506 #[test]
507 fn replacement_preserves_target_end_after_self_closing_child() {
508 let source = r#"<ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#"><ds:SignedInfo><ds:Reference><ds:DigestValue><marker/></ds:DigestValue></ds:Reference></ds:SignedInfo></ds:Signature>"#;
509 let filled = fill_digest_values(source, ["digest"]).expect("fill digest");
510 let document = roxmltree::Document::parse(&filled).expect("parse output");
511 let digest_value = document
512 .descendants()
513 .find(|node| node.has_tag_name((XMLDSIG_NS, "DigestValue")))
514 .expect("DigestValue");
515 assert_eq!(digest_value.text(), Some("digest"));
516 assert_eq!(
517 digest_value
518 .next_sibling_element()
519 .map(|node| node.tag_name().name()),
520 None
521 );
522 }
523
524 #[test]
525 fn replacement_fails_when_nested_dsig_values_are_skipped() {
526 let source = r#"<ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#"><ds:SignedInfo><ds:Reference><ds:DigestValue><ds:DigestValue>nested</ds:DigestValue></ds:DigestValue></ds:Reference></ds:SignedInfo></ds:Signature>"#;
527 let err =
528 fill_digest_values(source, ["outer", "nested"]).expect_err("nested target skipped");
529 assert!(matches!(
530 err,
531 XmlMutationError::ValueCountMismatch {
532 element: "DigestValue",
533 expected: 2,
534 actual: 1
535 }
536 ));
537 }
538}