microsoft-webui-handler 3.0.5

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

//! Minimal HTML entity encoding for safe embedding of dynamic values.
//!
//! A focused, zero-dependency implementation that covers the six characters
//! needed for XSS prevention in HTML text and attribute contexts.

use std::borrow::Cow;

use crate::{ResponseWriter, Result};

/// Characters escaped and their replacements:
///
/// * `&` → `&`
/// * `<` → `&lt;`
/// * `>` → `&gt;`
/// * `"` → `&quot;`
/// * `'` → `&#x27;`
/// * `/` → `&#x2F;`
///
/// Returns [`Cow::Borrowed`] when the input contains no characters that need
/// escaping (zero-allocation fast path).
pub fn encode_safe(input: &str) -> Cow<'_, str> {
    let bytes = input.as_bytes();

    // Fast path: find the first byte that needs escaping.
    let first = bytes.iter().position(|&byte| needs_escape(byte));

    let Some(pos) = first else {
        return Cow::Borrowed(input);
    };

    // Slow path: size the output exactly, since URL-like values grow by five
    // bytes per `/` and would otherwise reallocate partway through.
    let extra: usize = bytes[pos..].iter().map(|&b| escaped_extra_len(b)).sum();
    let mut out = String::with_capacity(input.len() + extra);
    out.push_str(&input[..pos]);

    let mut start = pos;
    for (i, &b) in bytes[pos..].iter().enumerate() {
        let Some(replacement) = escaped_replacement(b) else {
            continue;
        };
        // Flush the unescaped run before this match.
        out.push_str(&input[start..pos + i]);
        out.push_str(replacement);
        start = pos + i + 1;
    }
    // Flush any remaining unescaped tail.
    out.push_str(&input[start..]);

    Cow::Owned(out)
}

#[inline]
const fn needs_escape(byte: u8) -> bool {
    matches!(byte, b'&' | b'<' | b'>' | b'"' | b'\'' | b'/')
}

#[inline]
const fn escaped_replacement(byte: u8) -> Option<&'static str> {
    match byte {
        b'&' => Some("&amp;"),
        b'<' => Some("&lt;"),
        b'>' => Some("&gt;"),
        b'"' => Some("&quot;"),
        b'\'' => Some("&#x27;"),
        b'/' => Some("&#x2F;"),
        _ => None,
    }
}

#[inline]
fn escaped_chunks(input: &str, mut append: impl FnMut(&str)) {
    let Some(first) = input.as_bytes().iter().position(|&byte| needs_escape(byte)) else {
        append(input);
        return;
    };
    if first > 0 {
        append(&input[..first]);
    }
    let mut start = first;
    for (offset, &byte) in input.as_bytes()[first..].iter().enumerate() {
        let index = first + offset;
        if let Some(replacement) = escaped_replacement(byte) {
            if start < index {
                append(&input[start..index]);
            }
            append(replacement);
            start = index + 1;
        }
    }
    if start < input.len() {
        append(&input[start..]);
    }
}

/// Append HTML-escaped text directly to a string using [`crate::encode_safe`]'s entity rules.
#[doc(hidden)]
#[inline]
pub fn append_escaped_to_string(output: &mut String, input: &str) {
    escaped_chunks(input, |chunk| output.push_str(chunk));
}

/// Append HTML-escaped text directly to a byte buffer using [`crate::encode_safe`]'s entity rules.
#[doc(hidden)]
#[inline]
pub fn append_escaped_to_bytes(output: &mut Vec<u8>, input: &str) {
    escaped_chunks(input, |chunk| output.extend_from_slice(chunk.as_bytes()));
}

/// Bytes [`encode_safe`] adds when it replaces `b` with its entity.
#[inline]
const fn escaped_extra_len(b: u8) -> usize {
    match b {
        b'&' => 4,
        b'<' | b'>' => 3,
        b'"' | b'\'' | b'/' => 5,
        _ => 0,
    }
}

#[inline]
fn is_style_end_tag(bytes: &[u8], index: usize) -> bool {
    bytes[index] == b'<'
        && bytes[index + 1] == b'/'
        && bytes[index + 2].eq_ignore_ascii_case(&b's')
        && bytes[index + 3].eq_ignore_ascii_case(&b't')
        && bytes[index + 4].eq_ignore_ascii_case(&b'y')
        && bytes[index + 5].eq_ignore_ascii_case(&b'l')
        && bytes[index + 6].eq_ignore_ascii_case(&b'e')
}

/// Return whether CSS needs escaping before it is written into a `<style>`.
pub(crate) fn style_text_needs_escape(css: &str) -> bool {
    let bytes = css.as_bytes();
    let mut index = 0usize;
    while index + 7 <= bytes.len() {
        if is_style_end_tag(bytes, index) {
            return true;
        }
        index += 1;
    }
    false
}

/// Write CSS into an HTML `<style>` raw-text element without allowing an
/// authored, case-insensitive `</style` sequence to terminate the element.
///
/// Escaping only the slash (`<\/style`) preserves the CSS string value and
/// keeps the ordinary path allocation-free.
pub(crate) fn write_style_text(writer: &mut dyn ResponseWriter, css: &str) -> Result<()> {
    let bytes = css.as_bytes();
    let mut chunk_start = 0usize;
    let mut index = 0usize;
    while index + 7 <= bytes.len() {
        if is_style_end_tag(bytes, index) {
            writer.write(&css[chunk_start..index + 1])?;
            writer.write("\\/")?;
            chunk_start = index + 2;
            index += 7;
        } else {
            index += 1;
        }
    }
    writer.write(&css[chunk_start..])
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn escaped_append_preserves_existing_content_and_encoding() {
        let ascii: String = (0..=127).map(char::from).collect();
        for input in ["", "plain text", "&<>\"'/", "&amp;", "é水🙂/", &ascii] {
            let expected = format!("prefix{}", encode_safe(input));
            let mut text = String::from("prefix");
            let mut bytes = b"prefix".to_vec();
            append_escaped_to_string(&mut text, input);
            append_escaped_to_bytes(&mut bytes, input);
            assert_eq!(text, expected);
            assert_eq!(bytes, expected.as_bytes());
        }
    }

    #[derive(Default)]
    struct StyleWriter {
        output: String,
        writes: usize,
    }

    impl ResponseWriter for StyleWriter {
        fn write(&mut self, content: &str) -> Result<()> {
            self.output.push_str(content);
            self.writes += 1;
            Ok(())
        }

        fn end(&mut self) -> Result<()> {
            Ok(())
        }
    }

    #[test]
    fn no_escaping_returns_borrowed() {
        let input = "Hello World 123";
        let result = encode_safe(input);
        assert!(matches!(result, Cow::Borrowed(_)));
        assert_eq!(result, "Hello World 123");
    }

    #[test]
    fn empty_string_returns_borrowed() {
        let result = encode_safe("");
        assert!(matches!(result, Cow::Borrowed(_)));
        assert_eq!(result, "");
    }

    #[test]
    fn escapes_ampersand() {
        assert_eq!(encode_safe("a&b"), "a&amp;b");
    }

    #[test]
    fn escapes_less_than() {
        assert_eq!(encode_safe("<script>"), "&lt;script&gt;");
    }

    #[test]
    fn escapes_greater_than() {
        assert_eq!(encode_safe("a>b"), "a&gt;b");
    }

    #[test]
    fn escapes_double_quote() {
        assert_eq!(encode_safe(r#"a"b"#), "a&quot;b");
    }

    #[test]
    fn escapes_single_quote() {
        assert_eq!(encode_safe("a'b"), "a&#x27;b");
    }

    #[test]
    fn escapes_forward_slash() {
        assert_eq!(encode_safe("a/b"), "a&#x2F;b");
    }

    #[test]
    fn escapes_all_special_chars() {
        assert_eq!(encode_safe(r#"&<>"'/"#), "&amp;&lt;&gt;&quot;&#x27;&#x2F;");
    }

    #[test]
    fn preserves_unicode() {
        assert_eq!(encode_safe("こんにちは"), "こんにちは");
        assert!(matches!(encode_safe("こんにちは"), Cow::Borrowed(_)));
    }

    #[test]
    fn mixed_unicode_and_special() {
        assert_eq!(encode_safe("日本語&テスト"), "日本語&amp;テスト");
    }

    #[test]
    fn escapes_at_start() {
        assert_eq!(encode_safe("&start"), "&amp;start");
    }

    #[test]
    fn escapes_at_end() {
        assert_eq!(encode_safe("end&"), "end&amp;");
    }

    #[test]
    fn multiple_consecutive_escapes() {
        assert_eq!(encode_safe("&&"), "&amp;&amp;");
    }

    #[test]
    fn realistic_nonce() {
        // CSP nonces are base64 — no special chars expected.
        let nonce = "YWJjZGVmZ2hpamtsbW5vcA==";
        let result = encode_safe(nonce);
        assert!(matches!(result, Cow::Borrowed(_)));
    }

    #[test]
    fn realistic_attribute_value() {
        assert_eq!(
            encode_safe("John O'Brien & Associates"),
            "John O&#x27;Brien &amp; Associates"
        );
    }

    #[test]
    fn realistic_xss_attempt() {
        assert_eq!(
            encode_safe("<script>alert('xss')</script>"),
            "&lt;script&gt;alert(&#x27;xss&#x27;)&lt;&#x2F;script&gt;"
        );
    }

    #[test]
    fn escapes_all_known_special_chars() {
        // Verify all six special characters are escaped correctly.
        let test_cases = [
            ("", ""),
            ("hello", "hello"),
            ("&", "&amp;"),
            ("<", "&lt;"),
            (">", "&gt;"),
            ("\"", "&quot;"),
            ("'", "&#x27;"),
            ("/", "&#x2F;"),
            ("a&b<c>d\"e'f/g", "a&amp;b&lt;c&gt;d&quot;e&#x27;f&#x2F;g"),
        ];
        for (input, expected) in test_cases {
            assert_eq!(encode_safe(input), expected, "Failed for input: {input:?}");
        }
    }

    #[test]
    fn style_text_preserves_normal_bytes_in_one_write() {
        let mut writer = StyleWriter::default();
        write_style_text(&mut writer, ".card{content:'normal / style'}").unwrap();
        assert_eq!(writer.output, ".card{content:'normal / style'}");
        assert_eq!(writer.writes, 1);
    }

    #[test]
    fn style_text_escapes_mixed_case_end_tags() {
        let mut writer = StyleWriter::default();
        write_style_text(
            &mut writer,
            ".a{content:'</style>'}.b{content:'</StYlE attr>'}",
        )
        .unwrap();
        assert_eq!(
            writer.output,
            ".a{content:'<\\/style>'}.b{content:'<\\/StYlE attr>'}"
        );
    }

    #[test]
    fn style_text_escape_detection_matches_writer() {
        let cases = [
            ("", false),
            ("</styl", false),
            ("< /style", false),
            ("<\\/style", false),
            ("</style", true),
            ("prefix</STYLE>suffix", true),
            ("a</StYlE attr>b", true),
        ];

        for (css, expected) in cases {
            assert_eq!(
                style_text_needs_escape(css),
                expected,
                "unexpected detection for {css:?}"
            );

            let mut writer = StyleWriter::default();
            write_style_text(&mut writer, css).unwrap();
            assert_eq!(
                writer.output != css,
                expected,
                "writer and detector diverged for {css:?}"
            );
        }
    }
}