Skip to content

Commit 0f06b13

Browse files
committed
Add bail-out handler API for flushing buffered state on graceful bail-out.
Graceful bail-out (`MemoryLimitExceeded` or `ContentHandlerError` with the matching flag on) flushes the unparsed input remainder raw to the sink and propagates the error. That is enough for handlers that only transform tokens they see, but handlers that buffer state across the document (e.g. ROFL's email-obfuscation module, which holds up to ~128 chars in a text buffer while deciding whether they belong to an email) lose that state on bail-out and produce a response with a gap. This commit adds a hook that fires once on a graceful bail-out, immediately before the raw flush, and lets handlers append final bytes to the sink: 1. New rewritable unit `BailOut` with a single method `append(content, content_type)`, modelled after `DocumentEnd::append`. The wrapper carries the rewriter's current encoding (after any `<meta charset>`-driven change), so encoding-correctness is automatic. 2. New builder method `Settings::append_bail_out_handler` (and the `RewriteStrSettings` mirror) plus `bail_out!` macro for type-hint ergonomics, parallel to the existing `element!` / `end!` macros. 3. `HandlerTypes` grows a `BailOutHandler<'h>` associated type, with `LocalHandlerTypes` aliasing `BailOutHandler<'h>` and `SendHandlerTypes` aliasing `BailOutHandlerSend<'h>`. Matching `IntoHandler` impls cover both bare-closure cases. 4. `TransformController` grows a `handle_bail_out` method with an empty default impl so existing implementors (test fixtures, parser-trace tool) keep compiling. `HtmlRewriteController` overrides it to iterate the user-registered handlers in registration order. 5. `Dispatcher::run_bail_out_handlers` constructs the `BailOut` wrapper and delegates to the controller. It is invoked from every existing graceful bail-out site in `TransformStream::write()` (3 sites: `Arena::append`, `Parser::parse`, `Arena::init_with`) and `TransformStream::end()` (1 site), gated on `should_bail_out_for(&err)`. Hook output therefore lands in the sink as `[transformed prefix] + [hook output] + [raw remainder]`. 6. `RewritingError` is marked `#[non_exhaustive]` so we can add variants in future minor releases. `match`es still work; only exhaustive external matches need a catch-all arm. The `end()` bail-out site is defensive: it is symmetric with the `write()` sites but is not reachable through normal input. EOF-in-tag / -attribute / -comment emits as text per HTML5, so content-handler errors don't fire from `parser.parse(_, true)`, and memory errors fire earlier in `write()`. Tested implicitly via the shared call path.
1 parent ba25359 commit 0f06b13

9 files changed

Lines changed: 517 additions & 21 deletions

File tree

CHANGELOG.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,17 @@
1111
rewriter flushes remaining input bytes before propagating a handler error, preserving
1212
the response. Currently exposed via the Rust API only; the C API still uses the original
1313
behavior.
14+
- Added `Settings::append_bail_out_handler()` and the matching `bail_out!` macro,
15+
`BailOut` rewritable unit, and `BailOutHandler` / `BailOutHandlerSend` type aliases.
16+
Bail-out handlers fire immediately before the raw flush of remaining unparsed input on a
17+
graceful bail-out (memory or content-handler error). Handlers receive the
18+
`RewritingError` and a `BailOut` through which they can append final bytes to the sink
19+
via `BailOut::append(content, content_type)`. Intended for handlers that buffer state
20+
across the document (e.g. text-buffering handlers that defer emission) and need to
21+
flush that state on bail-out.
22+
- Marked `RewritingError` `#[non_exhaustive]` so future error variants can be added without
23+
a major version bump. External callers can still `match` on it, but must include a
24+
catch-all `_ =>` arm.
1425
- Reworked `Settings`, `MemorySettings` and `RewriteStrSettings` to use a consuming-builder
1526
API. Fields are now private; construction is via `::new()` plus chained `with_*` setters
1627
and `append_*` methods for the content-handler vectors. This makes future field additions

src/lib.rs

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,10 @@ mod transform_stream;
4141
use cfg_if::cfg_if;
4242

4343
pub use self::rewriter::{
44-
AsciiCompatibleEncoding, CommentHandler, DoctypeHandler, DocumentContentHandlers,
45-
ElementContentHandlers, ElementHandler, EndHandler, EndTagHandler, HandlerResult, HandlerTypes,
46-
HtmlRewriter, LocalHandlerTypes, MemorySettings, RewriteStrSettings, Settings, TextHandler,
47-
rewrite_str,
44+
AsciiCompatibleEncoding, BailOutHandler, CommentHandler, DoctypeHandler,
45+
DocumentContentHandlers, ElementContentHandlers, ElementHandler, EndHandler, EndTagHandler,
46+
HandlerResult, HandlerTypes, HtmlRewriter, LocalHandlerTypes, MemorySettings,
47+
RewriteStrSettings, Settings, TextHandler, rewrite_str,
4848
};
4949
pub use self::selectors_vm::Selector;
5050
pub use self::transform_stream::OutputSink;
@@ -56,9 +56,10 @@ pub use self::transform_stream::OutputSink;
5656
/// Rewriting is sequential, so there's no benefit from using the `Send`-compatible rewriter.
5757
pub mod send {
5858
pub use crate::rewriter::{
59-
CommentHandlerSend as CommentHandler, DoctypeHandlerSend as DoctypeHandler,
60-
ElementHandlerSend as ElementHandler, EndHandlerSend as EndHandler,
61-
EndTagHandlerSend as EndTagHandler, TextHandlerSend as TextHandler,
59+
BailOutHandlerSend as BailOutHandler, CommentHandlerSend as CommentHandler,
60+
DoctypeHandlerSend as DoctypeHandler, ElementHandlerSend as ElementHandler,
61+
EndHandlerSend as EndHandler, EndTagHandlerSend as EndTagHandler,
62+
TextHandlerSend as TextHandler,
6263
};
6364
pub use crate::rewriter::{IntoHandler, SendHandlerTypes};
6465

@@ -95,7 +96,7 @@ pub mod errors {
9596
/// HTML content descriptors that can be produced and modified by a rewriter.
9697
pub mod html_content {
9798
pub use super::rewritable_units::{
98-
Attribute, Comment, ContentType, Doctype, DocumentEnd, Element, EndTag, StartTag,
99+
Attribute, BailOut, Comment, ContentType, Doctype, DocumentEnd, Element, EndTag, StartTag,
99100
StreamingHandler, StreamingHandlerSink, TextChunk, UserData,
100101
};
101102

src/rewritable_units/bail_out.rs

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
use super::{ContentType, StreamingHandlerSink};
2+
use crate::transform_stream::OutputSink;
3+
use encoding_rs::Encoding;
4+
5+
/// A rewritable unit that represents the moment the rewriter is about to abandon
6+
/// processing through a graceful bail-out.
7+
///
8+
/// Bail-out handlers registered via [`Settings::append_bail_out_handler()`] receive a
9+
/// `&mut BailOut` and can emit final bytes into the output sink via [`append()`]. This
10+
/// is the only opportunity for content other handlers have buffered (e.g. text withheld
11+
/// pending a future chunk) to land in the response when the rewriter aborts.
12+
///
13+
/// Bytes appended via this unit are written *before* the rewriter's own raw flush of
14+
/// remaining unparsed input. The resulting sink order is:
15+
///
16+
/// 1. Transformed bytes the rewriter already emitted normally.
17+
/// 2. Bytes appended by bail-out handlers, in registration order.
18+
/// 3. The rewriter's raw flush of the chunk's unparsed suffix.
19+
///
20+
/// [`Settings::append_bail_out_handler()`]:
21+
/// crate::Settings::append_bail_out_handler
22+
/// [`append()`]: Self::append
23+
pub struct BailOut<'a> {
24+
output_sink: &'a mut dyn OutputSink,
25+
encoding: &'static Encoding,
26+
}
27+
28+
impl<'a> BailOut<'a> {
29+
#[inline]
30+
#[must_use]
31+
pub(crate) fn new(output_sink: &'a mut dyn OutputSink, encoding: &'static Encoding) -> Self {
32+
Self {
33+
output_sink,
34+
encoding,
35+
}
36+
}
37+
38+
/// Appends `content` at the bail-out point.
39+
///
40+
/// Subsequent calls to this method append `content` to the previously inserted
41+
/// content within the same bail-out invocation. When multiple bail-out handlers are
42+
/// registered, their `append` calls are concatenated in registration order.
43+
///
44+
/// `content_type` controls how the content is interpreted before being written to
45+
/// the sink. See [`ContentType`].
46+
///
47+
/// # Example
48+
///
49+
/// ```
50+
/// use lol_html::{bail_out, Settings};
51+
/// use lol_html::errors::RewritingError;
52+
/// use lol_html::html_content::ContentType;
53+
///
54+
/// // A handler that, on content-handler-error bail-out, drops a notice into the sink
55+
/// // before the rewriter's own raw flush of remaining unparsed input.
56+
/// let settings = Settings::new()
57+
/// .with_graceful_bail_out_on_content_handler_error(true)
58+
/// .append_bail_out_handler(bail_out!(|err, bail_out| {
59+
/// if matches!(err, RewritingError::ContentHandlerError(_)) {
60+
/// bail_out.append("<!-- bailed out -->", ContentType::Html);
61+
/// }
62+
/// }));
63+
/// # let _ = settings;
64+
/// ```
65+
#[inline]
66+
pub fn append(&mut self, content: &str, content_type: ContentType) {
67+
StreamingHandlerSink::new(self.encoding, &mut |c| {
68+
self.output_sink.handle_chunk(c);
69+
})
70+
.write_str(content, content_type);
71+
}
72+
}

src/rewritable_units/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ pub(crate) use self::mutations::{Mutations, StringChunk};
44
pub(crate) use self::text_decoder::TextDecoder;
55
pub(crate) use self::text_encoder::{IncompleteUtf8Resync, TextEncoder};
66

7+
pub use self::bail_out::*;
78
pub use self::document_end::*;
89
pub use self::element::*;
910
pub use self::mutations::{ContentType, StreamingHandler};
@@ -83,6 +84,7 @@ macro_rules! impl_user_data {
8384
#[macro_use]
8485
mod mutations;
8586

87+
mod bail_out;
8688
mod document_end;
8789
mod element;
8890
mod streaming_sink;

0 commit comments

Comments
 (0)