frontend_environment/
lib.rs

1use lol_html::html_content::ContentType;
2use lol_html::{element, HtmlRewriter, Settings};
3use std::collections::HashMap;
4use std::fmt::Write;
5use std::io;
6#[cfg(feature = "axum")]
7pub use crate::axum::serve_files_with_script;
8
9#[derive(Debug, Clone)]
10/// Map of values that will be provided as environment-variable-like global variables to the frontend.
11pub struct FrontedEnvironment(pub HashMap<String, String>);
12
13/// Rewrites HTML to inject a `<script>` tag (which contains global JS variables that act like environment variables)
14/// into the `<head>` tag.
15pub fn inject_environment_script_tag(
16    input: &[u8],
17    output: &mut Vec<u8>,
18    frontend_env: &FrontedEnvironment,
19) -> io::Result<()> {
20    let mut script_tag = String::new();
21    script_tag.write_str("<script>\n").unwrap();
22    // Writes a line with the content `window.KEY = "VALUE";` for every entry
23    for (key, value) in &frontend_env.0 {
24        script_tag.write_str("window.").unwrap();
25        script_tag.write_str(&key).unwrap();
26        script_tag.write_str(" = \"").unwrap();
27        script_tag.write_str(&value).unwrap();
28        script_tag.write_str("\";\n").unwrap();
29    }
30    script_tag.write_str("</script>").unwrap();
31
32    let mut rewriter = HtmlRewriter::new(
33        Settings {
34            element_content_handlers: vec![element!("head", |el| {
35                el.append(&script_tag, ContentType::Html);
36                Ok(())
37            })],
38            ..Settings::default()
39        },
40        |c: &[u8]| output.extend_from_slice(c),
41    );
42
43    rewriter.write(input).unwrap();
44    rewriter.end().unwrap();
45    Ok(())
46}
47
48#[cfg(feature = "axum")]
49pub mod axum {
50    use ::axum::body::{Body, Bytes, HttpBody};
51    use ::axum::headers::HeaderName;
52    use ::axum::http::{HeaderValue, Request};
53    use ::axum::response::Response;
54    use ::axum::{http, BoxError, Extension};
55    use http_body::combinators::UnsyncBoxBody;
56    use std::convert::Infallible;
57    use tower_http::services::{ServeDir, ServeFile};
58    use super::*;
59
60    /// Static file handler that injects a script tag with environment variables into HTML files.
61    pub async fn serve_files_with_script(
62        Extension(frontend_environment): Extension<FrontedEnvironment>,
63        req: Request<Body>,
64    ) -> Result<Response<UnsyncBoxBody<Bytes, BoxError>>, Infallible> {
65        let mut static_files_service =
66            ServeDir::new("public").not_found_service(ServeFile::new("public/404.html"));
67
68        let res = static_files_service.try_call(req).await.unwrap();
69
70        let headers = res.headers().clone();
71        if headers.get(http::header::CONTENT_TYPE) == Some(&HeaderValue::from_static("text/html")) {
72            let mut res = res.map(move |body| {
73                let body_bytes = body.map_err(Into::into).boxed_unsync();
74                // Inject variables into HTML files
75                body_bytes
76                    .map_data(move |bytes| {
77                        let mut output = Vec::with_capacity(bytes.len() * 2);
78                        inject_environment_script_tag(
79                            &bytes.as_ref(),
80                            &mut output,
81                            &frontend_environment,
82                        )
83                            .unwrap();
84                        output.into()
85                    })
86                    .boxed_unsync()
87            });
88            res.headers_mut()
89                .remove(HeaderName::from_static("content-length"));
90            Ok(res)
91        } else {
92            Ok(res.map(|body| body.map_err(Into::into).boxed_unsync()))
93        }
94    }
95}