fluxer/fluxer_unfurl/src/sanitizer.rs

// SPDX-License-Identifier: AGPL-3.0-or-later

use ammonia::{Builder, UrlRelative};
use std::borrow::Cow;
use std::collections::HashSet;

const TRUSTED_HOSTS: &[&str] = &[
    "youtube.com",
    "youtu.be",
    "youtube-nocookie.com",
    "vimeo.com",
    "soundcloud.com",
    "spotify.com",
    "twitter.com",
    "x.com",
    "reddit.com",
    "tiktok.com",
    "codepen.io",
];

pub fn sanitize_oembed_html_with_context(
    html: &str,
    source_url: &url::Url,
    provider_url: Option<&str>,
) -> Option<String> {
    let source_trusted = source_url.host_str().is_some_and(is_trusted_oembed_host);

    let provider_trusted = provider_url
        .and_then(|u| url::Url::parse(u).ok())
        .and_then(|u| u.host_str().map(is_trusted_oembed_host))
        .unwrap_or(false);

    if !source_trusted && !provider_trusted {
        return None;
    }

    sanitize_oembed_html(html)
}

pub fn sanitize_oembed_html(html: &str) -> Option<String> {
    let trimmed = html.trim();
    if trimmed.is_empty() {
        return None;
    }

    let mut allowed_tags: HashSet<&str> = HashSet::new();
    allowed_tags.extend([
        "blockquote",
        "div",
        "span",
        "p",
        "strong",
        "em",
        "b",
        "i",
        "br",
        "a",
        "iframe",
    ]);

    let mut tag_attributes: std::collections::HashMap<&str, HashSet<&str>> =
        std::collections::HashMap::new();

    let mut iframe_attrs: HashSet<&str> = HashSet::new();
    iframe_attrs.extend([
        "src",
        "width",
        "height",
        "allow",
        "allowfullscreen",
        "frameborder",
        "title",
        "loading",
        "referrerpolicy",
        "sandbox",
    ]);
    tag_attributes.insert("iframe", iframe_attrs);

    let mut a_attrs: HashSet<&str> = HashSet::new();
    a_attrs.extend(["href", "target"]);
    tag_attributes.insert("a", a_attrs);

    let mut bq_attrs: HashSet<&str> = HashSet::new();
    bq_attrs.insert("cite");
    tag_attributes.insert("blockquote", bq_attrs);

    let sanitized = Builder::new()
        .tags(allowed_tags)
        .tag_attributes(tag_attributes)
        .url_schemes(HashSet::from(["http", "https"]))
        .url_relative(UrlRelative::Deny)
        .attribute_filter(|element, attribute, value| {
            if element == "iframe" && attribute == "src" && !is_trusted_oembed_src(value) {
                return None;
            }
            Some(Cow::Borrowed(value))
        })
        .link_rel(Some("noopener noreferrer"))
        .clean(trimmed)
        .to_string();

    let result = sanitized.trim().to_owned();
    if result.is_empty() {
        return None;
    }

    Some(result)
}

pub fn is_trusted_oembed_host(hostname: &str) -> bool {
    let lower = hostname.to_ascii_lowercase();
    TRUSTED_HOSTS
        .iter()
        .any(|&trusted| lower == trusted || lower.ends_with(&format!(".{trusted}")))
}

fn is_trusted_oembed_src(src: &str) -> bool {
    url::Url::parse(src).ok().is_some_and(|url| {
        matches!(url.scheme(), "http" | "https")
            && url.host_str().is_some_and(is_trusted_oembed_host)
    })
}

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

    #[test]
    fn strips_script_tags() {
        let html = r#"<script>alert(1)</script><p>Hello</p>"#;
        let result = sanitize_oembed_html(html);
        assert_eq!(result, Some("<p>Hello</p>".to_owned()));
    }

    #[test]
    fn empty_input_returns_none() {
        assert_eq!(sanitize_oembed_html(""), None);
        assert_eq!(sanitize_oembed_html("   "), None);
    }

    #[test]
    fn trusted_host_check() {
        assert!(is_trusted_oembed_host("youtube.com"));
        assert!(is_trusted_oembed_host("www.youtube.com"));
        assert!(!is_trusted_oembed_host("evil.com"));
    }

    #[test]
    fn keeps_trusted_iframe_src() {
        let html = r#"<iframe src="https://www.youtube.com/embed/abc"></iframe>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(result.contains(r#"src="https://www.youtube.com/embed/abc""#));
    }

    #[test]
    fn strips_untrusted_iframe_src() {
        let html = r#"<iframe src="https://youtube.com.evil.test/embed"></iframe>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(!result.contains("youtube.com.evil.test"));
        assert!(!result.contains("src="));
    }

    #[test]
    fn strips_relative_iframe_src() {
        let html = r#"<iframe src="/embed"></iframe>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(!result.contains("src="));
    }

    #[test]
    fn keeps_basic_formatting_tags() {
        let html = r#"<p><strong>Bold</strong> <em>italic</em> <b>b</b> <i>i</i></p>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(result.contains("<strong>Bold</strong>"));
        assert!(result.contains("<em>italic</em>"));
    }

    #[test]
    fn strips_style_and_img_tags() {
        let html = r#"<style>body{}</style><img src="x"><p>Keep</p>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(!result.contains("style"));
        assert!(!result.contains("img"));
        assert!(result.contains("<p>Keep</p>"));
    }

    #[test]
    fn keeps_blockquote_with_cite() {
        let html = r#"<blockquote cite="https://e.com/post">Quote text</blockquote>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(result.contains("blockquote"));
        assert!(result.contains(r#"cite="https://e.com/post""#));
    }

    #[test]
    fn anchor_gets_rel_noopener() {
        let html = r#"<a href="https://e.com">Link</a>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(result.contains("noopener"));
    }

    #[test]
    fn iframe_keeps_sandbox_and_dimensions() {
        let html = r#"<iframe src="https://www.youtube.com/embed/abc" width="640" height="360" sandbox="allow-scripts"></iframe>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(result.contains(r#"width="640""#));
        assert!(result.contains(r#"height="360""#));
        assert!(result.contains("sandbox"));
    }

    #[test]
    fn all_trusted_hosts_are_recognised() {
        for host in [
            "youtube.com",
            "www.youtube.com",
            "youtu.be",
            "youtube-nocookie.com",
            "vimeo.com",
            "player.vimeo.com",
            "soundcloud.com",
            "w.soundcloud.com",
            "spotify.com",
            "open.spotify.com",
            "twitter.com",
            "x.com",
            "reddit.com",
            "www.reddit.com",
            "tiktok.com",
            "www.tiktok.com",
            "codepen.io",
            "cdpn.io",
        ] {
            if host == "cdpn.io" {
                assert!(
                    !is_trusted_oembed_host(host),
                    "{host} should NOT be trusted"
                );
            } else {
                assert!(is_trusted_oembed_host(host), "{host} should be trusted");
            }
        }
    }

    #[test]
    fn context_sanitizer_rejects_untrusted_source_and_provider() {
        let source = url::Url::parse("https://untrusted.com/page").unwrap();
        let html = r#"<p>Hello</p>"#;
        assert!(sanitize_oembed_html_with_context(html, &source, None).is_none());
        assert!(
            sanitize_oembed_html_with_context(html, &source, Some("https://also-untrusted.com"))
                .is_none()
        );
    }

    #[test]
    fn context_sanitizer_allows_trusted_source() {
        let source = url::Url::parse("https://www.youtube.com/watch?v=abc").unwrap();
        let html = r#"<p>Hello</p>"#;
        assert!(sanitize_oembed_html_with_context(html, &source, None).is_some());
    }

    #[test]
    fn context_sanitizer_allows_trusted_provider() {
        let source = url::Url::parse("https://untrusted.com/page").unwrap();
        let html = r#"<p>Hello</p>"#;
        assert!(
            sanitize_oembed_html_with_context(html, &source, Some("https://youtube.com")).is_some()
        );
    }

    #[test]
    fn strips_javascript_uri_in_iframe() {
        let html = r#"<iframe src="javascript:alert(1)"></iframe>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(!result.contains("javascript"));
    }

    #[test]
    fn strips_data_uri_in_iframe() {
        let html = r#"<iframe src="data:text/html,<h1>pwned</h1>"></iframe>"#;
        let result = sanitize_oembed_html(html).unwrap();
        assert!(!result.contains("data:"));
    }
}