lol-async 0.3.0

async html-rewriter
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

#![forbid(unsafe_code)]
#![deny(
    missing_copy_implementations,
    rustdoc::missing_crate_level_docs,
    missing_debug_implementations,
    missing_docs,
    nonstandard_style,
    unused_qualifications
)]

/*!
Async wrapper for
[`cloudflare/lol-html`](https://github.com/cloudflare/lol-html).

Since `lol-html` 2.x, [`lol_html::HtmlRewriter`] supports [`Send`]
via the [`lol_html::send`] types. This crate wraps it in an
[`AsyncRead`] implementation, feeding input from an inner async reader
and producing rewritten HTML output.

```
# use async_global_executor as your_async_executor;
# use futures_lite::{io::Cursor, AsyncReadExt};
use lol_async::html::{element, html_content::ContentType, send::Settings};

# your_async_executor::block_on(async {
let mut reader = lol_async::rewrite(
    Cursor::new(r#"<html>
<head><title>hello lol</title></head>
<body><h1>hey there</h1></body>
</html>"#),
    Settings {
        element_content_handlers: vec![element!("h1", |el| {
            el.append("<span>this was inserted</span>", ContentType::Html);
            Ok(())
        })],
        ..Settings::new_send()
    }
);

let mut buf = String::new();
reader.read_to_string(&mut buf).await?;

assert_eq!(buf, r#"<html>
<head><title>hello lol</title></head>
<body><h1>hey there<span>this was inserted</span></h1></body>
</html>"#);
# Result::<_, Box<dyn std::error::Error>>::Ok(()) }).unwrap();
```
*/

use futures_lite::{AsyncRead, ready};
use lol_html::{OutputSink, send::HtmlRewriter};
use pin_project_lite::pin_project;
use std::{
    collections::VecDeque,
    fmt::{self, Debug},
    io::{self, Read},
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll},
};

pub use lol_html as html;
pub use lol_html::send::Settings;

#[derive(Clone)]
struct OutputBuffer(Arc<Mutex<VecDeque<u8>>>);

impl OutputSink for OutputBuffer {
    fn handle_chunk(&mut self, chunk: &[u8]) {
        self.0.lock().unwrap().extend(chunk);
    }
}

pin_project! {
    /// An [`AsyncRead`] adapter that streams rewritten HTML.
    ///
    /// Created by [`rewrite`]. Reads from the inner source, feeds data
    /// through the [`lol_html::HtmlRewriter`], and yields rewritten output.
    pub struct Rewriter<'h, R> {
        #[pin] source: R,
        rewriter: Option<HtmlRewriter<'h, OutputBuffer>>,
        output: OutputBuffer,
        read_buf: Vec<u8>,
    }
}

impl<R> Debug for Rewriter<'_, R> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Rewriter")
            .field("output_len", &self.output.0.lock().unwrap().len())
            .field("done", &self.rewriter.is_none())
            .finish()
    }
}

/// Creates an [`AsyncRead`] that streams HTML rewritten according to the
/// provided [`Settings`](lol_html::send::Settings).
///
/// See the [crate-level documentation](crate) for an example.
pub fn rewrite<'h, R: AsyncRead>(source: R, settings: Settings<'h, '_>) -> Rewriter<'h, R> {
    let output = OutputBuffer(Arc::new(Mutex::new(VecDeque::new())));
    let rewriter = HtmlRewriter::new(settings, output.clone());
    Rewriter {
        source,
        rewriter: Some(rewriter),
        output,
        read_buf: vec![0; 1024],
    }
}

impl<R: AsyncRead> AsyncRead for Rewriter<'_, R> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        let mut this = self.project();

        loop {
            // Drain any buffered output first
            {
                let mut output = this.output.0.lock().unwrap();
                if !output.is_empty() {
                    let n = Read::read(&mut *output, buf)?;
                    return Poll::Ready(Ok(n));
                }
            }

            // If the rewriter is done and no output remains, signal EOF
            let rewriter = match this.rewriter.as_mut() {
                Some(r) => r,
                None => return Poll::Ready(Ok(0)),
            };

            // Read from the source and feed through the rewriter
            match ready!(this.source.as_mut().poll_read(cx, this.read_buf)) {
                Ok(0) => {
                    this.rewriter
                        .take()
                        .unwrap()
                        .end()
                        .map_err(io::Error::other)?;
                    // Loop back to drain any final output
                }

                Ok(n) => {
                    rewriter
                        .write(&this.read_buf[..n])
                        .map_err(io::Error::other)?;
                    // Loop back to drain output
                }

                Err(e) => return Poll::Ready(Err(e)),
            }
        }
    }
}