lol_async/
lib.rs

1#![forbid(unsafe_code)]
2#![deny(
3    missing_copy_implementations,
4    rustdoc::missing_crate_level_docs,
5    missing_debug_implementations,
6    missing_docs,
7    nonstandard_style,
8    unused_qualifications
9)]
10
11/*!
12Async wrapper for
13[`cloudflare/lol-html`](https://github.com/cloudflare/lol-html).
14
15Since `lol-html` 2.x, [`lol_html::HtmlRewriter`] supports [`Send`]
16via the [`lol_html::send`] types. This crate wraps it in an
17[`AsyncRead`] implementation, feeding input from an inner async reader
18and producing rewritten HTML output.
19
20```
21# use async_global_executor as your_async_executor;
22# use futures_lite::{io::Cursor, AsyncReadExt};
23use lol_async::html::{element, html_content::ContentType, send::Settings};
24
25# your_async_executor::block_on(async {
26let mut reader = lol_async::rewrite(
27    Cursor::new(r#"<html>
28<head><title>hello lol</title></head>
29<body><h1>hey there</h1></body>
30</html>"#),
31    Settings {
32        element_content_handlers: vec![element!("h1", |el| {
33            el.append("<span>this was inserted</span>", ContentType::Html);
34            Ok(())
35        })],
36        ..Settings::new_send()
37    }
38);
39
40let mut buf = String::new();
41reader.read_to_string(&mut buf).await?;
42
43assert_eq!(buf, r#"<html>
44<head><title>hello lol</title></head>
45<body><h1>hey there<span>this was inserted</span></h1></body>
46</html>"#);
47# Result::<_, Box<dyn std::error::Error>>::Ok(()) }).unwrap();
48```
49*/
50
51use futures_lite::{AsyncRead, ready};
52use lol_html::{OutputSink, send::HtmlRewriter};
53use pin_project_lite::pin_project;
54use std::{
55    collections::VecDeque,
56    fmt::{self, Debug},
57    io::{self, Read},
58    pin::Pin,
59    sync::{Arc, Mutex},
60    task::{Context, Poll},
61};
62
63pub use lol_html as html;
64pub use lol_html::send::Settings;
65
66#[derive(Clone)]
67struct OutputBuffer(Arc<Mutex<VecDeque<u8>>>);
68
69impl OutputSink for OutputBuffer {
70    fn handle_chunk(&mut self, chunk: &[u8]) {
71        self.0.lock().unwrap().extend(chunk);
72    }
73}
74
75pin_project! {
76    /// An [`AsyncRead`] adapter that streams rewritten HTML.
77    ///
78    /// Created by [`rewrite`]. Reads from the inner source, feeds data
79    /// through the [`lol_html::HtmlRewriter`], and yields rewritten output.
80    pub struct Rewriter<'h, R> {
81        #[pin] source: R,
82        rewriter: Option<HtmlRewriter<'h, OutputBuffer>>,
83        output: OutputBuffer,
84        read_buf: Vec<u8>,
85    }
86}
87
88impl<R> Debug for Rewriter<'_, R> {
89    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90        f.debug_struct("Rewriter")
91            .field("output_len", &self.output.0.lock().unwrap().len())
92            .field("done", &self.rewriter.is_none())
93            .finish()
94    }
95}
96
97/// Creates an [`AsyncRead`] that streams HTML rewritten according to the
98/// provided [`Settings`](lol_html::send::Settings).
99///
100/// See the [crate-level documentation](crate) for an example.
101pub fn rewrite<'h, R: AsyncRead>(source: R, settings: Settings<'h, '_>) -> Rewriter<'h, R> {
102    let output = OutputBuffer(Arc::new(Mutex::new(VecDeque::new())));
103    let rewriter = HtmlRewriter::new(settings, output.clone());
104    Rewriter {
105        source,
106        rewriter: Some(rewriter),
107        output,
108        read_buf: vec![0; 1024],
109    }
110}
111
112impl<R: AsyncRead> AsyncRead for Rewriter<'_, R> {
113    fn poll_read(
114        self: Pin<&mut Self>,
115        cx: &mut Context<'_>,
116        buf: &mut [u8],
117    ) -> Poll<io::Result<usize>> {
118        let mut this = self.project();
119
120        loop {
121            // Drain any buffered output first
122            {
123                let mut output = this.output.0.lock().unwrap();
124                if !output.is_empty() {
125                    let n = Read::read(&mut *output, buf)?;
126                    return Poll::Ready(Ok(n));
127                }
128            }
129
130            // If the rewriter is done and no output remains, signal EOF
131            let rewriter = match this.rewriter.as_mut() {
132                Some(r) => r,
133                None => return Poll::Ready(Ok(0)),
134            };
135
136            // Read from the source and feed through the rewriter
137            match ready!(this.source.as_mut().poll_read(cx, this.read_buf)) {
138                Ok(0) => {
139                    this.rewriter
140                        .take()
141                        .unwrap()
142                        .end()
143                        .map_err(io::Error::other)?;
144                    // Loop back to drain any final output
145                }
146
147                Ok(n) => {
148                    rewriter
149                        .write(&this.read_buf[..n])
150                        .map_err(io::Error::other)?;
151                    // Loop back to drain output
152                }
153
154                Err(e) => return Poll::Ready(Err(e)),
155            }
156        }
157    }
158}
lol_async/lib.rs

lol_async/
lib.rs
