Allow using element! in a separate expression from rewrite_str (#69)

cloudflare · Oct 18, 2020 · e0eaf6c · e0eaf6c
1 parent 3c53085
commit e0eaf6c
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 27 deletions.
diff --git a/c-api/src/rewriter_builder.rs b/c-api/src/rewriter_builder.rs
@@ -1,5 +1,6 @@
 use super::*;
 use libc::c_void;
+use std::borrow::Cow;
 
 #[repr(C)]
 pub enum RewriterDirective {
@@ -89,7 +90,7 @@ impl ExternElementContentHandlers {
 
 pub struct SafeContentHandlers<'b> {
     pub document: Vec<DocumentContentHandlers<'b>>,
-    pub element: Vec<(&'b Selector, ElementContentHandlers<'b>)>,
+    pub element: Vec<(Cow<'b, Selector>, ElementContentHandlers<'b>)>,
 }
 
 #[derive(Default)]
@@ -109,7 +110,7 @@ impl HtmlRewriterBuilder {
             element: self
                 .element_content_handlers
                 .iter()
-                .map(|(s, h)| (*s, h.as_safe_element_content_handlers()))
+                .map(|(s, h)| (Cow::Borrowed(*s), h.as_safe_element_content_handlers()))
                 .collect(),
         }
     }

diff --git a/src/rewritable_units/mod.rs b/src/rewritable_units/mod.rs
@@ -90,6 +90,7 @@ mod test_utils {
     use crate::test_utils::{Output, ASCII_COMPATIBLE_ENCODINGS};
     use crate::*;
     use encoding_rs::Encoding;
+    use std::borrow::Cow;
 
     pub fn encoded(input: &str) -> Vec<(Vec<u8>, &'static Encoding)> {
         ASCII_COMPATIBLE_ENCODINGS
@@ -113,10 +114,10 @@ mod test_utils {
             .collect()
     }
 
-    pub fn rewrite_html(
+    pub fn rewrite_html<'a>(
         html: &[u8],
         encoding: &'static Encoding,
-        element_content_handlers: Vec<(&Selector, ElementContentHandlers)>,
+        element_content_handlers: Vec<(Cow<'a, Selector>, ElementContentHandlers)>,
         document_content_handlers: Vec<DocumentContentHandlers>,
     ) -> String {
         let mut output = Output::new(encoding);

diff --git a/src/rewriter/mod.rs b/src/rewriter/mod.rs
@@ -156,7 +156,7 @@ impl<'h, O: OutputSink> HtmlRewriter<'h, O> {
         for (selector, handlers) in settings.element_content_handlers {
             let locator = dispatcher.add_selector_associated_handlers(handlers);
 
-            selectors_ast.add_selector(selector, locator);
+            selectors_ast.add_selector(&selector, locator);
         }
 
         for handlers in settings.document_content_handlers {
@@ -251,22 +251,23 @@ impl<O: OutputSink> Debug for HtmlRewriter<'_, O> {
 /// ```
 /// use lol_html::{rewrite_str, element, RewriteStrSettings};
 ///
+/// let element_content_handlers = vec![
+///     // Rewrite insecure hyperlinks
+///     element!("a[href]", |el| {
+///         let href = el
+///             .get_attribute("href")
+///             .unwrap()
+///             .replace("http:", "https:");
+///
+///          el.set_attribute("href", &href).unwrap();
+///
+///          Ok(())
+///     })
+/// ];
 /// let output = rewrite_str(
 ///     r#"<div><a href="http://example.com"></a></div>"#,
 ///     RewriteStrSettings {
-///         element_content_handlers: vec![
-///             // Rewrite insecure hyperlinks
-///             element!("a[href]", |el| {
-///                 let href = el
-///                     .get_attribute("href")
-///                     .unwrap()
-///                     .replace("http:", "https:");
-///
-///                  el.set_attribute("href", &href).unwrap();
-///
-///                  Ok(())
-///             })
-///         ],
+///         element_content_handlers,
 ///         ..RewriteStrSettings::default()
 ///     }
 /// ).unwrap();
@@ -650,9 +651,14 @@ mod tests {
                 document_handlers: DocumentContentHandlers,
                 expected_err: &'static str,
             ) {
+                use std::borrow::Cow;
+
                 let mut rewriter = HtmlRewriter::try_new(
                     Settings {
-                        element_content_handlers: vec![(&"*".parse().unwrap(), element_handlers)],
+                        element_content_handlers: vec![(
+                            Cow::Owned("*".parse().unwrap()),
+                            element_handlers,
+                        )],
                         document_content_handlers: vec![document_handlers],
                         ..Settings::default()
                     },

diff --git a/src/rewriter/settings.rs b/src/rewriter/settings.rs
@@ -1,5 +1,6 @@
 use crate::rewritable_units::{Comment, Doctype, DocumentEnd, Element, EndTag, TextChunk};
 use crate::selectors_vm::Selector;
+use std::borrow::Cow;
 use std::error::Error;
 
 pub(super) type HandlerResult = Result<(), Box<dyn Error + Send + Sync>>;
@@ -108,7 +109,7 @@ impl<'h> DocumentContentHandlers<'h> {
 macro_rules! __element_content_handler {
     ($selector:expr, $handler_name:ident, $handler:expr) => {
         (
-            &$selector.parse::<$crate::Selector>().unwrap(),
+            ::std::borrow::Cow::Owned($selector.parse::<$crate::Selector>().unwrap()),
             $crate::ElementContentHandlers::default().$handler_name($handler),
         )
     };
@@ -417,20 +418,21 @@ pub struct Settings<'h, 's> {
     ///
     /// ### Example
     /// ```
+    /// use std::borrow::Cow;
     /// use lol_html::{ElementContentHandlers, Settings};
     ///
     /// let settings = Settings {
     ///     element_content_handlers: vec! [
     ///         (
-    ///             &"div[foo]".parse().unwrap(),
+    ///             Cow::Owned("div[foo]".parse().unwrap()),
     ///             ElementContentHandlers::default().element(|el| {
     ///                 // ...
     ///
     ///                 Ok(())
     ///             })
     ///         ),
     ///         (
-    ///             &"body".parse().unwrap(),
+    ///             Cow::Owned("body".parse().unwrap()),
     ///             ElementContentHandlers::default().comments(|c| {
     ///                 // ...
     ///
@@ -445,7 +447,7 @@ pub struct Settings<'h, 's> {
     /// [`element`]: macro.element.html
     /// [`comments`]: macro.comments.html
     /// [`text`]: macro.text.html
-    pub element_content_handlers: Vec<(&'s Selector, ElementContentHandlers<'h>)>,
+    pub element_content_handlers: Vec<(Cow<'s, Selector>, ElementContentHandlers<'h>)>,
 
     /// Specifies rewriting handlers for the content without associating it to a particular
     /// CSS selector.
@@ -552,20 +554,21 @@ pub struct RewriteStrSettings<'h, 's> {
     ///
     /// ### Example
     /// ```
+    /// use std::borrow::Cow;
     /// use lol_html::{ElementContentHandlers, RewriteStrSettings};
     ///
     /// let settings = RewriteStrSettings {
     ///     element_content_handlers: vec! [
     ///         (
-    ///             &"div[foo]".parse().unwrap(),
+    ///             Cow::Owned("div[foo]".parse().unwrap()),
     ///             ElementContentHandlers::default().element(|el| {
     ///                 // ...
     ///
     ///                 Ok(())
     ///             })
     ///         ),
     ///         (
-    ///             &"body".parse().unwrap(),
+    ///             Cow::Owned("div[foo]".parse().unwrap()),
     ///             ElementContentHandlers::default().comments(|c| {
     ///                 // ...
     ///
@@ -580,7 +583,7 @@ pub struct RewriteStrSettings<'h, 's> {
     /// [`element`]: macro.element.html
     /// [`comments`]: macro.comments.html
     /// [`text`]: macro.text.html
-    pub element_content_handlers: Vec<(&'s Selector, ElementContentHandlers<'h>)>,
+    pub element_content_handlers: Vec<(Cow<'s, Selector>, ElementContentHandlers<'h>)>,
 
     /// Specifies rewriting handlers for the content without associating it to a particular
     /// CSS selector.

diff --git a/src/selectors_vm/parser.rs b/src/selectors_vm/parser.rs
@@ -195,7 +195,7 @@ impl<'i> Parser<'i> for SelectorsParser {
 /// [`parse`]: https://doc.rust-lang.org/std/primitive.str.html#method.parse
 /// [element content handlers]: struct.Settings.html#structfield.element_content_handlers
 /// [`FromStr`]: https://doc.rust-lang.org/std/str/trait.FromStr.html
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct Selector(pub(crate) SelectorList<SelectorImplDescriptor>);
 
 impl FromStr for Selector {