rama_http/protocols/html/rewrite/
element.rs

1//! The [`Element`] handle passed to rewrite handlers, plus the
2//! [`ElementContentHandler`] trait.
3
4use std::borrow::Cow;
5use std::fmt;
6use std::str::FromStr;
7
8use rama_core::error::BoxError;
9use rama_utils::byte_set::{set_each, set_range};
10
11use super::super::tokenizer::{HtmlTag, StartTag};
12use super::super::{IntoHtml, escape_attr_value_into};
13
14/// The result of an element content handler. An error aborts the rewrite.
15pub type HandlerResult = Result<(), BoxError>;
16
17/// Bytes that may not appear in an HTML attribute name: C0 controls, `DEL`,
18/// space, and the delimiters that would let a crafted name break out of the
19/// tag (`"` `'` `>` `/` `=` `<`).
20const FORBIDDEN_NAME_BYTE: [bool; 256] = set_each(
21    set_each(set_range([false; 256], 0, 0x20), &[0x7f]),
22    b" \"'<>/=",
23);
24
25/// A validated HTML attribute name.
26///
27/// Constructing one is the only way to name an attribute in
28/// [`Element::set_attribute`], so a name can never inject markup. Build it
29/// from a literal with [`from_static`](Self::from_static), or validate
30/// untrusted input with [`TryFrom`]/[`FromStr`].
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct AttributeName(Cow<'static, str>);
33
34/// Why a string is not a valid [`AttributeName`].
35#[derive(Debug, Clone, PartialEq, Eq)]
36#[non_exhaustive]
37pub enum InvalidAttributeName {
38    /// The name was empty.
39    Empty,
40    /// The name contained a byte that is forbidden in an attribute name.
41    ForbiddenByte(u8),
42}
43
44impl fmt::Display for InvalidAttributeName {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        match self {
47            Self::Empty => f.write_str("empty attribute name"),
48            Self::ForbiddenByte(b) => write!(f, "forbidden byte {b:#04x} in attribute name"),
49        }
50    }
51}
52
53impl std::error::Error for InvalidAttributeName {}
54
55impl AttributeName {
56    /// Validates a literal attribute name at construction.
57    ///
58    /// # Panics
59    ///
60    /// Panics if `name` is empty or contains a byte forbidden in an attribute
61    /// name. Since the input is `'static`, a bad name is a programming error
62    /// caught on first use (and at compile time when used in a `const`).
63    #[must_use]
64    pub const fn from_static(name: &'static str) -> Self {
65        assert!(
66            is_valid_name(name.as_bytes()),
67            "invalid HTML attribute name"
68        );
69        Self(Cow::Borrowed(name))
70    }
71
72    /// The validated name.
73    #[must_use]
74    pub fn as_str(&self) -> &str {
75        &self.0
76    }
77
78    fn into_name_bytes(self) -> Cow<'static, [u8]> {
79        match self.0 {
80            Cow::Borrowed(s) => Cow::Borrowed(s.as_bytes()),
81            Cow::Owned(s) => Cow::Owned(s.into_bytes()),
82        }
83    }
84}
85
86impl TryFrom<&str> for AttributeName {
87    type Error = InvalidAttributeName;
88
89    fn try_from(name: &str) -> Result<Self, Self::Error> {
90        validate_name(name.as_bytes())?;
91        Ok(Self(Cow::Owned(name.to_owned())))
92    }
93}
94
95impl FromStr for AttributeName {
96    type Err = InvalidAttributeName;
97
98    fn from_str(s: &str) -> Result<Self, Self::Err> {
99        Self::try_from(s)
100    }
101}
102
103const fn is_valid_name(bytes: &[u8]) -> bool {
104    if bytes.is_empty() {
105        return false;
106    }
107    let mut i = 0;
108    while i < bytes.len() {
109        if FORBIDDEN_NAME_BYTE[bytes[i] as usize] {
110            return false;
111        }
112        i += 1;
113    }
114    true
115}
116
117fn validate_name(bytes: &[u8]) -> Result<(), InvalidAttributeName> {
118    if bytes.is_empty() {
119        return Err(InvalidAttributeName::Empty);
120    }
121    match bytes
122        .iter()
123        .copied()
124        .find(|&b| FORBIDDEN_NAME_BYTE[b as usize])
125    {
126        Some(b) => Err(InvalidAttributeName::ForbiddenByte(b)),
127        None => Ok(()),
128    }
129}
130
131/// Handles matched elements during a rewrite.
132///
133/// Implement this on your own type — that type holds any state the handler
134/// accumulates, mutated through `&mut self`. `selector` is the index (in
135/// registration order) of the selector that matched `element`.
136///
137/// For one-off rewrites, the closure-based
138/// [`ElementContentHandlers`](super::ElementContentHandlers) builder
139/// implements this trait for you.
140pub trait ElementContentHandler {
141    /// Handles a matched element.
142    ///
143    /// # Errors
144    ///
145    /// Returning an error aborts the rewrite and surfaces the error from
146    /// [`HtmlRewriter::write`](super::HtmlRewriter::write) /
147    /// [`end`](super::HtmlRewriter::end).
148    fn handle_element(&mut self, selector: usize, element: &mut Element<'_>) -> HandlerResult;
149}
150
151/// One attribute in an [`Element`]'s edited attribute list. Unchanged
152/// attributes borrow the source tag; only set/added ones own their bytes.
153#[derive(Debug)]
154struct EditedAttribute<'t> {
155    name: Cow<'t, [u8]>,
156    /// `None` for a valueless attribute (e.g. `disabled`).
157    value: Option<Cow<'t, [u8]>>,
158}
159
160/// What happens to the element's start tag, children and end tag as a whole.
161///
162/// These are mutually exclusive dispositions (last call wins); `before` /
163/// `after` / attribute edits apply on top of any of them.
164#[derive(Debug, Default)]
165enum ElementMode {
166    /// Emit the element unchanged (modulo attribute / `prepend` / `append`
167    /// edits).
168    #[default]
169    Normal,
170    /// Replace the element's children with this (pre-escaped) content; keep
171    /// the start and end tags.
172    Inner(String),
173    /// Replace the whole element (start tag … end tag) with this content.
174    Replace(String),
175    /// Remove the whole element, children included.
176    Remove,
177    /// Remove only the start and end tags, keeping the children.
178    RemoveKeepContent,
179}
180
181/// A matched element, for inspection and mutation by a handler.
182///
183/// Start-anchored edits ([`before`], [`prepend`], attribute changes) take
184/// effect at the start tag; end-anchored edits ([`append`], [`after`],
185/// [`set_inner_content`], [`replace`], [`remove`],
186/// [`remove_and_keep_content`]) are deferred to the matching end tag by the
187/// rewriter.
188///
189/// [`before`]: Element::before
190/// [`prepend`]: Element::prepend
191/// [`append`]: Element::append
192/// [`after`]: Element::after
193/// [`set_inner_content`]: Element::set_inner_content
194/// [`replace`]: Element::replace
195/// [`remove`]: Element::remove
196/// [`remove_and_keep_content`]: Element::remove_and_keep_content
197pub struct Element<'t> {
198    tag: &'t StartTag<'t>,
199    before: String,
200    prepend: String,
201    append: String,
202    after: String,
203    /// `None` until an attribute is edited; then it is the full attribute
204    /// list to re-serialize.
205    attributes: Option<Vec<EditedAttribute<'t>>>,
206    mode: ElementMode,
207}
208
209impl<'t> Element<'t> {
210    pub(crate) fn new(tag: &'t StartTag<'t>) -> Self {
211        Self {
212            tag,
213            before: String::new(),
214            prepend: String::new(),
215            append: String::new(),
216            after: String::new(),
217            attributes: None,
218            mode: ElementMode::Normal,
219        }
220    }
221
222    /// The element's [`HtmlTag`] — a known element variant, or
223    /// [`HtmlTag::Other`] (the original-case name) for a custom tag.
224    #[must_use]
225    pub fn tag(&self) -> HtmlTag<'_> {
226        self.tag.tag()
227    }
228
229    /// The value of the attribute with the given (ASCII case-insensitive)
230    /// name, or `None` if absent. A valueless attribute reports `Some(b"")`.
231    #[must_use]
232    pub fn attribute(&self, name: &str) -> Option<&[u8]> {
233        let name = name.as_bytes();
234        match &self.attributes {
235            Some(edited) => edited
236                .iter()
237                .find(|a| a.name.eq_ignore_ascii_case(name))
238                .map(|a| a.value.as_deref().unwrap_or(b"")),
239            None => self
240                .tag
241                .attributes()
242                .find(|a| a.name().eq_ignore_ascii_case(name))
243                .map(|a| if a.has_value() { a.value() } else { b"" }),
244        }
245    }
246
247    /// Whether the element has the given attribute.
248    #[must_use]
249    pub fn has_attribute(&self, name: &str) -> bool {
250        self.attribute(name).is_some()
251    }
252
253    /// Sets (or adds) an attribute. The [`AttributeName`] is pre-validated, so
254    /// it cannot inject markup; the value is escaped on render.
255    pub fn set_attribute(&mut self, name: AttributeName, value: &str) {
256        self.ensure_attributes();
257        let Some(attributes) = self.attributes.as_mut() else {
258            return;
259        };
260        let name = name.into_name_bytes();
261        if let Some(existing) = attributes
262            .iter_mut()
263            .find(|a| a.name.eq_ignore_ascii_case(&name))
264        {
265            existing.value = Some(Cow::Owned(value.as_bytes().to_vec()));
266        } else {
267            attributes.push(EditedAttribute {
268                name,
269                value: Some(Cow::Owned(value.as_bytes().to_vec())),
270            });
271        }
272    }
273
274    /// Removes the attribute with the given (ASCII case-insensitive) name.
275    pub fn remove_attribute(&mut self, name: &str) {
276        self.ensure_attributes();
277        let Some(attributes) = self.attributes.as_mut() else {
278            return;
279        };
280        let name = name.as_bytes();
281        attributes.retain(|a| !a.name.eq_ignore_ascii_case(name));
282    }
283
284    /// Inserts content immediately before the element's start tag.
285    ///
286    /// Accepts any [`IntoHtml`] value: plain strings are escaped; wrap
287    /// trusted HTML in [`PreEscaped`](super::super::PreEscaped) to emit it
288    /// verbatim.
289    pub fn before(&mut self, content: impl IntoHtml) {
290        reserve_html(&mut self.before, &content);
291        content.escape_and_write(&mut self.before);
292    }
293
294    /// Inserts content as the element's first children (immediately after the
295    /// start tag). Escaping follows [`before`](Self::before).
296    pub fn prepend(&mut self, content: impl IntoHtml) {
297        reserve_html(&mut self.prepend, &content);
298        content.escape_and_write(&mut self.prepend);
299    }
300
301    /// Inserts content as the element's last children (immediately before the
302    /// end tag). Escaping follows [`before`](Self::before).
303    pub fn append(&mut self, content: impl IntoHtml) {
304        reserve_html(&mut self.append, &content);
305        content.escape_and_write(&mut self.append);
306    }
307
308    /// Inserts content immediately after the element's end tag. Escaping
309    /// follows [`before`](Self::before).
310    pub fn after(&mut self, content: impl IntoHtml) {
311        reserve_html(&mut self.after, &content);
312        content.escape_and_write(&mut self.after);
313    }
314
315    /// Replaces the element's children, keeping the start and end tags (and
316    /// any attribute edits). Escaping follows [`before`](Self::before).
317    pub fn set_inner_content(&mut self, content: impl IntoHtml) {
318        let mut inner = String::with_capacity(html_capacity(&content));
319        content.escape_and_write(&mut inner);
320        self.mode = ElementMode::Inner(inner);
321    }
322
323    /// Replaces the whole element (start tag through end tag). Escaping
324    /// follows [`before`](Self::before).
325    pub fn replace(&mut self, content: impl IntoHtml) {
326        let mut replacement = String::with_capacity(html_capacity(&content));
327        content.escape_and_write(&mut replacement);
328        self.mode = ElementMode::Replace(replacement);
329    }
330
331    /// Removes the whole element, children included.
332    pub fn remove(&mut self) {
333        self.mode = ElementMode::Remove;
334    }
335
336    /// Removes only the element's start and end tags, leaving its children in
337    /// place.
338    pub fn remove_and_keep_content(&mut self) {
339        self.mode = ElementMode::RemoveKeepContent;
340    }
341
342    /// Whether a [`remove`](Self::remove) /
343    /// [`remove_and_keep_content`](Self::remove_and_keep_content) /
344    /// [`replace`](Self::replace) disposition is in effect.
345    #[must_use]
346    pub fn is_removed(&self) -> bool {
347        matches!(
348            self.mode,
349            ElementMode::Remove | ElementMode::RemoveKeepContent | ElementMode::Replace(_)
350        )
351    }
352
353    fn ensure_attributes(&mut self) {
354        if self.attributes.is_none() {
355            let attributes = self
356                .tag
357                .attributes()
358                .map(|a| EditedAttribute {
359                    name: Cow::Borrowed(a.name()),
360                    value: a.has_value().then(|| Cow::Borrowed(a.value())),
361                })
362                .collect();
363            self.attributes = Some(attributes);
364        }
365    }
366
367    /// Emits the element's start-side output into `out` (only when `visible`
368    /// — i.e. not swallowed by an enclosing removed/replaced ancestor) and
369    /// returns the [`EndActions`] to apply at the matching end tag. Consumes
370    /// `self` so the owned edit buffers move out without copying.
371    pub(crate) fn serialize(self, out: &mut Vec<u8>, visible: bool) -> EndActions {
372        let Self {
373            tag,
374            before,
375            prepend,
376            append,
377            after,
378            attributes,
379            mode,
380        } = self;
381
382        match mode {
383            ElementMode::Normal => {
384                if visible {
385                    out.extend_from_slice(before.as_bytes());
386                    emit_start_tag(out, tag, attributes.as_deref());
387                    out.extend_from_slice(prepend.as_bytes());
388                }
389                EndActions {
390                    append,
391                    after,
392                    suppress_content: false,
393                    suppress_end_tag: false,
394                }
395            }
396            ElementMode::Inner(inner) => {
397                if visible {
398                    out.extend_from_slice(before.as_bytes());
399                    emit_start_tag(out, tag, attributes.as_deref());
400                    out.extend_from_slice(prepend.as_bytes());
401                    out.extend_from_slice(inner.as_bytes());
402                }
403                EndActions {
404                    append,
405                    after,
406                    suppress_content: true,
407                    suppress_end_tag: false,
408                }
409            }
410            ElementMode::Replace(replacement) => {
411                if visible {
412                    out.extend_from_slice(before.as_bytes());
413                    out.extend_from_slice(replacement.as_bytes());
414                }
415                // The element (and its children/end tag) is gone; only
416                // `after` still applies, at the end-tag position.
417                EndActions {
418                    append: String::new(),
419                    after,
420                    suppress_content: true,
421                    suppress_end_tag: true,
422                }
423            }
424            ElementMode::Remove => {
425                if visible {
426                    out.extend_from_slice(before.as_bytes());
427                }
428                EndActions {
429                    append: String::new(),
430                    after,
431                    suppress_content: true,
432                    suppress_end_tag: true,
433                }
434            }
435            ElementMode::RemoveKeepContent => {
436                if visible {
437                    out.extend_from_slice(before.as_bytes());
438                    // Start tag dropped; `prepend` sits where it was.
439                    out.extend_from_slice(prepend.as_bytes());
440                }
441                EndActions {
442                    append,
443                    after,
444                    suppress_content: false,
445                    suppress_end_tag: true,
446                }
447            }
448        }
449    }
450}
451
452/// Output to emit at an element's matching end tag, returned by
453/// [`Element::serialize`].
454pub(crate) struct EndActions {
455    /// Emitted just before the end tag (the element's last children).
456    pub(crate) append: String,
457    /// Emitted just after the end tag.
458    pub(crate) after: String,
459    /// Whether the element's children must be suppressed from the output.
460    pub(crate) suppress_content: bool,
461    /// Whether the end tag itself must be suppressed.
462    pub(crate) suppress_end_tag: bool,
463}
464
465impl EndActions {
466    /// The no-op actions for an opened element that needs no end-side edits
467    /// (an unmatched element, or one with only start-anchored edits).
468    pub(crate) fn passthrough() -> Self {
469        Self {
470            append: String::new(),
471            after: String::new(),
472            suppress_content: false,
473            suppress_end_tag: false,
474        }
475    }
476}
477
478/// Emits an element's start tag, re-serializing when attributes were edited
479/// (`edited` is `Some`) or passing the original bytes through verbatim.
480fn emit_start_tag(out: &mut Vec<u8>, tag: &StartTag<'_>, edited: Option<&[EditedAttribute<'_>]>) {
481    match edited {
482        None => out.extend_from_slice(tag.raw()),
483        Some(attributes) => {
484            out.push(b'<');
485            out.extend_from_slice(tag.name());
486            for attr in attributes {
487                out.push(b' ');
488                out.extend_from_slice(&attr.name);
489                if let Some(value) = &attr.value {
490                    out.extend_from_slice(b"=\"");
491                    escape_attr_value_into(out, value);
492                    out.push(b'"');
493                }
494            }
495            if tag.is_self_closing() {
496                out.extend_from_slice(b" />");
497            } else {
498                out.push(b'>');
499            }
500        }
501    }
502}
503
504fn reserve_html(buf: &mut String, content: &impl IntoHtml) {
505    buf.reserve(html_capacity(content));
506}
507
508fn html_capacity(content: &impl IntoHtml) -> usize {
509    let hint = content.size_hint();
510    hint + (hint / 10)
511}
512
513#[cfg(test)]
514mod tests {
515    use super::{AttributeName, InvalidAttributeName};
516
517    #[test]
518    fn valid_names() {
519        for name in ["class", "data-x", "aria-label", "x", "ns:attr"] {
520            assert_eq!(AttributeName::from_static(name).as_str(), name);
521            assert_eq!(AttributeName::try_from(name).unwrap().as_str(), name);
522            assert_eq!(name.parse::<AttributeName>().unwrap().as_str(), name);
523        }
524    }
525
526    #[test]
527    fn rejects_invalid_names() {
528        assert_eq!(
529            AttributeName::try_from(""),
530            Err(InvalidAttributeName::Empty)
531        );
532        // Every byte that could break out of the start tag is rejected.
533        for (input, byte) in [
534            ("a b", b' '),
535            ("a\tb", b'\t'),
536            ("a=b", b'='),
537            ("a>b", b'>'),
538            ("a/b", b'/'),
539            ("a<b", b'<'),
540            ("a\"b", b'"'),
541            ("a'b", b'\''),
542            ("a\x7fb", 0x7f),
543        ] {
544            assert_eq!(
545                AttributeName::try_from(input),
546                Err(InvalidAttributeName::ForbiddenByte(byte)),
547                "{input:?}"
548            );
549        }
550    }
551
552    #[test]
553    #[should_panic = "invalid HTML attribute name"]
554    fn from_static_panics_on_injection() {
555        let _name = AttributeName::from_static("x onload=alert(1)");
556    }
557}
rama_http/protocols/html/rewrite/element.rs

rama_http/protocols/html/rewrite/
element.rs
