webui_dev_server/
inject.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT license.
3
4//! HTML script-injection helper.
5//!
6//! Inserts a `<script>` snippet immediately before the closing `</body>`
7//! tag (case-insensitive). If no closing tag is found, appends to the end.
8
9/// Inject `script` immediately before the closing `</body>` tag in `html`.
10///
11/// The search is case-insensitive (`</BODY>` and `</body>` both match).
12/// If `html` has no closing body tag, the script is appended.
13#[must_use]
14pub fn inject_before_body_close(html: &str, script: &str) -> String {
15    if let Some(idx) = find_close_body(html) {
16        let mut out = String::with_capacity(html.len() + script.len() + 2);
17        out.push_str(&html[..idx]);
18        out.push_str(script);
19        out.push_str(&html[idx..]);
20        out
21    } else {
22        let mut out = String::with_capacity(html.len() + script.len());
23        out.push_str(html);
24        out.push_str(script);
25        out
26    }
27}
28
29/// Case-insensitive search for `</body>`. Returns the byte offset of the
30/// `<` character, or `None` if not present. `</body>` is ASCII-only, so a
31/// fixed-window byte-slice scan is sound (no UTF-8 boundary concerns) and
32/// gives us a simple `eq_ignore_ascii_case` per window without allocating.
33fn find_close_body(html: &str) -> Option<usize> {
34    html.as_bytes()
35        .windows(7)
36        .position(|w| w.eq_ignore_ascii_case(b"</body>"))
37}
38
39#[cfg(test)]
40#[allow(clippy::disallowed_methods)]
41mod tests {
42    use super::*;
43
44    #[test]
45    fn inject_before_close_body() {
46        let html = "<html><body><p>hi</p></body></html>";
47        let injected = inject_before_body_close(html, "<script>x</script>");
48        let close_idx = injected.find("</body>").unwrap();
49        let script_idx = injected.find("<script>").unwrap();
50        assert!(script_idx < close_idx);
51        assert!(injected.starts_with("<html><body><p>hi</p>"));
52        assert!(injected.ends_with("</body></html>"));
53    }
54
55    #[test]
56    fn inject_handles_missing_close_body() {
57        let html = "<h1>hi</h1>";
58        let injected = inject_before_body_close(html, "<script>x</script>");
59        assert!(injected.starts_with(html));
60        assert!(injected.ends_with("<script>x</script>"));
61    }
62
63    #[test]
64    fn inject_case_insensitive_close_body() {
65        let html = "<HTML><BODY>x</BODY></HTML>";
66        let injected = inject_before_body_close(html, "<script>x</script>");
67        let script_idx = injected.find("<script>").unwrap();
68        let close_idx = injected.find("</BODY>").unwrap();
69        assert!(script_idx < close_idx);
70    }
71
72    #[test]
73    fn inject_uses_last_close_body_when_multiple() {
74        // `</body>` appears once in the document; this test ensures we
75        // pick the first match (which is also the only match in valid HTML).
76        let html = "<body>a</body>b";
77        let injected = inject_before_body_close(html, "S");
78        assert_eq!(injected, "<body>aS</body>b");
79    }
80}
webui_dev_server/inject.rs

webui_dev_server/
inject.rs
