1use crate::{Diagnostic, Document, Limits, Node, parse};
2use sha2::{Digest, Sha256};
3use std::{collections::BTreeMap, fmt::Write, ops::Range};
4use xmlparser::{ElementEnd, Token, Tokenizer};
5
6pub fn budget_signature(limits: &Limits) -> Option<String> {
11 if limits.max_tokens.is_none() && limits.per_item.is_none() {
12 return None;
13 }
14 let mut hash = Sha256::new();
15 hash.update(b"htmlp-budget-v1\0");
16 for value in [limits.max_tokens, limits.per_item] {
17 hash.update([u8::from(value.is_some())]);
18 if let Some(value) = value {
19 hash.update(value.to_be_bytes());
20 }
21 }
22 let reason = limits.reason.as_deref().unwrap_or("");
23 hash.update((reason.len() as u64).to_be_bytes());
24 hash.update(reason.as_bytes());
25 let mut sig = String::with_capacity(16);
26 for byte in &hash.finalize()[..8] {
27 write!(sig, "{byte:02x}").expect("writing to String cannot fail");
28 }
29 Some(sig)
30}
31
32impl Document {
33 pub fn sign(&mut self) {
37 self.limits.sig = budget_signature(&self.limits);
38 let mut pending: Vec<_> = self.children.iter_mut().collect();
40 while let Some(node) = pending.pop() {
41 if let Node::Section(section) = node {
42 section.limits.sig = budget_signature(§ion.limits);
43 pending.extend(section.children.iter_mut());
44 }
45 }
46 }
47}
48
49pub fn sign_source(source: &str) -> Result<String, Diagnostic> {
54 let doc = parse(source)?;
55 let mut signatures = BTreeMap::new();
56 signatures.insert(doc.position.offset, budget_signature(&doc.limits));
57 let mut pending: Vec<_> = doc.children.iter().collect();
58 while let Some(node) = pending.pop() {
59 if let Node::Section(section) = node {
60 signatures.insert(section.position.offset, budget_signature(§ion.limits));
61 pending.extend(section.children.iter());
62 }
63 }
64 let mut current = None;
65 let mut insertion = 0;
66 let mut attribute: Option<(Range<usize>, Range<usize>)> = None;
67 let mut edits = Vec::new();
68 for token in Tokenizer::from(source) {
70 match token.expect("source already parsed") {
71 Token::ElementStart { span, .. } => {
72 current = signatures.get(&span.start());
73 attribute = None;
74 insertion = span.end();
75 }
76 Token::Attribute {
77 local, value, span, ..
78 } => {
79 insertion = span.end();
80 if local.as_str() == "sig" {
81 attribute = Some((value.range(), span.range()));
82 }
83 }
84 Token::ElementEnd {
85 end: ElementEnd::Open | ElementEnd::Empty,
86 ..
87 } => match (current.expect("element has AST entry"), attribute.take()) {
88 (Some(sig), Some((value, _))) => edits.push((value, sig.clone())),
89 (Some(sig), None) => edits.push((insertion..insertion, format!(" sig=\"{sig}\""))),
90 (None, Some((_, whole))) => edits.push((whole, String::new())),
91 (None, None) => {}
92 },
93 _ => {}
94 }
95 }
96 let mut output = source.to_owned();
97 for (range, replacement) in edits.into_iter().rev() {
98 output.replace_range(range, &replacement);
99 }
100 if output.len() > 4 * 1024 * 1024 {
101 return Err(Diagnostic::new(
102 "source-size",
103 "Signed source exceeds 4 MiB",
104 doc.position,
105 ));
106 }
107 Ok(output)
108}