lol_html 3.0.0

use crate::AsciiCompatibleEncoding;
use crate::base::{Bytes, Range, SharedEncoding, SourceLocation};
use crate::html::{LocalName, Namespace};
use crate::html_content::{TextChunk, TextType};
use crate::parser::{
    ActionError, ActionResult, AttributeBuffer, Lexeme, LexemeSink, NonTagContentLexeme,
    NonTagContentTokenOutline, ParserDirective, ParserOutputSink, TagHintSink, TagLexeme,
    TagTokenOutline,
};
use crate::rewritable_units::TextDecoder;
use crate::rewritable_units::ToTokenResult;
use crate::rewritable_units::{BailOut, DocumentEnd, Serialize, ToToken, Token, TokenCaptureFlags};
use crate::rewriter::RewritingError;
use encoding_rs::Encoding;

pub(crate) struct AuxStartTagInfo<'i> {
    pub input: &'i Bytes<'i>,
    pub attr_buffer: &'i AttributeBuffer,
    pub self_closing: bool,
}

type AuxStartTagInfoRequest<C> =
    Box<dyn FnOnce(&mut C, AuxStartTagInfo<'_>) -> ActionResult<TokenCaptureFlags> + Send>;

// Pub only for integration tests
#[allow(private_interfaces)]
pub enum DispatcherError<C> {
    InfoRequest(AuxStartTagInfoRequest<C>),
    RewritingError(RewritingError),
}

// Pub only for integration tests
pub type StartTagHandlingResult<C> = Result<TokenCaptureFlags, DispatcherError<C>>;

// Pub only for integration tests
pub trait TransformController: Sized {
    fn initial_capture_flags(&self) -> TokenCaptureFlags;
    fn handle_start_tag(
        &mut self,
        name: LocalName<'_>,
        ns: Namespace,
    ) -> StartTagHandlingResult<Self>;
    fn handle_end_tag(&mut self, name: LocalName<'_>) -> TokenCaptureFlags;
    fn handle_token(&mut self, token: &mut Token<'_>) -> Result<(), RewritingError>;
    fn handle_end(&mut self, document_end: &mut DocumentEnd<'_>) -> Result<(), RewritingError>;
    fn should_emit_content(&self) -> bool;

    /// Invoked when the rewriter triggers a graceful bail-out. Default impl does nothing;
    /// the production `HtmlRewriteController` overrides this to run the user-registered
    /// bail-out handlers.
    fn handle_bail_out(&mut self, _error: &RewritingError, _bail_out: &mut BailOut<'_>) {}
}

/// Defines an interface for the [`HtmlRewriter`]'s output.
///
/// Implemented for [`Fn`] and [`FnMut`].
///
/// [`HtmlRewriter`]: struct.HtmlRewriter.html
/// [`Fn`]: https://doc.rust-lang.org/std/ops/trait.Fn.html
/// [`FnMut`]: https://doc.rust-lang.org/std/ops/trait.FnMut.html
pub trait OutputSink {
    /// Handles rewriter's output chunk.
    ///
    /// # Note
    /// The last chunk of the output has zero length.
    fn handle_chunk(&mut self, chunk: &[u8]);

    /// Called before the first `handle_chunk` and once when `<meta charset>` is applied.
    ///
    /// Implementations for closures won't be able to access this method
    fn set_encoding(&mut self, _new_encoding: AsciiCompatibleEncoding) {}
}

impl<F: FnMut(&[u8])> OutputSink for F {
    fn handle_chunk(&mut self, chunk: &[u8]) {
        self(chunk);
    }
}

// Pub only for integration tests
pub struct Dispatcher<C, O> {
    delegate: DispatcherDelegate<C, O>,
    text_decoder: TextDecoder,
    last_text_type: TextType,
    got_flags_from_hint: bool,
    pending_element_aux_info_req: Option<AuxStartTagInfoRequest<C>>,
    encoding: AsciiCompatibleEncoding,
    next_encoding: SharedEncoding,
}

/// Fields split out of `Dispatcher` for borrow checking of event handlers
struct DispatcherDelegate<C, O> {
    transform_controller: C,
    output_sink: O,
    remaining_content_start: usize,
    capture_flags: TokenCaptureFlags,
    emission_enabled: bool,
}

impl<C, O> DispatcherDelegate<C, O>
where
    C: TransformController,
    O: OutputSink,
{
    fn flush_remaining_input(&mut self, input: &[u8], consumed_byte_count: usize) {
        if self.emission_enabled {
            let output = input
                .get(self.remaining_content_start..consumed_byte_count)
                .unwrap_or_default();

            if !output.is_empty() {
                self.output_sink.handle_chunk(output);
            }
        }

        self.remaining_content_start = 0;
    }

    fn finish(&mut self, encoding: &'static Encoding, input: &[u8]) -> Result<(), RewritingError> {
        self.flush_remaining_input(input, input.len());

        let mut document_end = DocumentEnd::new(&mut self.output_sink, encoding);

        self.transform_controller.handle_end(&mut document_end)?;

        // NOTE: output the finalizing chunk.
        self.output_sink.handle_chunk(&[]);

        Ok(())
    }

    /// Emit the raw chunk of input that lies between the previously emitted byte and the
    /// start of `lexeme`, and move `remaining_content_start` to the start of `lexeme`.
    ///
    /// This is the first half of what used to be `lexeme_consumed()`. It is split out so that
    /// the second half ([`consume_lexeme()`]) can be deferred until after the lexeme's token
    /// has been successfully emitted; that way, if token emission fails (e.g. a content
    /// handler returns an error), `remaining_content_start` still points at the failing
    /// lexeme's start and a graceful bail-out can flush the lexeme bytes raw rather than
    /// losing them.
    ///
    /// [`consume_lexeme()`]: Self::consume_lexeme
    fn emit_chunk_before_lexeme<T>(&mut self, lexeme: &Lexeme<'_, T>) {
        let lexeme_range = lexeme.raw_range();

        let chunk_range = Range {
            start: self.remaining_content_start,
            end: lexeme_range.start,
        };

        let chunk = lexeme.input().slice(chunk_range);

        if self.emission_enabled && !chunk.is_empty() {
            self.output_sink.handle_chunk(&chunk);
        }

        self.remaining_content_start = lexeme_range.start;
    }

    /// Advance `remaining_content_start` past the end of `lexeme`, marking it as committed.
    /// Must only be called after the lexeme's token has been successfully emitted; see
    /// [`emit_chunk_before_lexeme()`].
    ///
    /// [`emit_chunk_before_lexeme()`]: Self::emit_chunk_before_lexeme
    fn consume_lexeme<T>(&mut self, lexeme: &Lexeme<'_, T>) {
        self.remaining_content_start = lexeme.raw_range().end;
    }

    #[inline]
    fn token_produced(&mut self, mut token: Token<'_>) -> Result<(), RewritingError> {
        trace!(@output token);

        self.transform_controller.handle_token(&mut token)?;

        if self.emission_enabled {
            token.into_bytes(&mut |c| self.output_sink.handle_chunk(c))?;
        }
        Ok(())
    }

    fn text_token_produced(
        &mut self,
        text: &str,
        encoding: &'static Encoding,
        text_type: TextType,
        is_last_in_node: bool,
        source_location: SourceLocation,
    ) -> Result<(), RewritingError> {
        let mut token = Token::TextChunk(TextChunk::new(
            text,
            text_type,
            is_last_in_node,
            encoding,
            source_location,
        ));

        trace!(@output token);

        self.transform_controller.handle_token(&mut token)?;

        if self.emission_enabled {
            token.into_bytes(&mut |c| self.output_sink.handle_chunk(c))?;
        }
        Ok(())
    }

    #[inline]
    fn should_stop_removing_element_content(&self) -> bool {
        !self.emission_enabled && self.transform_controller.should_emit_content()
    }
}

impl<C, O> Dispatcher<C, O>
where
    C: TransformController,
    O: OutputSink,
{
    #[inline]
    pub fn new(
        transform_controller: C,
        mut output_sink: O,
        encoding: AsciiCompatibleEncoding,
        next_encoding: SharedEncoding,
    ) -> Self {
        let capture_flags = transform_controller.initial_capture_flags();
        output_sink.set_encoding(encoding);

        Self {
            delegate: DispatcherDelegate {
                transform_controller,
                output_sink,
                capture_flags,
                remaining_content_start: 0,
                emission_enabled: true,
            },
            text_decoder: TextDecoder::new(encoding),
            last_text_type: TextType::Data,
            encoding,
            got_flags_from_hint: false,
            pending_element_aux_info_req: None,
            next_encoding,
        }
    }

    fn flush_encoding_change(&mut self) {
        if let Some(&next_encoding) = self.next_encoding.get()
            && next_encoding != self.encoding
        {
            self.encoding = next_encoding;
            self.text_decoder.set_encoding(next_encoding);
            self.delegate.output_sink.set_encoding(next_encoding);
        }
    }

    #[inline(never)]
    fn try_produce_token_from_lexeme<'i, T>(&mut self, lexeme: &Lexeme<'i, T>) -> ActionResult
    where
        Lexeme<'i, T>: ToToken,
    {
        match lexeme.to_token(&mut self.delegate.capture_flags, self.encoding.get()) {
            ToTokenResult::Token(token) => {
                self.delegate.emit_chunk_before_lexeme(lexeme);
                self.delegate.token_produced(token)?;
                self.delegate.consume_lexeme(lexeme);
                // handler for <meta charset> may have changed encoding
                self.flush_encoding_change();
            }
            ToTokenResult::Text(text_type) => {
                self.delegate.emit_chunk_before_lexeme(lexeme);
                self.last_text_type = text_type;
                self.text_decoder.feed_text(
                    lexeme.spanned(),
                    false,
                    &mut |text, is_last, encoding, source_location| {
                        self.delegate.text_token_produced(
                            text,
                            encoding,
                            self.last_text_type,
                            is_last,
                            source_location,
                        )
                    },
                )?;
                self.delegate.consume_lexeme(lexeme);
            }
            ToTokenResult::None => {}
        }
        Ok(())
    }

    #[inline]
    const fn get_next_parser_directive(&self) -> ParserDirective {
        if !self.delegate.capture_flags.is_empty() {
            ParserDirective::Lex
        } else {
            ParserDirective::WherePossibleScanForTagsOnly
        }
    }

    fn adjust_capture_flags_for_tag_lexeme(&mut self, lexeme: &TagLexeme<'_>) -> ActionResult {
        let input = lexeme.input();

        macro_rules! get_flags_from_aux_info_res {
            ($handler:expr, $attributes:expr, $self_closing:expr) => {
                $handler(
                    &mut self.delegate.transform_controller,
                    AuxStartTagInfo {
                        input,
                        attr_buffer: $attributes,
                        self_closing: $self_closing,
                    },
                )
            };
        }

        let capture_flags = match self.pending_element_aux_info_req.take() {
            // NOTE: tag hint was produced for the tag, but
            // attributes and self closing flag were requested.
            Some(aux_info_req) => match *lexeme.token_outline() {
                TagTokenOutline::StartTag {
                    ref attributes,
                    self_closing,
                    ..
                } => {
                    get_flags_from_aux_info_res!(aux_info_req, &attributes, self_closing)
                }
                TagTokenOutline::EndTag { .. } => Err(ActionError::internal(
                    "Tag should be a start tag at this point",
                )),
            },

            // NOTE: tag hint hasn't been produced for the tag, because
            // parser is not in the tag scan mode.
            None => match *lexeme.token_outline() {
                TagTokenOutline::StartTag {
                    name,
                    name_hash,
                    ns,
                    ref attributes,
                    self_closing,
                } => {
                    let name = LocalName::new(*input, name, name_hash);

                    match self
                        .delegate
                        .transform_controller
                        .handle_start_tag(name, ns)
                    {
                        Ok(flags) => Ok(flags),
                        Err(DispatcherError::InfoRequest(aux_info_req)) => {
                            get_flags_from_aux_info_res!(aux_info_req, &attributes, self_closing)
                        }
                        Err(DispatcherError::RewritingError(e)) => Err(e.into()),
                    }
                }

                TagTokenOutline::EndTag { name, name_hash } => {
                    let name = LocalName::new(*input, name, name_hash);
                    Ok(self.delegate.transform_controller.handle_end_tag(name))
                }
            },
        };

        match capture_flags {
            Ok(flags) => {
                self.delegate.capture_flags = flags;
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    #[inline]
    fn apply_capture_flags_from_hint_and_get_next_parser_directive(
        &mut self,
        flags: TokenCaptureFlags,
    ) -> ParserDirective {
        self.delegate.capture_flags = flags;
        self.got_flags_from_hint = true;
        self.get_next_parser_directive()
    }

    /// Emit text chunk with is_last_in_node
    #[inline]
    fn flush_pending_captured_text(&mut self) -> Result<(), RewritingError> {
        self.text_decoder
            .flush_pending(&mut |text, is_last, encoding, source_location| {
                self.delegate.text_token_produced(
                    text,
                    encoding,
                    self.last_text_type,
                    is_last,
                    source_location,
                )
            })
    }

    pub fn flush_remaining_input(&mut self, input: &[u8], consumed_byte_count: usize) {
        self.delegate
            .flush_remaining_input(input, consumed_byte_count);
    }

    /// Flushes all input bytes the dispatcher has received but not yet emitted, *as-is*. Used
    /// when `MemorySettings::graceful_bail_out_on_memory_limit_exceeded` is enabled and a memory
    /// error is being propagated: the caller wants to preserve enough of the input in the sink
    /// to be able to continue the response without breaking it.
    ///
    /// Unlike [`flush_remaining_input()`], this ignores `emission_enabled`. If a handler was
    /// removing element content at the time of the bail-out, those bytes will still be flushed
    /// raw. The alternative (skipping the flush) would lose bytes from the input chunk entirely,
    /// which the caller cannot recover from since they don't buffer their input.
    pub fn flush_for_bail_out(&mut self, input: &[u8]) {
        let output = input
            .get(self.delegate.remaining_content_start..)
            .unwrap_or_default();

        if !output.is_empty() {
            self.delegate.output_sink.handle_chunk(output);
        }

        self.delegate.remaining_content_start = 0;
    }

    /// Invokes the transform controller's bail-out handlers (in registration order),
    /// constructing a [`BailOut`] wrapper around the output sink and the current encoding.
    /// Must be called *before* [`flush_for_bail_out()`] so that handler emissions land in
    /// the sink ahead of the raw flush of remaining unparsed input.
    ///
    /// [`flush_for_bail_out()`]: Self::flush_for_bail_out
    pub fn run_bail_out_handlers(&mut self, error: &RewritingError) {
        let mut bail_out = BailOut::new(&mut self.delegate.output_sink, self.encoding.get());
        self.delegate
            .transform_controller
            .handle_bail_out(error, &mut bail_out);
    }

    pub fn finish(&mut self, input: &[u8]) -> Result<(), RewritingError> {
        self.delegate.finish(self.encoding.get(), input)
    }
}

impl<C, O> LexemeSink for Dispatcher<C, O>
where
    C: TransformController,
    O: OutputSink,
{
    fn handle_tag(&mut self, lexeme: &TagLexeme<'_>) -> ActionResult<ParserDirective> {
        // NOTE: flush pending text before reporting tag to the transform controller.
        // Otherwise, transform controller can enable or disable text handlers too early.
        // In case of start tag, newly matched element text handlers
        // will receive leftovers from the previous match. And, in case of end tag,
        // handlers will be disabled before the receive the finalizing chunk.
        self.flush_pending_captured_text()?;

        if self.got_flags_from_hint {
            self.got_flags_from_hint = false;
        } else {
            self.adjust_capture_flags_for_tag_lexeme(lexeme)?;
        }

        if let TagTokenOutline::EndTag { .. } = lexeme.token_outline() {
            if self.delegate.should_stop_removing_element_content() {
                self.delegate.emission_enabled = true;
                self.delegate.remaining_content_start = lexeme.raw_range().start;
            }
        }

        self.try_produce_token_from_lexeme(lexeme)?;
        self.delegate.emission_enabled = self.delegate.transform_controller.should_emit_content();

        Ok(self.get_next_parser_directive())
    }

    #[inline]
    fn handle_non_tag_content(&mut self, lexeme: &NonTagContentLexeme<'_>) -> ActionResult {
        match lexeme.token_outline() {
            Some(NonTagContentTokenOutline::Text(_)) => {}
            // when it's None, it still needs a flush for CDATA
            _ => self.flush_pending_captured_text()?,
        }
        self.try_produce_token_from_lexeme(lexeme)
    }
}

impl<C, O> TagHintSink for Dispatcher<C, O>
where
    C: TransformController,
    O: OutputSink,
{
    fn handle_start_tag_hint(
        &mut self,
        name: LocalName<'_>,
        ns: Namespace,
    ) -> Result<ParserDirective, RewritingError> {
        match self
            .delegate
            .transform_controller
            .handle_start_tag(name, ns)
        {
            Ok(flags) => {
                Ok(self.apply_capture_flags_from_hint_and_get_next_parser_directive(flags))
            }
            Err(DispatcherError::InfoRequest(aux_info_req)) => {
                self.got_flags_from_hint = false;
                self.pending_element_aux_info_req = Some(aux_info_req);

                Ok(ParserDirective::Lex)
            }
            Err(DispatcherError::RewritingError(e)) => Err(e),
        }
    }

    fn handle_end_tag_hint(
        &mut self,
        name: LocalName<'_>,
    ) -> Result<ParserDirective, RewritingError> {
        self.flush_pending_captured_text()?;

        let mut flags = self.delegate.transform_controller.handle_end_tag(name);

        // NOTE: if emission was disabled (i.e. we've been removing element content)
        // we need to request the end tag lexeme, to ensure that we have it.
        // Otherwise, if we have unfinished end tag in the end of input we'll emit
        // it where we shouldn't.
        if self.delegate.should_stop_removing_element_content() {
            flags |= TokenCaptureFlags::NEXT_END_TAG;
        }

        Ok(self.apply_capture_flags_from_hint_and_get_next_parser_directive(flags))
    }
}

impl<C, O> ParserOutputSink for Dispatcher<C, O>
where
    C: TransformController,
    O: OutputSink,
{
}