kvarn_extensions/
push.rs

1use std::collections::HashSet;
2
3use crate::*;
4
5/// Mounts a push extension with priority `-32`, overriding any other [`Post`] extension with that
6/// priority.
7///
8/// This only pushes new content to each connection every 2 minutes (if you use
9/// [`SmartPush::default`]), not every time.
10pub fn mount(extensions: &mut Extensions, manager: SmartPush) -> &mut Extensions {
11    let manager = Mutex::new(manager);
12
13    extensions.add_post(
14        post!(
15            request,
16            host,
17            response_pipe,
18            identity_body,
19            addr,
20            move |manager: Mutex<SmartPush>| {
21                push(
22                    request,
23                    host,
24                    response_pipe,
25                    identity_body,
26                    addr,
27                    Some(manager),
28                )
29                .await
30            }
31        ),
32        Id::new(-32, "HTTP/2 push"),
33    );
34    extensions
35}
36
37/// Always push all links on page.
38///
39/// # Examples
40///
41/// ```rust
42/// # use kvarn::prelude::*;
43/// # let mut extensions = Extensions::new();
44/// use kvarn_extensions::push::always;
45/// extensions.add_post(
46///     Box::new(always),
47///     Id::new(-32, "HTTP/2 push"),
48/// );
49/// ```
50pub fn always<'a>(
51    request: &'a FatRequest,
52    host: &'a Host,
53    response_pipe: &'a mut application::ResponseBodyPipe,
54    bytes: Bytes,
55    addr: SocketAddr,
56) -> RetFut<'a, ()> {
57    Box::pin(push(request, host, response_pipe, bytes, addr, None))
58}
59
60pub struct SmartPush {
61    db: HashSet<SocketAddr>,
62    last_clear: Instant,
63    clear_interval: Duration,
64    check_every_request: u32,
65    iteration: u32,
66}
67impl SmartPush {
68    /// `clear_interval` is the duration between clearing the log of who's been pushed content.
69    ///
70    /// `check_every_request` is the number of requests between checks if the duration has
71    /// expired.
72    pub fn new(clear_interval: Duration, check_every_request: u32) -> Self {
73        Self {
74            db: HashSet::new(),
75            last_clear: Instant::now(),
76            clear_interval,
77            check_every_request,
78            iteration: 0,
79        }
80    }
81    fn accept(&mut self, remote: SocketAddr) -> bool {
82        if self.iteration >= self.check_every_request {
83            let now = Instant::now();
84            let elapsed = now - self.last_clear;
85            if elapsed > self.clear_interval {
86                self.last_clear = now;
87                self.db.clear();
88            }
89        }
90        self.iteration += 1;
91
92        !self.db.contains(&remote)
93    }
94    fn register(&mut self, remote: SocketAddr) {
95        self.db.insert(remote);
96    }
97}
98impl Default for SmartPush {
99    fn default() -> Self {
100        Self::new(Duration::from_secs(60 * 2), 8)
101    }
102}
103
104async fn push<'a>(
105    request: &'a FatRequest,
106    host: &'a Host,
107    response_pipe: &'a mut application::ResponseBodyPipe,
108    bytes: Bytes,
109    addr: SocketAddr,
110    manager: Option<&'a Mutex<SmartPush>>,
111) {
112    use internals::*;
113    // let request = unsafe { request.get_inner() };
114    // let response_pipe = unsafe { response_pipe.get_inner() };
115
116    // If it is not HTTP/2
117    #[allow(irrefutable_let_patterns)]
118    if !matches!(response_pipe, ResponseBodyPipe::Http2(_, _)) {
119        return;
120    }
121
122    if let Some(manager) = manager {
123        // let manager = unsafe { manager.get() };
124        let mut lock = manager.lock().await;
125        if !lock.accept(addr) {
126            return;
127        }
128    }
129
130    // If user agent is Firefox, return.
131    // This implementations of push doesn not work with Firefox!
132    // I do not know why. Any help is appreciated.
133    // Kvarn follows the HTTP/2 spec completely, according to h2spec.
134    if request
135        .headers()
136        .get("user-agent")
137        .and_then(|user_agent| user_agent.to_str().ok())
138        .map_or(false, |user_agent| user_agent.contains("Firefox/"))
139    {
140        return;
141    }
142
143    const HTML_START: &str = "<!DOCTYPE html>";
144
145    match str::from_utf8(&bytes) {
146        // If it is HTML
147        Ok(string)
148            if string
149                .get(..HTML_START.len())
150                .map_or(false, |s| s.eq_ignore_ascii_case(HTML_START)) =>
151        {
152            let mut urls: Vec<_> = url_crawl::get_urls(string).map(String::from).collect();
153
154            // remove images
155            urls.retain(|url| {
156                !url.contains(".jpg")
157                    && !url.contains(".avif")
158                    && !url.contains("png")
159                    && !url.contains(".webp")
160                    && !url.contains(".gif")
161            });
162
163            for url in &mut urls {
164                if !url.starts_with('/') && !url.contains(':') {
165                    let path = request.uri().path();
166                    let mut last_slash = 0;
167                    for (pos, c) in path.chars().enumerate() {
168                        if c == '/' {
169                            last_slash = pos;
170                        }
171                    }
172                    url.insert_str(0, &path[..=last_slash]);
173                }
174            }
175
176            debug!("Pushing urls {:?}", urls);
177
178            urls.sort_unstable();
179            urls.dedup();
180
181            for url in urls {
182                let mut uri = request.uri().clone().into_parts();
183                if let Some(uri) =
184                    uri::PathAndQuery::from_maybe_shared::<Bytes>(url.into_bytes().into())
185                        .ok()
186                        .and_then(|path| {
187                            uri.path_and_query = Some(path);
188                            Uri::from_parts(uri).ok()
189                        })
190                {
191                    let mut push_request = Request::builder().uri(uri);
192                    macro_rules! copy_header {
193                        ($builder: expr, $headers: expr, $name: expr) => {
194                            for header in $headers.get_all($name) {
195                                $builder = $builder.header($name, header);
196                            }
197                        };
198                    }
199                    let headers = request.headers();
200
201                    copy_header!(push_request, headers, "accept-encoding");
202                    copy_header!(push_request, headers, "accept-language");
203                    copy_header!(push_request, headers, "user-agent");
204                    copy_header!(push_request, headers, "host");
205                    copy_header!(push_request, headers, "origin");
206                    copy_header!(push_request, headers, "cookies");
207
208                    let push_request = push_request
209                        .body(())
210                        .expect("failed to construct a request only from another valid request.");
211
212                    let empty_request = utils::empty_clone_request(&push_request);
213
214                    let response_pipe = match response_pipe.push_request(empty_request) {
215                        Ok(pipe) => pipe,
216                        Err(_) => return,
217                    };
218
219                    let mut push_request =
220                        push_request.map(|_| kvarn::application::Body::Bytes(Bytes::new().into()));
221
222                    let response = kvarn::handle_cache(&mut push_request, addr, host).await;
223
224                    if let Err(err) = kvarn::SendKind::Push(response_pipe)
225                        .send(response, request, host, addr)
226                        .await
227                    {
228                        error!("Error occurred when pushing request. {:?}", err);
229                    }
230                }
231            }
232
233            debug!("Push done.");
234        }
235        // Else, do nothing
236        _ => {}
237    }
238    if let Some(manager) = manager {
239        // let manager = unsafe { manager.get() };
240        let mut lock = manager.lock().await;
241        lock.register(addr);
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248    #[tokio::test]
249    async fn run() {
250        let mut extensions = Extensions::new();
251        mount(&mut extensions, SmartPush::default());
252        let _server = kvarn_testing::ServerBuilder::from(extensions).run().await;
253    }
254    #[test]
255    fn exclusive() {
256        let mut extensions = new();
257        extensions.add_post(Box::new(always), Id::new(-32, "HTTP/2 push"));
258
259        let debug = format!("{extensions:?}");
260        assert_eq!(debug.match_indices("push").count(), 1);
261        mount(&mut extensions, SmartPush::default());
262        assert_eq!(debug.match_indices("push").count(), 1);
263    }
264}
kvarn_extensions/push.rs

kvarn_extensions/
push.rs
