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/*!
12This crate is an effort to adapt
13[`cloudflare/lol-html`](https://github.com/cloudflare/lol-html) for an
14async context. Unfortunately, due to lol-html's design, the wrapped
15api is not ideal. In particular, the [`lol_html::HtmlRewriter`] is
16`!Send`, which means that we have some contortions in order to use
17this library in an async context that expects most futures to be
18`Send`. This crate addresses this by returning two types: A
19[`LolFuture`] that must be polled on whatever thread calls [`rewrite`]
20and a [`LolReader`] that is [`AsyncRead`] and can be moved around and
21sent between threads as needed.
22
23❗ **Due to this design, it is necessary to poll the [`LolFuture`] in
24addition to reading from the [`LolReader`]** ❗
25
26## Improvements
27
28Improvements to the design of this crate are very welcome. I don't
29have a lot of experience working around `!Send` types in an async
30context, and although this crate achieved the result I needed, I hate
31it. Please open a PR or write a better crate! Alternatively, if you
32know someone at cloudflare, maybe arms can be twisted to adapt
33`lol-html` for async rust.
34
35```
36# use async_global_executor as your_async_executor;
37# use futures_lite::{io::Cursor, AsyncReadExt};
38use lol_async::html::{element, html_content::ContentType, Settings};
39
40# your_async_executor::block_on(async {
41let (fut, mut reader) = lol_async::rewrite(
42    Cursor::new(r#"<html>
43<head><title>hello lol</title></head>
44<body><h1>hey there</h1></body>
45</html>"#),
46    Settings {
47        element_content_handlers: vec![element!("h1", |el| {
48            el.append("<span>this was inserted</span>", ContentType::Html);
49            Ok(())
50        })],
51        ..Settings::default()
52    }
53);
54
55let handle = your_async_executor::spawn_local(fut);
56
57let mut buf = String::new();
58reader.read_to_string(&mut buf).await?;
59
60handle.await?;
61assert_eq!(buf, r#"<html>
62<head><title>hello lol</title></head>
63<body><h1>hey there<span>this was inserted</span></h1></body>
64</html>"#);
65# Result::<_, Box<dyn std::error::Error>>::Ok(()) }).unwrap();
66```
67
68
69*/
70use atomic_waker::AtomicWaker;
71use futures_lite::{ready, AsyncRead};
72use lockfree::queue::Queue;
73use lol_html::{HtmlRewriter, OutputSink, Settings};
74use pin_project_lite::pin_project;
75use std::{
76    fmt::{self, Debug},
77    future::Future,
78    io::Result,
79    pin::Pin,
80    sync::{
81        atomic::{AtomicBool, Ordering},
82        Arc,
83    },
84    task::{Context, Poll},
85};
86
87pub use lol_html as html;
88
89impl<R> Debug for LolFuture<'_, R> {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        f.debug_struct("LolFuture")
92            .field("buffer", &self.buffer)
93            .finish()
94    }
95}
96
97impl Debug for LolReader {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        f.debug_struct("LolReader")
100            .field("waker", &self.waker)
101            .field("done", &self.done)
102            .finish()
103    }
104}
105
106/**
107An [`AsyncRead`] type that will yield the rewritten html. LolReader is
108Send, allowing it to be used on a different thread from the paired
109[`LolFuture`]. Please note that reading from LolReader does not drive the
110LolHtml rewriter. Awaiting the [`LolFuture`] is also necessary.
111*/
112
113pub struct LolReader {
114    waker: Arc<AtomicWaker>,
115    done: Arc<AtomicBool>,
116    data: Arc<Queue<u8>>,
117}
118
119struct LolOutputter {
120    done: Arc<AtomicBool>,
121    waker: Arc<AtomicWaker>,
122    data: Arc<Queue<u8>>,
123}
124
125pin_project! {
126    /**
127    `await` this [`Future`] to drive the html rewriting process. The
128    LolFuture contains the [`HtmlRewriter`] and as a result is
129    `!Send`, so it must be spawned locally.
130    */
131    pub struct LolFuture<'h, R> {
132        #[pin] source: R,
133        rewriter: Option<HtmlRewriter<'h, LolOutputter>>,
134        buffer: Vec<u8>
135    }
136}
137
138/**
139This function is the primary entrypoint for `lol-async`. It takes a
140data `Source` that is [`AsyncRead`] and a [`Settings`] that describes
141the desired rewriting logic. It returns a !Send [`LolFuture`] future
142that drives the rewriter on the current thread and a [`LolReader`]
143that is Send and can be used anywhere an [`AsyncRead`] would be
144used. The html content yielded by the [`LolReader`] will be rewritten
145according to the rules specified in the Settings.
146*/
147pub fn rewrite<'h, Source>(
148    source: Source,
149    settings: Settings<'h, '_>,
150) -> (LolFuture<'h, Source>, LolReader)
151where
152    Source: AsyncRead,
153{
154    let waker = Arc::new(AtomicWaker::new());
155    let done = Arc::new(AtomicBool::new(false));
156    let data = Arc::new(Queue::new());
157
158    let output_sink = LolOutputter {
159        waker: waker.clone(),
160        done: done.clone(),
161        data: data.clone(),
162    };
163
164    let rewriter = HtmlRewriter::new(settings, output_sink);
165
166    let future = LolFuture {
167        source,
168        rewriter: Some(rewriter),
169        buffer: vec![0; 1024],
170    };
171
172    let reader = LolReader { waker, done, data };
173
174    (future, reader)
175}
176
177impl<Source: AsyncRead> Future for LolFuture<'_, Source> {
178    type Output = Result<()>;
179    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
180        let mut this = self.project();
181        loop {
182            match this.rewriter {
183                Some(rewriter) => match ready!(this.source.as_mut().poll_read(cx, this.buffer))? {
184                    0 => {
185                        this.rewriter.take().unwrap().end().unwrap();
186                    }
187
188                    new_bytes => {
189                        rewriter.write(&this.buffer[..new_bytes]).unwrap();
190                    }
191                },
192
193                None => return Poll::Ready(Ok(())),
194            }
195        }
196    }
197}
198
199impl OutputSink for LolOutputter {
200    fn handle_chunk(&mut self, chunk: &[u8]) {
201        if chunk.is_empty() {
202            self.done.store(true, Ordering::SeqCst);
203        } else {
204            self.data.extend(chunk.to_vec());
205        }
206
207        self.waker.wake();
208    }
209}
210
211impl AsyncRead for LolReader {
212    fn poll_read(
213        self: Pin<&mut Self>,
214        cx: &mut Context<'_>,
215        buf: &mut [u8],
216    ) -> Poll<Result<usize>> {
217        let data: Vec<u8> = self.data.pop_iter().take(buf.len()).collect();
218        let len = data.len();
219        if len > 0 {
220            buf[..len].copy_from_slice(&data);
221            Poll::Ready(Ok(len))
222        } else if self.done.load(Ordering::SeqCst) {
223            Poll::Ready(Ok(0))
224        } else {
225            self.waker.register(cx.waker());
226            Poll::Pending
227        }
228    }
229}
lol_async/lib.rs

lol_async/
lib.rs
