Skip to main content

trillium_html_rewriter/
lib.rs

1#![forbid(unsafe_code)]
2#![deny(
3    clippy::dbg_macro,
4    missing_copy_implementations,
5    rustdoc::missing_crate_level_docs,
6    missing_debug_implementations,
7    missing_docs,
8    nonstandard_style,
9    unused_qualifications
10)]
11#![doc = include_str!("../README.md")]
12
13use lol_async::rewrite;
14pub use lol_async::{Settings, html};
15use mime::Mime;
16use std::{
17    fmt::{self, Debug, Formatter},
18    str::FromStr,
19    sync::Arc,
20};
21use trillium::{
22    Body, Conn, Handler,
23    KnownHeaderName::{ContentLength, ContentType},
24};
25
26/// A trillium [`Handler`] that rewrites HTML response bodies with
27/// [`lol-html`](https://docs.rs/lol-html), using [`lol-async`](https://docs.rs/lol-async).
28///
29/// It wraps the response produced by other handlers: in [`before_send`](Handler::before_send) it
30/// inspects the outgoing `Content-Type` and, if the mime subtype is `html` (e.g. `text/html`),
31/// replaces the response body with a streaming rewrite driven by the [`Settings`] returned from the
32/// closure passed to [`HtmlRewriter::new`]. Responses with any other content type (or none) are
33/// passed through unchanged.
34pub struct HtmlRewriter {
35    settings: Arc<dyn Fn() -> Settings<'static, 'static> + Send + Sync + 'static>,
36}
37
38impl Debug for HtmlRewriter {
39    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
40        f.debug_struct("HtmlRewriter").finish()
41    }
42}
43
44impl Handler for HtmlRewriter {
45    async fn before_send(&self, mut conn: Conn) -> Conn {
46        let html = conn
47            .response_headers()
48            .get_str(ContentType)
49            .and_then(|c| Mime::from_str(c).ok())
50            .map(|m| m.subtype() == "html")
51            .unwrap_or_default();
52
53        if html && let Some(body) = conn.take_response_body() {
54            let reader = rewrite(body, (self.settings)());
55            conn.response_headers_mut().remove(ContentLength); // we no longer know the content length, if we ever did
56            conn.with_body(Body::new_streaming(reader, None))
57        } else {
58            conn
59        }
60    }
61}
62
63impl HtmlRewriter {
64    /// Construct a new html rewriter from a closure that builds [`Settings`].
65    ///
66    /// The closure — rather than a `Settings` value — is required because `lol-html`'s content
67    /// handlers are single-use; it is invoked once per rewritten response to produce a fresh set of
68    /// handlers. Build the settings with [`Settings::new_send()`] as the base (its handlers are
69    /// `Send`, as required here) and populate `element_content_handlers` /
70    /// `document_content_handlers`. See [`lol_async::html::Settings`] and the
71    /// [`lol-html`](https://docs.rs/lol-html) docs for the full rewriting API.
72    pub fn new(f: impl Fn() -> Settings<'static, 'static> + Send + Sync + 'static) -> Self {
73        Self {
74            settings: Arc::new(f)
75                as Arc<dyn Fn() -> Settings<'static, 'static> + Send + Sync + 'static>,
76        }
77    }
78}