selma 0.0.1

data/ext/selma/src/rewriter.rs

@@ -0,0 +1,441 @@
+use std::{borrow::Cow, cell::RefCell, rc::Rc};
+
+use lol_html::{
+    doc_comments, doctype, element,
+    html_content::{ContentType, Element, EndTag, TextChunk},
+    text, DocumentContentHandlers, ElementContentHandlers, HtmlRewriter, Selector, Settings,
+};
+use magnus::{exception, function, method, scan_args, Module, Object, RArray, RModule, Value};
+
+use crate::{
+    html::{element::SelmaHTMLElement, end_tag::SelmaHTMLEndTag},
+    sanitizer::SelmaSanitizer,
+    selector::SelmaSelector,
+    tags::Tag,
+    wrapped_struct::WrappedStruct,
+};
+
+#[derive(Clone, Debug)]
+pub struct Handler {
+    rb_handler: Value,
+    rb_selector: WrappedStruct<SelmaSelector>,
+
+    total_element_handler_calls: usize,
+    total_elapsed_element_handlers: f64,
+
+    total_text_handler_calls: usize,
+    total_elapsed_text_handlers: f64,
+}
+
+pub struct Rewriter {
+    sanitizer: Option<SelmaSanitizer>,
+    handlers: Vec<Handler>,
+
+    total_elapsed: f64,
+}
+
+#[magnus::wrap(class = "Selma::Rewriter")]
+pub struct SelmaRewriter(std::cell::RefCell<Rewriter>);
+
+/// SAFETY: This is safe because we only access this data when the GVL is held.
+unsafe impl Send for SelmaRewriter {}
+
+impl SelmaRewriter {
+    const SELMA_ON_END_TAG: &str = "on_end_tag";
+    const SELMA_HANDLE_ELEMENT: &str = "handle_element";
+    const SELMA_HANDLE_TEXT: &str = "handle_text";
+
+    /// @yard
+    /// @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
+    /// @return [Selma::Rewriter]
+    fn new(args: &[Value]) -> Result<Self, magnus::Error> {
+        let (rb_sanitizer, rb_handlers) = Self::scan_parse_args(args)?;
+
+        let sanitizer = match rb_sanitizer {
+            None => {
+                let default_sanitizer = SelmaSanitizer::new(&[])?;
+                let wrapped_sanitizer = WrappedStruct::from(default_sanitizer);
+                wrapped_sanitizer.funcall::<&str, (), Value>("setup", ())?;
+                Some(wrapped_sanitizer.get().unwrap().to_owned())
+            }
+            Some(sanitizer_value) => match sanitizer_value {
+                None => None,
+                Some(sanitizer) => {
+                    sanitizer.funcall::<&str, (), Value>("setup", ())?;
+                    Some(sanitizer.get().unwrap().to_owned())
+                }
+            },
+        };
+
+        let handlers = match rb_handlers {
+            None => vec![],
+            Some(rb_handlers) => {
+                let mut handlers: Vec<Handler> = vec![];
+
+                for h in rb_handlers.each() {
+                    let rb_handler = h.unwrap();
+
+                    // prevents missing #selector from ruining things
+                    if !rb_handler.respond_to("selector", true).unwrap() {
+                        let classname = unsafe { rb_handler.classname() };
+                        return Err(magnus::Error::new(
+                            exception::no_method_error(),
+                            format!(
+                                "Could not call #selector on {:?}; is this an object that defines it?",
+                                classname
+                            ),
+                        ));
+                    }
+
+                    let rb_selector: WrappedStruct<SelmaSelector> =
+                        match rb_handler.funcall("selector", ()) {
+                            Err(e) => {
+                                return Err(magnus::Error::new(
+                                    exception::type_error(),
+                                    format!("Error instantiating selector: {}", e),
+                                ));
+                            }
+                            Ok(rb_selector) => rb_selector,
+                        };
+                    let handler = Handler {
+                        rb_handler,
+                        rb_selector,
+                        total_element_handler_calls: 0,
+                        total_elapsed_element_handlers: 0.0,
+
+                        total_text_handler_calls: 0,
+                        total_elapsed_text_handlers: 0.0,
+                    };
+                    handlers.push(handler);
+                }
+                handlers
+            }
+        };
+
+        if sanitizer.is_none() && handlers.is_empty() {
+            return Err(magnus::Error::new(
+                exception::arg_error(),
+                "Must provide a sanitizer or a handler",
+            ));
+        }
+
+        Ok(Self(std::cell::RefCell::new(Rewriter {
+            sanitizer,
+            handlers,
+            total_elapsed: 0.0,
+        })))
+    }
+
+    #[allow(clippy::let_unit_value)]
+    fn scan_parse_args(
+        args: &[Value],
+    ) -> Result<
+        (
+            Option<Option<WrappedStruct<SelmaSanitizer>>>,
+            Option<RArray>,
+        ),
+        magnus::Error,
+    > {
+        let args = scan_args::scan_args(args)?;
+        let _: () = args.required;
+        let _: () = args.optional;
+        let _: () = args.splat;
+        let _: () = args.trailing;
+        let _: () = args.block;
+
+        let kw = scan_args::get_kwargs::<
+            _,
+            (),
+            (
+                Option<Option<WrappedStruct<SelmaSanitizer>>>,
+                Option<RArray>,
+            ),
+            (),
+        >(args.keywords, &[], &["sanitizer", "handlers"])?;
+        let (rb_sanitizer, rb_handlers) = kw.optional;
+
+        Ok((rb_sanitizer, rb_handlers))
+    }
+
+    /// Perform HTML rewrite sequence.
+    fn rewrite(&self, html: String) -> Result<String, magnus::Error> {
+        let sanitized_html = match &self.0.borrow().sanitizer {
+            None => html,
+            Some(sanitizer) => {
+                // let first_pass_html = Self::perform_initial_sanitization(sanitizer, &html).unwrap();
+
+                // due to malicious html crafting
+                // (e.g. <<foo>script>...</script>, or <div <!-- comment -->> as in tests),
+                // we need to run sanitization several times to truly remove unwanted tags,
+                // because lol-html happily accepts this garbage (by design?)
+                let sanitized_html = Self::perform_sanitization(sanitizer, &html).unwrap();
+
+                String::from_utf8(sanitized_html).unwrap()
+            }
+        };
+        let binding = self.0.borrow_mut();
+        let handlers = &binding.handlers;
+
+        match Self::perform_handler_rewrite(self, handlers, sanitized_html) {
+            Ok(rewritten_html) => Ok(String::from_utf8(rewritten_html).unwrap()),
+            Err(err) => Err(magnus::Error::new(
+                exception::runtime_error(),
+                format!("{}", err),
+            )),
+        }
+    }
+
+    fn perform_sanitization(
+        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(());
+                        }
+                        sanitizer.sanitize_attributes(el);
+
+                        Ok(())
+                    })],
+                    ..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);
+                    }
+
+                    Ok(())
+                }));
+            }
+
+            let mut rewriter = HtmlRewriter::new(
+                Settings {
+                    element_content_handlers,
+                    ..Settings::default()
+                },
+                |c: &[u8]| output.extend_from_slice(c),
+            );
+
+            let result = rewriter.write(first_pass_html.as_slice());
+            if result.is_err() {
+                return Err(magnus::Error::new(
+                    exception::runtime_error(),
+                    format!("Failed to sanitize HTML: {}", result.unwrap_err()),
+                ));
+            }
+        }
+
+        Ok(output)
+    }
+
+    pub fn perform_handler_rewrite(
+        &self,
+        handlers: &[Handler],
+        html: String,
+    ) -> Result<Vec<u8>, magnus::Error> {
+        // TODO: this should ideally be done ahead of time, not on every `#rewrite` call
+        let mut element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> = vec![];
+
+        handlers.iter().for_each(|handler| {
+            let element_stack: Rc<RefCell<Vec<String>>> = Rc::new(RefCell::new(vec![]));
+
+            let selector = handler.rb_selector.get_static().unwrap();
+
+            // TODO: test final raise by simulating errors
+            if selector.match_element().is_some() {
+                let closure_element_stack = element_stack.clone();
+
+                element_content_handlers.push(element!(
+                    selector.match_element().unwrap(),
+                    move |el| {
+                        match Self::process_element_handlers(
+                            handler.rb_handler,
+                            el,
+                            &closure_element_stack.borrow(),
+                        ) {
+                            Ok(_) => Ok(()),
+                            Err(err) => Err(err.to_string().into()),
+                        }
+                    }
+                ));
+            }
+
+            if selector.match_text_within().is_some() {
+                let closure_element_stack = element_stack.clone();
+
+                element_content_handlers.push(text!(
+                    selector.match_text_within().unwrap(),
+                    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();
+                            if selector
+                                .ignore_text_within()
+                                .unwrap()
+                                .iter()
+                                .any(|f| f == &head_tag_name)
+                            {
+                                return Ok(());
+                            }
+                        }
+
+                        match Self::process_text_handlers(handler.rb_handler, text) {
+                            Ok(_) => Ok(()),
+                            Err(err) => Err(err.to_string().into()),
+                        }
+                    }
+                ));
+            }
+
+            // we need to check *every* element we iterate over, to create a stack of elements
+            element_content_handlers.push(element!("*", move |el| {
+                let tag_name = el.tag_name().to_lowercase();
+
+                // no need to track self-closing tags
+                if Tag::tag_from_tag_name(&tag_name).self_closing {
+                    return Ok(());
+                };
+
+                element_stack.as_ref().borrow_mut().push(tag_name);
+
+                let closure_element_stack = element_stack.clone();
+                el.on_end_tag(move |_end_tag: &mut EndTag| {
+                    let mut stack = closure_element_stack.as_ref().borrow_mut();
+                    stack.pop();
+                    Ok(())
+                });
+                Ok(())
+            }));
+        });
+
+        let mut output = vec![];
+        {
+            let mut rewriter = HtmlRewriter::new(
+                Settings {
+                    element_content_handlers,
+                    ..Settings::default()
+                },
+                |c: &[u8]| output.extend_from_slice(c),
+            );
+            match rewriter.write(html.as_bytes()) {
+                Ok(_) => {}
+                Err(err) => {
+                    return Err(magnus::Error::new(
+                        exception::runtime_error(),
+                        format!("{}", err),
+                    ));
+                }
+            }
+        }
+        Ok(output)
+    }
+
+    fn process_element_handlers(
+        rb_handler: Value,
+        element: &mut Element,
+        ancestors: &Vec<String>,
+    ) -> Result<(), magnus::Error> {
+        // if `on_end_tag` function is defined, call it
+        if rb_handler.respond_to(Self::SELMA_ON_END_TAG, true).unwrap() {
+            element.on_end_tag(move |end_tag| {
+                let rb_end_tag = SelmaHTMLEndTag::new(end_tag);
+
+                rb_handler
+                    .funcall::<_, _, Value>(Self::SELMA_ON_END_TAG, (rb_end_tag,))
+                    .unwrap();
+                Ok(())
+            });
+        }
+
+        let rb_element = SelmaHTMLElement::new(element, ancestors);
+        let rb_result =
+            rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_ELEMENT, (rb_element,));
+        match rb_result {
+            Ok(_) => Ok(()),
+            Err(err) => Err(magnus::Error::new(
+                exception::runtime_error(),
+                format!("{}", err),
+            )),
+        }
+    }
+
+    fn process_text_handlers(rb_handler: Value, text: &mut TextChunk) -> Result<(), magnus::Error> {
+        // prevents missing `handle_text` function
+        let content = text.as_str();
+
+        // FIXME: why does this happen?
+        if content.is_empty() {
+            return Ok(());
+        }
+        let rb_result = rb_handler.funcall(Self::SELMA_HANDLE_TEXT, (content,));
+
+        if rb_result.is_err() {
+            return Err(magnus::Error::new(
+                exception::type_error(),
+                format!(
+                    "Expected #{:?} to return a string: {:?}",
+                    Self::SELMA_HANDLE_TEXT,
+                    rb_result.err().unwrap()
+                ),
+            ));
+        }
+
+        let new_content: String = rb_result.unwrap();
+        // TODO: can this be an option?
+        text.replace(&new_content, ContentType::Html);
+
+        Ok(())
+    }
+}
+
+pub fn init(m_selma: RModule) -> Result<(), magnus::Error> {
+    let c_rewriter = m_selma
+        .define_class("Rewriter", Default::default())
+        .expect("cannot find class Selma::Rewriter");
+
+    c_rewriter.define_singleton_method("new", function!(SelmaRewriter::new, -1))?;
+    c_rewriter
+        .define_method("rewrite", method!(SelmaRewriter::rewrite, 1))
+        .expect("cannot define method `rewrite`");
+
+    Ok(())
+}