[ofl] Use lol_html for injecting reload script.

Closes #127.
anp · Jul 1, 2020 · 0647d6c · 0647d6c
1 parent b44ddee
commit 0647d6c
Show file tree

Hide file tree

Showing 3 changed files with 170 additions and 9 deletions.
diff --git a/ofl/Cargo.lock b/ofl/Cargo.lock
diff --git a/ofl/Cargo.toml b/ofl/Cargo.toml
@@ -22,6 +22,7 @@ futures = "0.3"
 grcov = "0.5.15"
 gumdrop = { version = "0.8", features = ["default_expr"] }
 http = "0.2"
+lol_html = "0.2.0"
 mdbook = "0.3.1"
 notify = "5.0.0-pre.2"
 opener = "0.4"

diff --git a/ofl/src/server/inject.rs b/ofl/src/server/inject.rs
@@ -4,8 +4,13 @@ use actix_web::{
     web,
 };
 use futures::stream::{Stream, StreamExt};
+use lol_html::{element, html_content::ContentType, RewriteStrSettings};
 use tracing::*;
 
+const RELOAD_ON_CHANGES: &str =
+    include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/reloadOnChanges.js"));
+const CHANGES_URL: &str = "/ch-ch-ch-changes";
+
 pub async fn reload_on_changes_into_html<B>(
     mut response: ServiceResponse<B>,
 ) -> Result<ServiceResponse<B>, actix_web::error::Error>
@@ -19,21 +24,34 @@ where
     }
 
     let path = response.request().path();
-    if path == "/ch-ch-ch-changes" {
+    if path == CHANGES_URL {
         return Ok(response);
     }
     info!({ %path }, "serving html page");
 
     let mut body = collect_body(response.take_body()).await;
 
-    if let Some(head_end) = body.find("</head>") {
-        info!({ position = head_end }, "inserting script tag");
-        body.insert_str(head_end, RELOAD_ON_CHANGES);
+    if let Some(rewritten) = inject_script_tag(&body) {
+        body = rewritten;
     }
 
     Ok(response.map_body(|_, _| ResponseBody::Other(body.into())))
 }
 
+fn inject_script_tag(body: &str) -> Option<String> {
+    lol_html::rewrite_str(&body, RewriteStrSettings {
+        element_content_handlers: vec![element!("head", |head| {
+            info!("inserting script tag");
+            head.append("<script>", ContentType::Html);
+            head.append(RELOAD_ON_CHANGES, ContentType::Text);
+            head.append("</script>", ContentType::Html);
+            Ok(())
+        })],
+        ..Default::default()
+    })
+    .ok()
+}
+
 async fn collect_body(
     mut body: impl Stream<Item = Result<web::Bytes, WebError>> + Unpin,
 ) -> String {
@@ -46,8 +64,43 @@ async fn collect_body(
     String::from_utf8(buf).unwrap()
 }
 
-const RELOAD_ON_CHANGES: &str = concat!(
-    "<script>",
-    include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/reloadOnChanges.js")),
-    "</script>"
-);
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! assert_injected {
+        ($s:expr) => {{
+            let base = $s;
+            let injected = inject_script_tag(base).unwrap();
+            assert!(!base.contains(CHANGES_URL), "original shouldn't have reload URL");
+            assert!(injected.contains(CHANGES_URL), "must have injected reload URL");
+        }};
+    }
+
+    macro_rules! assert_untouched {
+        ($s:expr) => {{
+            let base = $s;
+            assert_eq!(inject_script_tag(base).unwrap(), base, "contents shouldn't change");
+        }};
+    }
+
+    #[test]
+    fn inject() {
+        assert_injected!(include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/../index.html")));
+    }
+
+    #[test]
+    fn inject_weird_whitespace() {
+        assert_injected!("<html> <head> </ head > </html>");
+    }
+
+    #[test]
+    fn inject_toml_nop() {
+        assert_untouched!(include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/../Cargo.toml")));
+    }
+
+    #[test]
+    fn inject_parse_fail_nop() {
+        assert_untouched!("<html> <head </html>");
+    }
+}