selma 0.2.2 → 0.4.10

data/ext/selma/src/rewriter.rs

@@ -1,19 +1,31 @@
 use lol_html::{
     doc_comments, doctype, element,
     html_content::{Element, TextChunk},
-    text, DocumentContentHandlers, ElementContentHandlers, HtmlRewriter, Selector, Settings,
+    text, DocumentContentHandlers, ElementContentHandlers, HtmlRewriter, MemorySettings, Selector,
+    Settings,
 };
 use magnus::{
-    exception, function, method, scan_args,
+    exception, function, gc, method,
+    r_hash::ForEach,
+    scan_args,
     typed_data::Obj,
     value::{Opaque, ReprValue},
-    Module, Object, RArray, RModule, Ruby, Value,
+    DataTypeFunctions, Integer, IntoValue, Module, Object, RArray, RHash, RModule, Ruby, Symbol,
+    TypedData, Value,
 };
 
-use std::{borrow::Cow, cell::RefCell, primitive::str, rc::Rc};
+use std::{
+    borrow::Cow,
+    cell::{Ref, RefCell},
+    mem,
+    ops::Deref,
+    primitive::str,
+    rc::Rc,
+};
 
 use crate::{
     html::{element::SelmaHTMLElement, end_tag::SelmaHTMLEndTag, text_chunk::SelmaHTMLTextChunk},
+    native_ref_wrap::NativeRefWrap,
     sanitizer::SelmaSanitizer,
     selector::SelmaSelector,
     tags::Tag,
@@ -30,16 +42,34 @@ pub struct Handler {
     // total_elapsed_text_handlers: f64,
 }
 
+struct RewriterOptions {
+    memory_options: MemorySettings,
+}
+
 pub struct Rewriter {
     sanitizer: Option<SelmaSanitizer>,
     handlers: Vec<Handler>,
+    options: RewriterOptions,
     // total_elapsed: f64,
 }
 
-#[magnus::wrap(class = "Selma::Rewriter")]
+#[derive(TypedData)]
+#[magnus(class = "Selma::Rewriter", free_immediately, mark)]
 pub struct SelmaRewriter(std::cell::RefCell<Rewriter>);
 
-type RewriterValues = (Option<Option<Obj<SelmaSanitizer>>>, Option<RArray>);
+impl DataTypeFunctions for SelmaRewriter {
+    fn mark(&self, marker: &gc::Marker) {
+        self.0.borrow().handlers.iter().for_each(|handler| {
+            marker.mark(handler.rb_handler);
+        });
+    }
+}
+
+type RewriterValues = (
+    Option<Option<Obj<SelmaSanitizer>>>,
+    Option<RArray>,
+    Option<RHash>,
+);
 
 impl SelmaRewriter {
     const SELMA_ON_END_TAG: &'static str = "on_end_tag";
@@ -50,25 +80,20 @@ impl SelmaRewriter {
     /// @def new(sanitizer: Selma::Sanitizer.new(Selma::Sanitizer::Config::DEFAULT), handlers: [])
     /// @param sanitizer [Selma::Sanitizer] The sanitizer which performs the initial cleanup
     /// @param handlers  [Array<Selma::Selector>] The handlers to use to perform HTML rewriting
+    /// @param options  [Hash] Any additional options to pass to the rewriter
     /// @return [Selma::Rewriter]
     fn new(args: &[Value]) -> Result<Self, magnus::Error> {
-        let (rb_sanitizer, rb_handlers) = Self::scan_parse_args(args)?;
+        let (rb_sanitizer, rb_handlers, rb_options) = Self::scan_parse_args(args)?;
 
         let sanitizer = match rb_sanitizer {
             None => {
-                // no `sanitizer:` provided, use default
+                // no `sanitizer:` kwarg provided, use default
                 let default_sanitizer = SelmaSanitizer::new(&[])?;
                 let wrapped_sanitizer = Obj::wrap(default_sanitizer);
-                wrapped_sanitizer.funcall::<&str, (), Value>("setup", ())?;
-                Some(wrapped_sanitizer.get().to_owned())
+                // wrapped_sanitizer.funcall::<&str, (), Value>("setup", ())?;
+                Some(wrapped_sanitizer.deref().to_owned())
             }
-            Some(sanitizer_value) => match sanitizer_value {
-                None => None, // no `sanitizer:` provided, use default
-                Some(sanitizer) => {
-                    sanitizer.funcall::<&str, (), Value>("setup", ())?;
-                    Some(sanitizer.get().to_owned())
-                }
-            },
+            Some(sanitizer_value) => sanitizer_value.map(|sanitizer| sanitizer.deref().to_owned()),
         };
 
         let handlers = match rb_handlers {
@@ -76,9 +101,7 @@ impl SelmaRewriter {
             Some(rb_handlers) => {
                 let mut handlers: Vec<Handler> = vec![];
 
-                for h in rb_handlers.each() {
-                    let rb_handler = h.unwrap();
-
+                for rb_handler in rb_handlers.into_iter() {
                     // prevents missing #selector from ruining things
                     if !rb_handler.respond_to("selector", true).unwrap() {
                         let classname = unsafe { rb_handler.classname() };
@@ -122,9 +145,88 @@ impl SelmaRewriter {
             ));
         }
 
+        let mut rewriter_options = RewriterOptions::new();
+
+        match rb_options {
+            None => {}
+            Some(options) => {
+                options.foreach(|key: Symbol, value: RHash| {
+                    let key = key.to_string();
+                    match key.as_str() {
+                        "memory" => {
+                            let max_allowed_memory_usage = value.get(Symbol::new("max_allowed_memory_usage"));
+                            if max_allowed_memory_usage.is_some() {
+                                let max_allowed_memory_usage = max_allowed_memory_usage.unwrap();
+                                let max_allowed_memory_usage =
+                                    Integer::from_value(max_allowed_memory_usage);
+                                if max_allowed_memory_usage.is_some() {
+                                    match max_allowed_memory_usage.unwrap().to_u64() {
+                                        Ok(max_allowed_memory_usage) => {
+                                            rewriter_options.memory_options.max_allowed_memory_usage =
+                                                max_allowed_memory_usage as usize;
+                                        }
+                                        Err(_e) => {
+                                            return Err(magnus::Error::new(
+                                                exception::arg_error(),
+                                                "max_allowed_memory_usage must be a positive integer",
+                                            ));
+                                        }
+                                    }
+                                } else {
+                                    rewriter_options.memory_options.max_allowed_memory_usage = MemorySettings::default().max_allowed_memory_usage;
+                                }
+                            }
+
+                            let preallocated_parsing_buffer_size = value.get(Symbol::new("preallocated_parsing_buffer_size"));
+                            if preallocated_parsing_buffer_size.is_some() {
+                                let preallocated_parsing_buffer_size = preallocated_parsing_buffer_size.unwrap();
+                                let preallocated_parsing_buffer_size =
+                                    Integer::from_value(preallocated_parsing_buffer_size);
+                                if preallocated_parsing_buffer_size.is_some() {
+                                    match preallocated_parsing_buffer_size.unwrap().to_u64() {
+                                        Ok(preallocated_parsing_buffer_size) => {
+                                            rewriter_options.memory_options.preallocated_parsing_buffer_size =
+                                                preallocated_parsing_buffer_size as usize;
+                                        }
+                                        Err(_e) => {
+                                            return Err(magnus::Error::new(
+                                                exception::arg_error(),
+                                                "preallocated_parsing_buffer_size must be a positive integer",
+                                            ));
+                                        }
+                                    }
+                                } else {
+                                    rewriter_options.memory_options.preallocated_parsing_buffer_size = MemorySettings::default().preallocated_parsing_buffer_size;
+                                }
+                            }
+                        }
+                        _ => {
+                            return Err(magnus::Error::new(
+                                exception::arg_error(),
+                                format!("Unknown option: {key:?}"),
+                            ));
+                        }
+                    }
+                    Ok(ForEach::Continue)
+                })?;
+            }
+        }
+
+        if rewriter_options
+            .memory_options
+            .preallocated_parsing_buffer_size
+            > rewriter_options.memory_options.max_allowed_memory_usage
+        {
+            return Err(magnus::Error::new(
+                exception::arg_error(),
+                "max_allowed_memory_usage must be greater than preallocated_parsing_buffer_size",
+            ));
+        }
+
         Ok(Self(std::cell::RefCell::new(Rewriter {
             sanitizer,
             handlers,
+            options: rewriter_options,
             // total_elapsed: 0.0,
         })))
     }
@@ -141,125 +243,126 @@ impl SelmaRewriter {
         let kwargs = scan_args::get_kwargs::<
             _,
             (),
-            (Option<Option<Obj<SelmaSanitizer>>>, Option<RArray>),
+            (
+                Option<Option<Obj<SelmaSanitizer>>>,
+                Option<RArray>,
+                Option<RHash>,
+            ),
             (),
-        >(args.keywords, &[], &["sanitizer", "handlers"])?;
-        let (rb_sanitizer, rb_handlers) = kwargs.optional;
+        >(args.keywords, &[], &["sanitizer", "handlers", "options"])?;
+        let (rb_sanitizer, rb_handlers, rb_options) = kwargs.optional;
 
-        Ok((rb_sanitizer, rb_handlers))
+        Ok((rb_sanitizer, rb_handlers, rb_options))
     }
 
     /// Perform HTML rewrite sequence.
     fn rewrite(&self, html: String) -> Result<String, magnus::Error> {
-        let sanitized_html = match &self.0.borrow().sanitizer {
-            None => Ok(html),
-            Some(sanitizer) => {
-                let sanitized_html = match Self::perform_sanitization(sanitizer, &html) {
-                    Ok(sanitized_html) => sanitized_html,
-                    Err(err) => return Err(err),
-                };
+        let binding = self.0.borrow();
 
-                String::from_utf8(sanitized_html)
+        let mut sanitizer_document_content_handlers: Vec<DocumentContentHandlers> = vec![];
+        let mut sanitizer_element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> =
+            vec![];
+
+        match &binding.sanitizer {
+            None => (),
+            Some(sanitizer) => {
+                if !sanitizer.get_allow_doctype() {
+                    sanitizer_document_content_handlers.push(doctype!(|d| {
+                        sanitizer.remove_doctype(d);
+                        Ok(())
+                    }));
+                }
+                if !sanitizer.get_allow_comments() {
+                    sanitizer_document_content_handlers.push(doc_comments!(|c| {
+                        sanitizer.remove_comment(c);
+                        Ok(())
+                    }));
+                }
+                sanitizer_element_content_handlers.push(element!("*", |el| {
+                    sanitizer.try_remove_element(el);
+                    if el.removed() {
+                        return Ok(());
+                    }
+                    // if it was removed, there are no attributes to sanitize
+                    match sanitizer.sanitize_attributes(el) {
+                        Ok(_) => Ok(()),
+                        Err(err) => Err(err.to_string().into()),
+                    }
+                }));
             }
         };
-        let binding = self.0.borrow_mut();
-        let handlers = &binding.handlers;
 
-        match Self::perform_handler_rewrite(self, handlers, sanitized_html.unwrap()) {
-            Ok(rewritten_html) => Ok(String::from_utf8(rewritten_html).unwrap()),
+        let handlers: &Vec<Handler> = &binding.handlers;
+
+        match Self::perform_handler_rewrite(
+            self,
+            sanitizer_document_content_handlers,
+            sanitizer_element_content_handlers,
+            handlers,
+            html,
+        ) {
+            Ok(rewritten_html) => match &binding.sanitizer {
+                None => match String::from_utf8(rewritten_html) {
+                    Ok(output) => Ok(output),
+                    Err(err) => Err(magnus::Error::new(
+                        exception::runtime_error(),
+                        format!("{err:?}"),
+                    )),
+                },
+                Some(sanitizer) => {
+                    Self::perform_final_sanitization(self, sanitizer, rewritten_html)
+                }
+            },
             Err(err) => Err(err),
         }
     }
 
-    fn perform_sanitization(
+    // to get rid of some really nasty edge cases with dangerous tags, we perform one more
+    // sanitization pass at the end
+    fn perform_final_sanitization(
+        &self,
         sanitizer: &SelmaSanitizer,
-        html: &String,
-    ) -> Result<Vec<u8>, magnus::Error> {
-        let mut first_pass_html = vec![];
-        {
-            let mut document_content_handlers: Vec<DocumentContentHandlers> = vec![];
-            if !sanitizer.get_allow_doctype() {
-                document_content_handlers.push(doctype!(|d| {
-                    sanitizer.remove_doctype(d);
-                    Ok(())
-                }));
-            }
-            if !sanitizer.get_allow_comments() {
-                document_content_handlers.push(doc_comments!(|c| {
-                    sanitizer.remove_comment(c);
-                    Ok(())
-                }));
-            }
-            let mut rewriter = HtmlRewriter::new(
-                Settings {
-                    document_content_handlers,
-                    element_content_handlers: vec![element!("*", |el| {
-                        sanitizer.try_remove_element(el);
-                        if el.removed() {
-                            return Ok(());
-                        }
-                        match sanitizer.sanitize_attributes(el) {
-                            Ok(_) => Ok(()),
-                            Err(err) => Err(err.to_string().into()),
-                        }
-                    })],
-                    // TODO: allow for MemorySettings to be defined
-                    ..Settings::default()
-                },
-                |c: &[u8]| first_pass_html.extend_from_slice(c),
-            );
-
-            let result = rewriter.write(html.as_bytes());
-            if result.is_err() {
-                return Err(magnus::Error::new(
-                    exception::runtime_error(),
-                    format!("Failed to sanitize HTML: {}", result.unwrap_err()),
-                ));
-            }
-        }
-
-        let mut output = vec![];
-        {
-            let mut element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> = vec![];
-            if sanitizer.get_escape_tagfilter() {
-                element_content_handlers.push(element!(Tag::ESCAPEWORTHY_TAGS_CSS, |el| {
-                    let should_remove = sanitizer.allow_element(el);
-                    if should_remove {
-                        sanitizer.force_remove_element(el);
-                    }
+        html: Vec<u8>,
+    ) -> Result<String, magnus::Error> {
+        // TODO: this should ideally be done ahead of time on `initialize`, not on every `#rewrite` call
+        let mut element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> = vec![];
 
-                    Ok(())
-                }));
-            }
+        if sanitizer.get_escape_tagfilter() {
+            element_content_handlers.push(element!(Tag::ESCAPEWORTHY_TAGS_CSS, |el| {
+                let should_remove = sanitizer.allow_element(el);
+                if should_remove {
+                    sanitizer.force_remove_element(el);
+                }
 
-            let mut rewriter = HtmlRewriter::new(
-                Settings {
-                    element_content_handlers,
-                    ..Settings::default()
-                },
-                |c: &[u8]| output.extend_from_slice(c),
-            );
+                Ok(())
+            }));
+        }
 
-            let result = rewriter.write(first_pass_html.as_slice());
-            if result.is_err() {
-                return Err(magnus::Error::new(
+        match Self::run_rewrite(self, vec![], element_content_handlers, html.as_slice()) {
+            Ok(rewritten_html) => match String::from_utf8(rewritten_html) {
+                Ok(output) => Ok(output),
+                Err(err) => Err(magnus::Error::new(
                     exception::runtime_error(),
-                    format!("Failed to sanitize HTML: {}", result.unwrap_err()),
-                ));
-            }
+                    format!("{err:?}"),
+                )),
+            },
+            Err(err) => Err(err),
         }
-
-        Ok(output)
     }
 
-    pub fn perform_handler_rewrite(
+    pub fn perform_handler_rewrite<'a>(
         &self,
-        handlers: &[Handler],
+        sanitizer_document_content_handlers: Vec<DocumentContentHandlers<'a>>,
+        sanitizer_element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers<'a>)>,
+        handlers: &'a [Handler],
         html: String,
     ) -> Result<Vec<u8>, magnus::Error> {
-        // TODO: this should ideally be done ahead of time, not on every `#rewrite` call
+        // TODO: this should ideally be done ahead of time on `initialize`, not on every `#rewrite` call
         let mut element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> = vec![];
 
+        // have sanitization happen first
+        element_content_handlers.extend(sanitizer_element_content_handlers);
+
         handlers.iter().for_each(|handler| {
             let element_stack: Rc<RefCell<Vec<String>>> = Rc::new(RefCell::new(vec![]));
 
@@ -275,7 +378,7 @@ impl SelmaRewriter {
                     selector.match_element().unwrap(),
                     move |el| {
                         match Self::process_element_handlers(
-                            ruby.get_inner(handler.rb_handler),
+                            handler,
                             el,
                             &closure_element_stack.borrow(),
                         ) {
@@ -294,21 +397,19 @@ impl SelmaRewriter {
                     move |text| {
                         let element_stack = closure_element_stack.as_ref().borrow();
                         if selector.ignore_text_within().is_some() {
-                            // check if current tag is a tag we should be ignoring text within
-                            let head_tag_name = element_stack.last().unwrap().to_string();
+                            // check if current tag is a tag we should be ignoring text within;
+                            // also checks if tag is within an ancestery of ignored tags
                             if selector
                                 .ignore_text_within()
                                 .unwrap()
                                 .iter()
-                                .any(|f| f == &head_tag_name)
+                                .any(|t| element_stack.contains(t))
                             {
                                 return Ok(());
                             }
                         }
 
-                        let ruby = Ruby::get().unwrap();
-                        match Self::process_text_handlers(ruby.get_inner(handler.rb_handler), text)
-                        {
+                        match Self::process_text_handlers(handler, text) {
                             Ok(_) => Ok(()),
                             Err(err) => Err(err.to_string().into()),
                         }
@@ -329,28 +430,46 @@ impl SelmaRewriter {
 
                 let closure_element_stack = element_stack.clone();
 
-                el.end_tag_handlers()
-                    .unwrap()
-                    .push(Box::new(move |_end_tag| {
-                        let mut stack = closure_element_stack.as_ref().borrow_mut();
-                        stack.pop();
-                        Ok(())
-                    }));
+                if let Some(end_tag_handlers) = el.end_tag_handlers() {
+                    end_tag_handlers.push(lol_html::EndTagHandler::into(Box::new(
+                        move |_end_tag| {
+                            closure_element_stack.as_ref().borrow_mut().pop();
+                            Ok(())
+                        },
+                    )));
+                }
 
                 Ok(())
             }));
         });
 
+        Self::run_rewrite(
+            self,
+            sanitizer_document_content_handlers,
+            element_content_handlers,
+            html.as_bytes(),
+        )
+    }
+
+    fn run_rewrite<'a>(
+        &self,
+        document_content_handlers: Vec<DocumentContentHandlers<'a>>,
+        element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers<'a>)>,
+        html: &[u8],
+    ) -> Result<Vec<u8>, magnus::Error> {
+        let binding = &self.0.borrow();
         let mut output = vec![];
         {
             let mut rewriter = HtmlRewriter::new(
                 Settings {
+                    document_content_handlers,
                     element_content_handlers,
+                    memory_settings: Self::get_memory_options(binding),
                     ..Settings::default()
                 },
                 |c: &[u8]| output.extend_from_slice(c),
             );
-            match rewriter.write(html.as_bytes()) {
+            match rewriter.write(html) {
                 Ok(_) => {}
                 Err(err) => {
                     return Err(magnus::Error::new(
@@ -364,10 +483,12 @@ impl SelmaRewriter {
     }
 
     fn process_element_handlers(
-        rb_handler: Value,
+        handler: &Handler,
         element: &mut Element,
         ancestors: &[String],
     ) -> Result<(), magnus::Error> {
+        let rb_handler = handler.rb_handler.into_value();
+
         // if `on_end_tag` function is defined, call it
         if rb_handler.respond_to(Self::SELMA_ON_END_TAG, true).unwrap() {
             // TODO: error here is an "EndTagError"
@@ -375,38 +496,61 @@ impl SelmaRewriter {
                 .end_tag_handlers()
                 .unwrap()
                 .push(Box::new(move |end_tag| {
-                    let rb_end_tag = SelmaHTMLEndTag::new(end_tag);
+                    let (ref_wrap, anchor) = NativeRefWrap::wrap(end_tag);
+
+                    let rb_end_tag = SelmaHTMLEndTag::new(ref_wrap);
 
-                    match rb_handler.funcall::<_, _, Value>(Self::SELMA_ON_END_TAG, (rb_end_tag,)) {
+                    let result =
+                        rb_handler.funcall::<_, _, Value>(Self::SELMA_ON_END_TAG, (rb_end_tag,));
+
+                    mem::drop(anchor);
+
+                    match result {
                         Ok(_) => Ok(()),
                         Err(err) => Err(err.to_string().into()),
                     }
                 }));
         }
 
-        let rb_element = SelmaHTMLElement::new(element, ancestors);
-        let rb_result =
-            rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_ELEMENT, (rb_element,));
-        match rb_result {
+        let (ref_wrap, anchor) = NativeRefWrap::wrap(element);
+        let rb_element = SelmaHTMLElement::new(ref_wrap, ancestors);
+        let result = rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_ELEMENT, (rb_element,));
+
+        mem::drop(anchor);
+
+        match result {
             Ok(_) => Ok(()),
-            Err(err) => Err(err),
+            Err(err) => Err(magnus::Error::new(
+                exception::runtime_error(),
+                format!("{err:?}"),
+            )),
         }
     }
 
     fn process_text_handlers(
-        rb_handler: Value,
+        handler: &Handler,
         text_chunk: &mut TextChunk,
     ) -> Result<(), magnus::Error> {
+        let rb_handler = handler.rb_handler.into_value();
+
         // prevents missing `handle_text_chunk` function
         let content = text_chunk.as_str();
 
-        // seems that sometimes lol-html returns blank text / EOLs?
+        // lol-html sometimes returns blank text if
+        // last_in_text_node() is true
         if content.is_empty() {
             return Ok(());
         }
 
-        let rb_text_chunk = SelmaHTMLTextChunk::new(text_chunk);
-        match rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_TEXT_CHUNK, (rb_text_chunk,)) {
+        let (ref_wrap, anchor) = NativeRefWrap::wrap(text_chunk);
+
+        let rb_text_chunk = SelmaHTMLTextChunk::new(ref_wrap);
+        let result =
+            rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_TEXT_CHUNK, (rb_text_chunk,));
+
+        mem::drop(anchor);
+
+        match result {
             Ok(_) => Ok(()),
             Err(err) => Err(magnus::Error::new(
                 exception::runtime_error(),
@@ -414,6 +558,22 @@ impl SelmaRewriter {
             )),
         }
     }
+
+    fn get_memory_options(binding: &Ref<Rewriter>) -> MemorySettings {
+        let options = &binding.options.memory_options;
+        MemorySettings {
+            max_allowed_memory_usage: options.max_allowed_memory_usage,
+            preallocated_parsing_buffer_size: options.preallocated_parsing_buffer_size,
+        }
+    }
+}
+
+impl RewriterOptions {
+    pub fn new() -> Self {
+        Self {
+            memory_options: MemorySettings::default(),
+        }
+    }
 }
 
 pub fn init(m_selma: RModule) -> Result<(), magnus::Error> {