Skip to main content

htmlp/
signing.rs

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
6/// First eight bytes of SHA-256 as lowercase hex. Covers normalized local
7/// limits and the decoded reason, not content, identity, or approval.
8/// Returns `None` when neither limit is declared. See the specification for
9/// the canonical byte encoding, allowing implementations in other languages.
10pub 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    /// Explicitly accept the current budgets, updating every local signature.
34    /// This does not validate the document or its budgets; call [`crate::lint`]
35    /// afterwards. Variables never have signatures.
36    pub fn sign(&mut self) {
37        self.limits.sig = budget_signature(&self.limits);
38        // Iterative traversal also supports trees built outside the parser.
39        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(&section.limits);
43                pending.extend(section.children.iter_mut());
44            }
45        }
46    }
47}
48
49/// Add or replace budget signatures without reformatting source. Existing
50/// quotes, Markdown, comments, line endings, and attribute order are preserved.
51/// Removes signatures from elements that no longer declare a local limit.
52/// Invalid syntax or missing reasons return an error without any output.
53pub 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(&section.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    // Parsing above guarantees valid tokens and an AST entry for each start.
69    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}