From 50629c4befa79d580a47eef79b64c673b4d5f582 Mon Sep 17 00:00:00 2001 From: gghez <6266247+gghez@users.noreply.github.com> Date: Sun, 21 Jun 2026 23:54:43 +0200 Subject: [PATCH] Expose ammonia's url_relative policy via url_relative kwarg Adds a `url_relative` keyword argument to `clean()` and `Cleaner()` that maps to ammonia's `UrlRelative` enum, so relative URLs in href/src/ can be passed through (default), denied, rewritten against a base or root URL, or rewritten by a custom callback. Accepted values: - None (default) / "pass_through" -> PassThrough - "deny" -> Deny - ("rewrite_with_base", base_url) -> RewriteWithBase - ("rewrite_with_root", root_url, path) -> RewriteWithRoot - callable (url) -> str | None -> Custom Input is validated eagerly at construction (ValueError for bad mode/URL/tuple, TypeError for unsupported types). A custom callback that raises or returns a non-str/non-None value strips the URL and reports the error via sys.unraisablehook, keeping clean() infallible. Closes #129 --- nh3.pyi | 21 ++++- src/lib.rs | 193 +++++++++++++++++++++++++++++++++++++++++++++- tests/test_nh3.py | 97 +++++++++++++++++++++++ 3 files changed, 308 insertions(+), 3 deletions(-) diff --git a/nh3.pyi b/nh3.pyi index 83e28c9..c06bfdb 100644 --- a/nh3.pyi +++ b/nh3.pyi @@ -1,10 +1,27 @@ -from typing import AbstractSet, Callable, Dict, Mapping, Optional, Set +from typing import ( + AbstractSet, + Callable, + Dict, + Literal, + Mapping, + Optional, + Set, + Tuple, + Union, +) ALLOWED_TAGS: Set[str] ALLOWED_ATTRIBUTES: Dict[str, Set[str]] ALLOWED_URL_SCHEMES: Set[str] CLEAN_CONTENT_TAGS: Set[str] +UrlRelative = Union[ + Literal["pass_through", "deny"], + Tuple[Literal["rewrite_with_base"], str], + Tuple[Literal["rewrite_with_root"], str, str], + Callable[[str], Optional[str]], +] + class Cleaner: def __init__( self, @@ -20,6 +37,7 @@ class Cleaner: url_schemes: Optional[AbstractSet[str]] = None, allowed_classes: Optional[Mapping[str, AbstractSet[str]]] = None, filter_style_properties: Optional[AbstractSet[str]] = None, + url_relative: Optional[UrlRelative] = None, ) -> None: ... def clean(self, html: str) -> str: ... @@ -37,6 +55,7 @@ def clean( url_schemes: Optional[AbstractSet[str]] = None, allowed_classes: Optional[Mapping[str, AbstractSet[str]]] = None, filter_style_properties: Optional[AbstractSet[str]] = None, + url_relative: Optional[UrlRelative] = None, ) -> str: ... def clean_text(html: str, tags: Optional[AbstractSet[str]] = None) -> str: ... def escape(html: str, tags: Optional[AbstractSet[str]] = None) -> str: ... diff --git a/src/lib.rs b/src/lib.rs index b9077b2..a4faf84 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,19 @@ use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyString, PyTuple}; +/// Internal representation of the parsed `url_relative` keyword argument. +/// +/// Parsing and validation happen eagerly when the `Cleaner` is constructed; this +/// enum is the validated result that gets converted to `ammonia::UrlRelative` in +/// `build_ammonia_from_config`. +enum UrlRelativeConfig { + PassThrough, + Deny, + RewriteWithBase(ammonia::Url), + RewriteWithRoot { root: ammonia::Url, path: String }, + Custom(Py), +} + struct Config { tags: Option>, clean_content_tags: Option>, @@ -19,6 +32,7 @@ struct Config { url_schemes: Option>, allowed_classes: Option>>, filter_style_properties: Option>, + url_relative: Option, } impl Default for Config { @@ -36,10 +50,85 @@ impl Default for Config { url_schemes: None, allowed_classes: None, filter_style_properties: None, + url_relative: None, } } } +/// Parse the Python `url_relative` argument into a validated [`UrlRelativeConfig`]. +/// +/// Accepts the strings ``"pass_through"`` / ``"deny"``, the tuples +/// ``("rewrite_with_base", base_url)`` / ``("rewrite_with_root", root_url, path)``, +/// or a callable. Any other value raises ``ValueError`` (bad mode / unparseable +/// URL / malformed tuple) or ``TypeError`` (unsupported type). +fn parse_url_relative(obj: &Bound<'_, PyAny>) -> PyResult { + if obj.cast::().is_ok() { + let s: String = obj.extract()?; + return match s.as_str() { + "pass_through" => Ok(UrlRelativeConfig::PassThrough), + "deny" => Ok(UrlRelativeConfig::Deny), + other => Err(PyValueError::new_err(format!( + "invalid url_relative string {other:?}; expected \"pass_through\" or \"deny\"" + ))), + }; + } + if let Ok(tuple) = obj.cast::() { + let mode: String = tuple + .get_item(0) + .map_err(|_| PyValueError::new_err("url_relative tuple must not be empty"))? + .extract() + .map_err(|_| PyValueError::new_err("url_relative tuple mode must be a string"))?; + return match mode.as_str() { + "rewrite_with_base" => { + if tuple.len() != 2 { + return Err(PyValueError::new_err( + "url_relative (\"rewrite_with_base\", base_url) expects exactly 2 elements", + )); + } + let base: String = tuple.get_item(1)?.extract().map_err(|_| { + PyValueError::new_err( + "url_relative rewrite_with_base base_url must be a string", + ) + })?; + let url = ammonia::Url::parse(&base).map_err(|e| { + PyValueError::new_err(format!("invalid url_relative base URL {base:?}: {e}")) + })?; + Ok(UrlRelativeConfig::RewriteWithBase(url)) + } + "rewrite_with_root" => { + if tuple.len() != 3 { + return Err(PyValueError::new_err( + "url_relative (\"rewrite_with_root\", root_url, path) expects exactly 3 elements", + )); + } + let root_url: String = tuple.get_item(1)?.extract().map_err(|_| { + PyValueError::new_err( + "url_relative rewrite_with_root root_url must be a string", + ) + })?; + let path: String = tuple.get_item(2)?.extract().map_err(|_| { + PyValueError::new_err("url_relative rewrite_with_root path must be a string") + })?; + let root = ammonia::Url::parse(&root_url).map_err(|e| { + PyValueError::new_err(format!( + "invalid url_relative root URL {root_url:?}: {e}" + )) + })?; + Ok(UrlRelativeConfig::RewriteWithRoot { root, path }) + } + other => Err(PyValueError::new_err(format!( + "invalid url_relative mode {other:?}; expected \"rewrite_with_base\" or \"rewrite_with_root\"" + ))), + }; + } + if obj.is_callable() { + return Ok(UrlRelativeConfig::Custom(obj.clone().unbind())); + } + Err(PyTypeError::new_err( + "url_relative must be a string, a tuple, or a callable", + )) +} + #[self_referencing] struct Inner { config: Config, @@ -101,6 +190,19 @@ struct Inner { /// invalid declarations and @rules will be removed, with only syntactically valid /// declarations kept. /// :type filter_style_properties: ``set[str]``, optional +/// :param url_relative: Configures how relative URLs (in ``href`` / ``src`` / +/// ````) are handled. Defaults to ``None`` (pass relative +/// URLs through unchanged). Accepted values: +/// +/// - ``"pass_through"``: keep relative URLs unchanged (explicit default). +/// - ``"deny"``: strip relative URLs entirely. +/// - ``("rewrite_with_base", base_url)``: resolve relative URLs against ``base_url``. +/// - ``("rewrite_with_root", root_url, path)``: force paths into a directory. +/// - a callable ``(url) -> str | None``: rewrite relative URLs; return a +/// string to replace, or ``None`` to strip. A callback that raises (or +/// returns a non-string, non-``None`` value) strips the URL, and the error +/// is reported via ``sys.unraisablehook``. +/// :type url_relative: ``str | tuple | Callable[[str], str | None]``, optional /// /// Example usage: /// @@ -262,6 +364,60 @@ impl Cleaner { .collect(), ); } + if let Some(url_relative) = config.url_relative.as_ref() { + let value = match url_relative { + UrlRelativeConfig::PassThrough => ammonia::UrlRelative::PassThrough, + UrlRelativeConfig::Deny => ammonia::UrlRelative::Deny, + UrlRelativeConfig::RewriteWithBase(url) => { + ammonia::UrlRelative::RewriteWithBase(url.clone()) + } + UrlRelativeConfig::RewriteWithRoot { root, path } => { + ammonia::UrlRelative::RewriteWithRoot { + root: root.clone(), + path: path.clone(), + } + } + UrlRelativeConfig::Custom(callback) => { + let callback = Python::attach(|py| callback.clone_ref(py)); + // Help the compiler infer the higher-ranked `Fn` bound that + // `UrlRelative::Custom` requires: the closure only ever returns + // owned/None values, so without this it cannot tie the output + // lifetime to the input `&str`. + fn constrain(f: F) -> F + where + F: for<'a> Fn(&'a str) -> Option> + Send + Sync + 'static, + { + f + } + let evaluate = constrain(move |url: &str| { + Python::attach(|py| { + let res = callback.call1(py, (url,)); + let err = match res { + Ok(val) => { + if val.is_none(py) { + return None; + } + match val.extract::(py) { + Ok(s) => return Some(Cow::Owned(s)), + Err(_) => PyTypeError::new_err( + "expected url_relative callback to return str or None", + ), + } + } + Err(err) => err, + }; + // A failing or mistyped callback strips the URL, keeping + // clean() infallible (unlike attribute_filter, which + // preserves the original value on error). + err.write_unraisable(py, None); + None + }) + }); + ammonia::UrlRelative::Custom(Box::new(evaluate)) + } + }; + builder.url_relative(value); + } builder } @@ -287,7 +443,8 @@ impl Cleaner { set_tag_attribute_values = None, url_schemes = None, allowed_classes = None, - filter_style_properties = None + filter_style_properties = None, + url_relative = None ))] fn py_new( py: Python, @@ -303,12 +460,17 @@ impl Cleaner { url_schemes: Option>, allowed_classes: Option>>, filter_style_properties: Option>, + url_relative: Option>, ) -> PyResult { if let Some(callback) = attribute_filter.as_ref() { if !callback.bind(py).is_callable() { return Err(PyTypeError::new_err("attribute_filter must be callable")); } } + let url_relative = match url_relative { + Some(obj) => Some(parse_url_relative(obj.bind(py))?), + None => None, + }; if link_rel.is_some() { if let Some(ref attrs) = attributes { for (tag, attr_set) in attrs.iter() { @@ -359,6 +521,7 @@ impl Cleaner { url_schemes, allowed_classes, filter_style_properties, + url_relative, }; Ok(Self::new(config)) } @@ -492,6 +655,29 @@ impl Cleaner { /// >>> nh3.clean("", /// ... link_rel=None, attributes=attributes) /// '' +/// +/// ``url_relative`` controls how relative URLs are handled. ``"deny"`` strips +/// them, while ``("rewrite_with_base", base)`` resolves them against a base URL: +/// +/// .. code-block:: pycon +/// +/// >>> nh3.clean('x', url_relative="deny") +/// 'x' +/// >>> nh3.clean( +/// ... 'x', +/// ... url_relative=("rewrite_with_base", "https://example.com"), +/// ... ) +/// 'x' +/// +/// A callable rewrites relative URLs (return ``None`` to strip): +/// +/// .. code-block:: pycon +/// +/// >>> nh3.clean( +/// ... '', +/// ... url_relative=lambda url: f"https://cdn.example.com{url}", +/// ... ) +/// '' #[pyfunction(signature = ( html, @@ -506,7 +692,8 @@ impl Cleaner { set_tag_attribute_values = None, url_schemes = None, allowed_classes = None, - filter_style_properties = None + filter_style_properties = None, + url_relative = None ))] #[allow(clippy::too_many_arguments)] fn clean( @@ -524,6 +711,7 @@ fn clean( url_schemes: Option>, allowed_classes: Option>>, filter_style_properties: Option>, + url_relative: Option>, ) -> PyResult { let cleaner = Cleaner::py_new( py, @@ -539,6 +727,7 @@ fn clean( url_schemes, allowed_classes, filter_style_properties, + url_relative, )?; Ok(py.detach(|| cleaner.clean(html))) } diff --git a/tests/test_nh3.py b/tests/test_nh3.py index e07c333..352ef76 100644 --- a/tests/test_nh3.py +++ b/tests/test_nh3.py @@ -217,6 +217,103 @@ def test_cleaner_frozenset_args(): assert cleaner.clean("hi") == 'hi' +def test_clean_url_relative_pass_through_is_default(): + html = 'x' + # Omitting url_relative keeps relative URLs (ammonia default), and the + # explicit "pass_through" string must behave identically. + assert nh3.clean(html) == 'x' + assert nh3.clean(html, url_relative="pass_through") == nh3.clean(html) + + +def test_clean_url_relative_deny(): + # Relative URLs are stripped, absolute URLs are kept. + assert ( + nh3.clean('x', url_relative="deny") + == 'x' + ) + assert ( + nh3.clean('x', url_relative="deny") + == 'x' + ) + + +def test_clean_url_relative_rewrite_with_base(): + assert ( + nh3.clean( + 'x', + url_relative=("rewrite_with_base", "https://example.com"), + ) + == 'x' + ) + + +def test_clean_url_relative_rewrite_with_root(): + out = nh3.clean( + 'x', + url_relative=( + "rewrite_with_root", + "https://github.com/rust-ammonia/ammonia/blob/master/", + "README.md", + ), + ) + assert ( + 'href="https://github.com/rust-ammonia/ammonia/blob/master/CONTRIBUTING.md"' + in out + ) + + +def test_clean_url_relative_custom_replace(): + def rewrite(url): + return f"https://cdn.example.com{url}" if url.startswith("/") else None + + assert ( + nh3.clean('', url_relative=rewrite) + == '' + ) + + +def test_clean_url_relative_custom_strip_on_none(): + assert ( + nh3.clean('y', url_relative=lambda _url: None) + == 'y' + ) + + +def test_clean_url_relative_custom_exception_strips(): + def boom(_url): + raise RuntimeError("nope") + + # A failing callback strips the URL; clean() itself stays infallible. The + # callback error is reported via sys.unraisablehook (surfaced by pytest as a + # PytestUnraisableExceptionWarning), mirroring attribute_filter's behaviour. + assert ( + nh3.clean('y', url_relative=boom) + == 'y' + ) + + +def test_clean_url_relative_invalid(): + with pytest.raises(ValueError): + nh3.clean("x", url_relative="bogus") + with pytest.raises(ValueError): + nh3.clean("x", url_relative=("bogus_mode", "https://example.com")) + with pytest.raises(ValueError): + nh3.clean("x", url_relative=("rewrite_with_base", "not a url")) + with pytest.raises(ValueError): + nh3.clean("x", url_relative=("rewrite_with_base",)) + with pytest.raises(TypeError): + nh3.clean("x", url_relative=123) + + +def test_cleaner_url_relative_reusable(): + cleaner = nh3.Cleaner(url_relative="deny") + assert cleaner.clean('x') == 'x' + assert ( + cleaner.clean('y') + == 'y' + ) + + def test_is_html(): assert not nh3.is_html("plain text") assert nh3.is_html("

html!

")