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

lol_async/
lib.rs
