Add support for image and link URL rewriting

kivikakk · Oct 31, 2024 · 4d37cc6 · 4d37cc6
1 parent b27a3dd
commit 4d37cc6
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 3 deletions.
diff --git a/src/html.rs b/src/html.rs
@@ -864,7 +864,11 @@ impl<'o, 'c: 'o> HtmlFormatter<'o, 'c> {
                         self.output.write_all(b" href=\"")?;
                         let url = nl.url.as_bytes();
                         if self.options.render.unsafe_ || !dangerous_url(url) {
-                            self.escape_href(url)?;
+                            if let Some(rewriter) = &self.options.extension.link_url_rewriter {
+                                self.escape_href(rewriter.0.lock().unwrap()(url).as_bytes())?;
+                            } else {
+                                self.escape_href(url)?;
+                            }
                         }
                         if !nl.title.is_empty() {
                             self.output.write_all(b"\" title=\"")?;
@@ -889,7 +893,11 @@ impl<'o, 'c: 'o> HtmlFormatter<'o, 'c> {
                     self.output.write_all(b" src=\"")?;
                     let url = nl.url.as_bytes();
                     if self.options.render.unsafe_ || !dangerous_url(url) {
-                        self.escape_href(url)?;
+                        if let Some(rewriter) = &self.options.extension.image_url_rewriter {
+                            self.escape_href(rewriter.0.lock().unwrap()(url).as_bytes())?;
+                        } else {
+                            self.escape_href(url)?;
+                        }
                     }
                     self.output.write_all(b"\" alt=\"")?;
                     return Ok(true);

diff --git a/src/lib.rs b/src/lib.rs
@@ -92,7 +92,7 @@ pub use parser::{
     parse_document, BrokenLinkCallback, BrokenLinkReference, ExtensionOptions,
     ExtensionOptionsBuilder, ListStyleType, Options, ParseOptions, ParseOptionsBuilder, Plugins,
     PluginsBuilder, RenderOptions, RenderOptionsBuilder, RenderPlugins, RenderPluginsBuilder,
-    ResolvedReference,
+    ResolvedReference, URLRewriterCallback,
 };
 pub use typed_arena::Arena;
 pub use xml::format_document as format_xml;

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -152,6 +152,16 @@ pub struct Options<'c> {
     pub render: RenderOptions,
 }
 
+#[derive(Clone)]
+/// Callback type for image and link rewrite extensions.
+pub struct URLRewriterCallback(pub Arc<Mutex<dyn Fn(&[u8]) -> String>>);
+
+impl Debug for URLRewriterCallback {
+    fn fmt(&self, formatter: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+        formatter.write_str("<callback>")
+    }
+}
+
 #[non_exhaustive]
 #[derive(Default, Debug, Clone, Builder)]
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
@@ -521,6 +531,40 @@ pub struct ExtensionOptions {
     /// ```
     #[builder(default)]
     pub greentext: bool,
+
+    /// Wraps embedded image URLs using a custom function.
+    ///
+    /// ```
+    /// # use ::core::str;
+    /// # use std::sync::{Arc, Mutex};
+    /// # use comrak::{markdown_to_html, ComrakOptions, URLRewriterCallback};
+    /// let mut options = ComrakOptions::default();
+    ///
+    /// let callback = |s: &[u8]| format!("https://safe.example.com?url={}", str::from_utf8(s).unwrap());
+    /// options.extension.image_url_rewriter = Some(URLRewriterCallback(Arc::new(Mutex::new(callback))));
+    ///
+    /// assert_eq!(markdown_to_html("![](http://unsafe.example.com/bad.png)", &options),
+    ///            "<p><img src=\"https://safe.example.com?url=http://unsafe.example.com/bad.png\" alt=\"\" /></p>\n");
+    /// ```
+    #[cfg_attr(feature = "arbitrary", arbitrary(value = None))]
+    pub image_url_rewriter: Option<URLRewriterCallback>,
+
+    /// Wraps link URLs using a custom function.
+    ///
+    /// ```
+    /// # use ::core::str;
+    /// # use std::sync::{Arc, Mutex};
+    /// # use comrak::{markdown_to_html, ComrakOptions, URLRewriterCallback};
+    /// let mut options = ComrakOptions::default();
+    ///
+    /// let callback = |s: &[u8]| format!("https://safe.example.com/norefer?url={}", str::from_utf8(s).unwrap());
+    /// options.extension.link_url_rewriter = Some(URLRewriterCallback(Arc::new(Mutex::new(callback))));
+    ///
+    /// assert_eq!(markdown_to_html("[my link](http://unsafe.example.com/bad)", &options),
+    ///            "<p><a href=\"https://safe.example.com/norefer?url=http://unsafe.example.com/bad\">my link</a></p>\n");
+    /// ```
+    #[cfg_attr(feature = "arbitrary", arbitrary(value = None))]
+    pub link_url_rewriter: Option<URLRewriterCallback>,
 }
 
 #[non_exhaustive]

diff --git a/src/tests.rs b/src/tests.rs
@@ -21,6 +21,7 @@ mod options;
 mod pathological;
 mod plugins;
 mod regressions;
+mod rewriter;
 mod shortcodes;
 mod spoiler;
 mod strikethrough;
@@ -120,6 +121,27 @@ fn html_opts_w(input: &str, expected: &str, roundtrip: bool, options: &Options)
     );
 }
 
+#[track_caller]
+fn html_opts_no_roundtrip<F>(input: &str, expected: &str, opts: F)
+where
+    F: Fn(&mut ComrakOptions),
+{
+    let mut options = ComrakOptions::default();
+    opts(&mut options);
+
+    let arena = Arena::new();
+
+    let root = parse_document(&arena, input, &options);
+    let mut output = vec![];
+    html::format_document(root, &options, &mut output).unwrap();
+    compare_strs(
+        &String::from_utf8(output).unwrap(),
+        expected,
+        "regular",
+        input,
+    );
+}
+
 macro_rules! html_opts {
     ([$($optclass:ident.$optname:ident),*], $lhs:expr, $rhs:expr) => {
         html_opts!([$($optclass.$optname),*], $lhs, $rhs,)

diff --git a/src/tests/rewriter.rs b/src/tests/rewriter.rs
@@ -0,0 +1,28 @@
+use ::core::str;
+use std::sync::{Arc, Mutex};
+
+use super::*;
+
+#[test]
+fn image_url_rewriter() {
+    html_opts_no_roundtrip(
+        "![](http://unsafe.example.com/bad.png)",
+        "<p><img src=\"https://safe.example.com?url=http://unsafe.example.com/bad.png\" alt=\"\" /></p>\n",
+        |opts| {
+            let callback = |s: &[u8]| format!("https://safe.example.com?url={}", str::from_utf8(s).unwrap());
+            opts.extension.image_url_rewriter = Some(URLRewriterCallback(Arc::new(Mutex::new(callback))));
+        }
+    );
+}
+
+#[test]
+fn link_url_rewriter() {
+    html_opts_no_roundtrip(
+        "[my link](http://unsafe.example.com/bad)",
+        "<p><a href=\"https://safe.example.com/norefer?url=http://unsafe.example.com/bad\">my link</a></p>\n",
+        |opts| {
+            let callback = |s: &[u8]| format!("https://safe.example.com/norefer?url={}", str::from_utf8(s).unwrap());
+            opts.extension.link_url_rewriter = Some(URLRewriterCallback(Arc::new(Mutex::new(callback))));
+        }
+    );
+}