src/workerd/api/html-rewriter.c++

File
Blob: src/workerd/api/html-rewriter.c++

49.0 KB
1 // Copyright (c) 2017-2022 Cloudflare, Inc.
2 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
3 //     https://opensource.org/licenses/Apache-2.0
4  
5 #include "html-rewriter.h"
6  
7 #include "util.h"
8  
9 #include <workerd/api/streams/common.h>
10 #include <workerd/api/streams/identity-transform-stream.h>
11 #include <workerd/io/features.h>
12 #include <workerd/io/io-context.h>
13  
14 #include <lol_html.h>
15  
16 struct lol_html_HtmlRewriter {};
17 struct lol_html_HtmlRewriterBuilder {};
18 struct lol_html_AttributesIterator {};
19 struct lol_html_Selector {};
20 // TODO(cleanup): These are defined internally in lol-html, but kj::Own<T> needs to check whether
21 //   or not T is polymorphic, so here's a dummy definition.
22  
23 namespace workerd::api {
24  
25 namespace {
26  
27 // =======================================================================================
28 // RAII helpers for lol-html
29  
30 // RAII helper for lol-html types which are managed by pointers and have straightforward _free()
31 // functions.
32 template <typename T, void (*lolhtmlFree)(T*)>
33 class LolHtmlDisposer: public kj::Disposer {
34  public:
35   static const LolHtmlDisposer INSTANCE;
36  
37  protected:
38   void disposeImpl(void* pointer) const override {
39     lolhtmlFree(reinterpret_cast<T*>(pointer));
40   }
41 };
42  
43 template <typename T, void (*lolhtmlFree)(T*)>
44 const LolHtmlDisposer<T, lolhtmlFree> LolHtmlDisposer<T, lolhtmlFree>::INSTANCE;
45  
46 #define LOL_HTML_OWN(name, ...)                                                                    \
47   ({                                                                                               \
48     using T = lol_html_##name##_t;                                                                 \
49     constexpr auto* lolhtmlFree = lol_html_##name##_free;                                          \
50     kj::Own<T>(&check(__VA_ARGS__), LolHtmlDisposer<T, lolhtmlFree>::INSTANCE);                    \
51   })
52  
53 // RAII helper for lol_html_str_t.
54 //
55 // We cannot use a kj::Own<T> because lol_html_str_t is a struct, not a pointer, so instead we
56 // have this LolString RAII wrapper.
57 //
58 // Use `kj::str(LolString.asChars())` to allocate your own copy of a LolString.
59 class LolString {
60  public:
61   explicit LolString(lol_html_str_t s): chars(s.data, s.len) {}
62   ~LolString() noexcept(false) {
63     lol_html_str_free({chars.begin(), chars.size()});
64   }
65   KJ_DISALLOW_COPY(LolString);
66  
67   kj::ArrayPtr<const char> asChars() const {
68     return chars;
69   }
70  
71   kj::Maybe<kj::String> asKjString() {
72     if (chars.begin() != nullptr) {
73       return kj::str(chars);
74     } else {
75       return kj::none;
76     }
77   }
78  
79  private:
80   kj::ArrayPtr<const char> chars;
81 };
82  
83 // =======================================================================================
84 // Error checking for lol-html
85  
86 kj::Maybe<kj::Exception> tryGetLastError() {
87   auto maybeErrorString = lol_html_take_last_error();
88   if (maybeErrorString.data == nullptr) {
89     return kj::none;
90   }
91   auto errorString = LolString(maybeErrorString);
92   return kj::Exception(kj::Exception::Type::FAILED, __FILE__, __LINE__,
93       kj::str(JSG_EXCEPTION(TypeError) ": Parser error: ", errorString.asChars()));
94 }
95  
96 void discardLastError() {
97   auto drop = LolString(lol_html_take_last_error());
98 }
99  
100 kj::Exception getLastError() {
101   return KJ_REQUIRE_NONNULL(tryGetLastError(),
102       "lol-html reported error through return value, but lol_html_take_last_error() is null");
103 }
104  
105 int check(int rc) {
106   if (rc == -1) {
107     kj::throwFatalException(getLastError());
108   }
109   return rc;
110 }
111  
112 template <typename T>
113 [[nodiscard]] T& check(T* ptr) {
114   // Nodiscard, because this function typically checks references that must later be freed.
115   if (ptr == nullptr) {
116     kj::throwFatalException(getLastError());
117   }
118   return *ptr;
119 }
120  
121 // Helper function to determine if a content token is still valid. Each content token has an
122 // implementation object inside a Maybe -- when HTMLRewriter::TokenScope (defined below)
123 // gets destroyed, that Maybe gets nullified, and the content token becomes a dead, useless,
124 // JavaScript object occupying space, waiting to get garbage collected.
125 //
126 // In other words, if you try to access a content token (Element, Text, etc.) outside of a
127 // content handler, you're gonna get this exception.
128 template <typename T>
129 decltype(auto) checkToken(kj::Maybe<T>& impl) {
130   return JSG_REQUIRE_NONNULL(impl, TypeError,
131       "This content token is no longer valid. Content tokens are only valid "
132       "during the execution of the relevant content handler.");
133 }
134  
135 }  // namespace
136  
137 // =======================================================================================
138 // HTMLRewriter::TokenScope
139  
140 class HTMLRewriter::TokenScope {
141  public:
142   template <typename T>
143   explicit TokenScope(jsg::Ref<T>& value): contentToken(value.addRef()) {}
144   ~TokenScope() noexcept(false) {
145     KJ_IF_SOME(token, contentToken) {
146       token->htmlContentScopeEnd();
147     }
148   }
149   TokenScope(TokenScope&& o): contentToken(kj::mv(o.contentToken)) {
150     o.contentToken = kj::none;
151   }
152   KJ_DISALLOW_COPY(TokenScope);
153  
154  private:
155   kj::Maybe<jsg::Ref<HTMLRewriter::Token>> contentToken;
156 };
157  
158 namespace {
159  
160 // =======================================================================================
161 // Rewriter
162 using ElementCallbackFunction = HTMLRewriter::ElementCallbackFunction;
163  
164 struct UnregisteredElementHandlers {
165   kj::Own<lol_html_Selector> selector;
166  
167   // The actual handler functions. We store them as jsg::Values for compatibility with GcVisitor.
168  
169   jsg::Optional<ElementCallbackFunction> element;
170   jsg::Optional<ElementCallbackFunction> comments;
171   jsg::Optional<ElementCallbackFunction> text;
172  
173   void visitForGc(jsg::GcVisitor& visitor) {
174     visitor.visit(element, comments, text);
175   }
176  
177   JSG_MEMORY_INFO(UnregisteredElementHandlers) {
178     tracker.trackField("element", element);
179     tracker.trackField("comments", comments);
180     tracker.trackField("text", text);
181   }
182 };
183  
184 struct UnregisteredDocumentHandlers {
185  
186   // The actual handler functions. We store them as jsg::Values for compatibility with GcVisitor.
187  
188   jsg::Optional<ElementCallbackFunction> doctype;
189   jsg::Optional<ElementCallbackFunction> comments;
190   jsg::Optional<ElementCallbackFunction> text;
191   jsg::Optional<ElementCallbackFunction> end;
192  
193   // The `this` object used to call the handler functions.
194  
195   void visitForGc(jsg::GcVisitor& visitor) {
196     visitor.visit(doctype, comments, text, end);
197   }
198  
199   JSG_MEMORY_INFO(UnregisteredDocumentHandlers) {
200     tracker.trackField("doctype", doctype);
201     tracker.trackField("comments", comments);
202     tracker.trackField("text", text);
203     tracker.trackField("end", end);
204   }
205 };
206  
207 using UnregisteredElementOrDocumentHandlers =
208     kj::OneOf<UnregisteredElementHandlers, UnregisteredDocumentHandlers>;
209  
210 }  // namespace
211  
212 // Wrapper around an actual rewriter (streaming parser).
213 class Rewriter final: public WritableStreamSink {
214  public:
215   explicit Rewriter(jsg::Lock& js,
216       kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
217       kj::ArrayPtr<const char> encoding,
218       kj::Own<WritableStreamSink> inner);
219   KJ_DISALLOW_COPY_AND_MOVE(Rewriter);
220  
221   // WritableStreamSink implementation. The input body pumpTo() operation calls these.
222   kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override;
223   kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override;
224   kj::Promise<void> end() override;
225   void abort(kj::Exception reason) override;
226  
227   // Implementation for `Element::onEndTag` to avoid exposing private details of Rewriter.
228   void onEndTag(lol_html_element_t* element, ElementCallbackFunction&& callback);
229  
230   lol_html_streaming_handler_t registerReplacer(jsg::Ref<ReadableStream> content, bool isHtml);
231  
232   ~Rewriter() {
233     KJ_ASSERT(registeredReplacers.size() == 0, "Some replacers were leaked by lol-html");
234   }
235  
236  private:
237   // Wait for the write promise (if any) produced by our `output()` callback, then, if there is a
238   // stored exception, abort the wrapped WritableStreamSink with it, then return the exception.
239   // Otherwise, just return.
240   kj::Promise<void> finishWrite();
241  
242   kj::Promise<void> flushWrite();
243  
244   static kj::Own<lol_html_HtmlRewriter> buildRewriter(jsg::Lock& js,
245       kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
246       kj::ArrayPtr<const char> encoding,
247       Rewriter& rewriterWrapper);
248  
249   static void output(const char* buffer, size_t size, void* userdata);
250   void outputImpl(kj::ArrayPtr<const byte> buffer);
251  
252   void tryHandleCancellation(int rc) {
253     if (canceled) {
254       canceled = false;
255  
256       // We canceled this, which means we used LOL_HTML_STOP. That means we have an error sitting
257       // in the error buffer in the lol-html C API. Let's make sure our return code is -1 and get
258       // rid of that error value to make sure nobody picks it up later on accident and thinks an
259       // error occurred.
260       KJ_ASSERT(rc == -1);
261       discardLastError();
262  
263       throw kj::CanceledException{};
264     }
265   }
266  
267   friend class ::workerd::api::HTMLRewriter;
268  
269   // Keeps track of streams currently being used as replacement content for tokens.
270   // lol-html will invoke replacerThunk with a pointer to the RegisteredReplacer to be used.
271   class RegisteredReplacer {
272    public:
273     Rewriter& rewriter;
274     bool isHtml;
275     jsg::Ref<ReadableStream> stream;
276   };
277  
278   struct RegisteredHandler {
279     // A back-reference to the rewriter which owns this particular registered handler.
280     Rewriter& rewriter;
281  
282     ElementCallbackFunction callback;
283   };
284  
285   kj::Vector<kj::Own<RegisteredHandler>> registeredHandlers;
286   // TODO(perf): Don't store Owns. We need to pass stable pointers as the userdata parameter to
287   //   lol_html_rewriter_builder_add_*_content_handlers(), but don't have a really easy way to
288   //   know precisely how many handlers we're going to register beforehand, so we need a vector. But
289   //   vectors can grow, moving their objects around, invalidating pointers into their storage.
290  
291   // This is separate from `registeredHandlers` so we can delete them more eagerly when EndTags are
292   // destroyed, and not have to look through all other handlers.
293   kj::Vector<kj::Own<RegisteredHandler>> registeredEndTagHandlers;
294   // TODO(perf) Don't store Owns, same as `registeredHandlers` above.
295  
296   template <typename T, typename CType = T::CType>
297   static lol_html_rewriter_directive_t thunk(CType* content, void* userdata);
298   template <typename T, typename CType = T::CType>
299   lol_html_rewriter_directive_t thunkImpl(CType* content, RegisteredHandler& registration);
300   template <typename T, typename CType = T::CType>
301   kj::Promise<void> thunkPromise(CType* content, RegisteredHandler& registration);
302  
303   // Eagerly free this handler. Should only be called if we're confident the handler will never be
304   // used again.
305   void removeEndTagHandler(RegisteredHandler& registration);
306  
307   // Field stores a list of readable streams that are being used by lol-html for replacements
308   kj::HashMap<void*, kj::Own<RegisteredReplacer>> registeredReplacers;
309  
310   static int replacerThunk(lol_html_streaming_sink_t* sink, void* userData);
311   kj::Promise<void> replacerThunkPromise(
312       lol_html_streaming_sink_t* sink, RegisteredReplacer& registration);
313   int replacerThunkImpl(lol_html_streaming_sink_t* sink, RegisteredReplacer& registration);
314   static void removeRegisteredReplacer(void* userData);
315  
316   // Must be constructed AFTER the registered handler vector, since the function which constructs
317   // this (buildRewriter()) modifies that vector.
318   kj::Own<lol_html_HtmlRewriter> rewriter;
319  
320   // Stores data written by lol-html, which will be periodically flushed to inner.
321   kj::Vector<kj::byte> outputBuffer;
322  
323   // Used to ensure memory usage is reported to V8 and cannot grow arbritrarily
324   jsg::ExternalMemoryAdjustment externalMemoryAdjustment;
325  
326   // True if we are currently flushing from outputBuffer to inner (in flushWrite)
327   bool flushing = false;
328  
329   // The destination for the output from lol-html
330   kj::Own<WritableStreamSink> inner;
331  
332   kj::Maybe<kj::Exception> maybeException;
333  
334   IoContext& ioContext;
335  
336   kj::Maybe<kj::WaitScope&> maybeWaitScope;
337  
338   bool canceled = false;
339  
340   kj::Maybe<jsg::Ref<jsg::AsyncContextFrame>> maybeAsyncContext;
341  
342   bool isPoisoned() {
343     // If a call to `lol-html` returned an error or propagated a user error from a handler
344     // (LOL_HTML_STOP for instance); we consider its instance as poisoned. Future calls to
345     // `lol_html_rewriter_write` and `lol_html_rewriter_end` will probably throw.
346     return maybeException != kj::none;
347   }
348  
349   void maybePoison(kj::Exception exception) {
350     // Ignore this error if maybeException is already populated -- this error is probably just a
351     // secondary effect.
352     if (maybeException == kj::none) {
353       maybeException = kj::mv(exception);
354     }
355   }
356 };
357  
358 kj::Own<lol_html_HtmlRewriter> Rewriter::buildRewriter(jsg::Lock& js,
359     kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
360     kj::ArrayPtr<const char> encoding,
361     Rewriter& rewriter) {
362   auto builder = LOL_HTML_OWN(rewriter_builder, lol_html_rewriter_builder_new());
363  
364   auto registerCallback = [&](ElementCallbackFunction& callback) {
365     auto registeredHandler = RegisteredHandler{rewriter, callback.addRef(js)};
366     return rewriter.registeredHandlers.add(kj::heap(kj::mv(registeredHandler))).get();
367   };
368  
369   for (auto& handlers: unregisteredHandlers) {
370     KJ_SWITCH_ONEOF(handlers) {
371       KJ_CASE_ONEOF(elementHandlers, UnregisteredElementHandlers) {
372         auto element = elementHandlers.element.map(registerCallback);
373         auto comments = elementHandlers.comments.map(registerCallback);
374         auto text = elementHandlers.text.map(registerCallback);
375  
376         check(lol_html_rewriter_builder_add_element_content_handlers(builder,
377             elementHandlers.selector, element == kj::none ? nullptr : &Rewriter::thunk<Element>,
378             element.orDefault(nullptr), comments == kj::none ? nullptr : &Rewriter::thunk<Comment>,
379             comments.orDefault(nullptr), text == kj::none ? nullptr : &Rewriter::thunk<Text>,
380             text.orDefault(nullptr)));
381       }
382       KJ_CASE_ONEOF(documentHandlers, UnregisteredDocumentHandlers) {
383         auto doctype = documentHandlers.doctype.map(registerCallback);
384         auto comments = documentHandlers.comments.map(registerCallback);
385         auto text = documentHandlers.text.map(registerCallback);
386         auto end = documentHandlers.end.map(registerCallback);
387  
388         // Adding document content handlers cannot fail, so no need for check().
389         lol_html_rewriter_builder_add_document_content_handlers(builder,
390             doctype == kj::none ? nullptr : &Rewriter::thunk<Doctype>, doctype.orDefault(nullptr),
391             comments == kj::none ? nullptr : &Rewriter::thunk<Comment>, comments.orDefault(nullptr),
392             text == kj::none ? nullptr : &Rewriter::thunk<Text>, text.orDefault(nullptr),
393             end == kj::none ? nullptr : &Rewriter::thunk<DocumentEnd>, end.orDefault(nullptr));
394       }
395     }
396   }
397  
398   // `strict` mode will bail out from tokenization process in cases when
399   // there is no way to determine correct parsing context. Recommended
400   // setting for safety reasons.
401   bool isStrict = true;
402  
403   // Configure a maximum memory limit that `lol-html` is allowed to use and
404   // preallocate some memory for its internal buffer.
405   lol_html_memory_settings_t memorySettings = {
406     .preallocated_parsing_buffer_size = 1024, .max_allowed_memory_usage = 3 * 1024 * 1024};
407  
408   if (FeatureFlags::get(js).getEsiIncludeIsVoidTag()) {
409     return LOL_HTML_OWN(rewriter,
410         unstable_lol_html_rewriter_build_with_esi_tags(builder, encoding.begin(), encoding.size(),
411             memorySettings, &Rewriter::output, &rewriter, isStrict));
412  
413   } else {
414     return LOL_HTML_OWN(rewriter,
415         lol_html_rewriter_build(builder, encoding.begin(), encoding.size(), memorySettings,
416             &Rewriter::output, &rewriter, isStrict));
417   }
418 }
419  
420 Rewriter::Rewriter(jsg::Lock& js,
421     kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
422     kj::ArrayPtr<const char> encoding,
423     kj::Own<WritableStreamSink> inner)
424     : rewriter(buildRewriter(js, unregisteredHandlers, encoding, *this)),
425       externalMemoryAdjustment(js.getExternalMemoryAdjustment()),
426       inner(kj::mv(inner)),
427       ioContext(IoContext::current()),
428       maybeAsyncContext(jsg::AsyncContextFrame::currentRef(js)) {}
429  
430 namespace {
431  
432 // The stack size floor enforced by kj. We could go lower,
433 // but it'd always be increased to this anyway.
434 const size_t FIBER_STACK_SIZE = 1024 * 64;
435  
436 const kj::FiberPool& getFiberPool() {
437   const static kj::FiberPool FIBER_POOL(FIBER_STACK_SIZE);
438   return FIBER_POOL;
439 }
440  
441 }  // namespace
442  
443 kj::Promise<void> Rewriter::write(kj::ArrayPtr<const byte> buffer) {
444   KJ_ASSERT(maybeWaitScope == kj::none);
445   // Defer fiber creation until the event loop runs. If this promise is dropped synchronously
446   // (e.g. by Canceler::cancel() during PumpToReader destruction), no fiber is created, avoiding
447   // a KJ assertion failure when destroying an unfired fiber. Once the event loop processes this,
448   // the fiber is created and immediately fires (armDepthFirst), so cancellation works normally.
449   return kj::evalLater([this, buffer]() {
450     return getFiberPool().startFiber([this, buffer](kj::WaitScope& scope) {
451       maybeWaitScope = scope;
452       if (!isPoisoned()) {
453         // Cannot use `check()` because `finishWrite()` implements the error path.
454         auto rc = lol_html_rewriter_write(rewriter, buffer.asChars().begin(), buffer.size());
455         tryHandleCancellation(rc);
456         if (rc == -1) {
457           maybePoison(getLastError());
458         }
459       }
460       return finishWrite();
461     });
462   });
463 }
464  
465 kj::Promise<void> Rewriter::write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) {
466   KJ_ASSERT(maybeWaitScope == kj::none);
467   return kj::evalLater([this, pieces]() {
468     return getFiberPool().startFiber([this, pieces](kj::WaitScope& scope) {
469       maybeWaitScope = scope;
470       if (!isPoisoned()) {
471         for (auto bytes: pieces) {
472           auto chars = bytes.asChars();
473           // Cannot use `check()` because `finishWrite()` implements the error path.
474           auto rc = lol_html_rewriter_write(rewriter, chars.begin(), chars.size());
475           tryHandleCancellation(rc);
476           if (rc == -1) {
477             maybePoison(getLastError());
478             // A handler threw an exception; stop calling `lol_html_rewriter_write()`.
479             break;
480           }
481         }
482       }
483       return finishWrite();
484     });
485   });
486 }
487  
488 kj::Promise<void> Rewriter::end() {
489   KJ_ASSERT(maybeWaitScope == kj::none);
490   return kj::evalLater([this]() {
491     return getFiberPool().startFiber([this](kj::WaitScope& scope) {
492       maybeWaitScope = scope;
493       if (!isPoisoned()) {
494         // Cannot use `check()` because `finishWrite()` implements the error path.
495         auto rc = lol_html_rewriter_end(rewriter);
496         tryHandleCancellation(rc);
497         if (rc == -1) {
498           maybePoison(getLastError());
499         }
500       }
501       return finishWrite().then([this]() { return inner->end(); });
502     });
503   });
504 }
505  
506 void Rewriter::abort(kj::Exception reason) {
507   // End the rewriter and forward the error to the wrapped output stream.
508   maybeException = reason.clone();
509  
510   inner->abort(kj::mv(reason));
511 }
512  
513 kj::Promise<void> Rewriter::finishWrite() {
514   maybeWaitScope = kj::none;
515   return flushWrite();
516 }
517  
518 kj::Promise<void> Rewriter::flushWrite() {
519   KJ_ASSERT(!flushing);
520  
521   if (!outputBuffer.empty()) {
522     KJ_DEFER({
523       externalMemoryAdjustment.set(0);
524       outputBuffer.clear();
525       flushing = false;
526     });
527  
528     flushing = true;
529     co_await inner->write(outputBuffer);
530   }
531  
532   KJ_IF_SOME(exception, maybeException) {
533     inner->abort(exception.clone());
534     kj::throwFatalException(exception.clone());
535   }
536 }
537 template <typename T, typename CType>
538 lol_html_rewriter_directive_t Rewriter::thunk(CType* content, void* userdata) {
539   auto& registration = *reinterpret_cast<RegisteredHandler*>(userdata);
540   return registration.rewriter.thunkImpl<T>(content, registration);
541 }
542  
543 template <typename T, typename CType>
544 lol_html_rewriter_directive_t Rewriter::thunkImpl(
545     CType* content, RegisteredHandler& registeredHandler) {
546   if (isPoisoned()) {
547     // Handlers disabled due to exception.
548     KJ_LOG(ERROR, "poisoned rewriter should not be able to call handlers");
549     return LOL_HTML_STOP;
550   }
551  
552   try {
553     KJ_IF_SOME(exception, kj::runCatchingExceptions([&] {
554       // V8 has a thread local pointer that points to where the stack limit is on this thread which
555       // is tested for overflows when we enter any JS code. However since we're running in a fiber
556       // here, we're in an entirely different stack that V8 doesn't know about, so it gets confused
557       // and may think we've overflowed our stack. evalLater will run thunkPromise on the main stack
558       // to keep V8 from getting confused.
559       auto promise = kj::evalLater([&]() { return thunkPromise<T>(content, registeredHandler); });
560       promise.wait(KJ_ASSERT_NONNULL(maybeWaitScope));
561       flushWrite().wait(KJ_ASSERT_NONNULL(maybeWaitScope));
562     })) {
563       // Exception in handler. We need to abort the streaming parser, but can't do so just yet: we
564       // need to unwind the stack because we're probably still inside a cool_thing_rewriter_write().
565       // We can't unwind with an exception across the Rust/C++ boundary, so instead we'll keep this
566       // exception around and disable all later handlers.
567       maybePoison(kj::mv(exception));
568       return LOL_HTML_STOP;
569     }
570   } catch (kj::CanceledException) {
571     // The fiber is being canceled. Same as runCatchingExceptions, we need to abort the parser,
572     // but can't since we're still inside cool_thing_rewriter_write(). This isn't handled by
573     // runCatchingExceptions since CanceledException isn't a kj exception, and we wouldn't want
574     // runCatchingExceptions to handle it anyway. We set canceled to true and once we leave Rust,
575     // we rethrow it to properly cancel the fiber.
576     canceled = true;
577     return LOL_HTML_STOP;
578   }
579   return LOL_HTML_CONTINUE;
580 }
581  
582 void Rewriter::removeEndTagHandler(RegisteredHandler& handler) {
583   auto size = registeredEndTagHandlers.size();
584   for (auto counter = size; counter != 0; --counter) {
585     auto idx = counter - 1;
586     if (registeredEndTagHandlers[idx].get() == &handler) {
587       // equivalent of `Vec::swap_remove` in Rust
588       if (counter != size) {
589         registeredEndTagHandlers[idx] = kj::mv(registeredEndTagHandlers[size - 1]);
590       }
591       registeredEndTagHandlers.removeLast();
592       break;
593     }
594   }
595 }
596  
597 template <typename T, typename CType>
598 kj::Promise<void> Rewriter::thunkPromise(CType* content, RegisteredHandler& registeredHandler) {
599   return ioContext.run(
600       [this, content, &registeredHandler](Worker::Lock& lock) -> kj::Promise<void> {
601     // We enter the AsyncContextFrame that was current when the Rewriter was created
602     // (when transform() was called). If someone wants, instead, to use the context
603     // that was current when on(...) is called, the ElementHandler can use AsyncResource
604     // (or eventually the standard AsyncContext once that lands).
605     jsg::Lock& js = lock;
606     jsg::AsyncContextFrame::Scope asyncContextScope(js, maybeAsyncContext);
607     auto jsContent = js.alloc<T>(*content, *this);
608     auto scope = HTMLRewriter::TokenScope(jsContent);
609     auto value = registeredHandler.callback(js, kj::mv(jsContent));
610  
611     if constexpr (kj::isSameType<T, EndTag>()) {
612       // TODO(someday): We can't unconditionally pop the top of `registeredEndTagHandlers`,
613       //   because that depends on https://github.com/cloudflare/lol-html/issues/110
614       //   being resolved. For now we let handles to end tag handlers tags live for the duration of
615       //   the response transformation, but eagerly release ones that we can.
616       //   In particular, note that `thunkPromise` is never called for implied end tags.
617       removeEndTagHandler(registeredHandler);
618     }
619  
620     return value.attach(kj::mv(scope));
621   });
622 }
623  
624 lol_html_streaming_handler_t Rewriter::registerReplacer(
625     jsg::Ref<ReadableStream> content, bool isHtml) {
626   auto replacer = kj::heap<RegisteredReplacer>(*this, isHtml, kj::mv(content));
627   auto userData = replacer.get();
628   registeredReplacers.insert(userData, kj::mv(replacer));
629  
630   return {
631     .user_data = userData,
632     .write_all_callback = Rewriter::replacerThunk,
633     .drop_callback = Rewriter::removeRegisteredReplacer,
634   };
635 }
636  
637 // Adapter that allows pumping a ReadableStream to a pre-established lol_html
638 // streaming sink, named `sink`. Writes of arbitrary bytes and sizes are allowed,
639 // but the content must be valid UTF-8 or lol_html will reject it.
640 class ReplacerStreamSink final: public WritableStreamSink {
641  public:
642   ReplacerStreamSink(lol_html_streaming_sink_t* sink, bool isHtml): sink(sink), isHtml(isHtml) {}
643  
644   kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override KJ_WARN_UNUSED_RESULT {
645     auto err = lol_html_streaming_sink_write_utf8_chunk(
646         sink, buffer.asChars().begin(), buffer.size(), isHtml);
647     if (err != 0) {
648       return getLastError();
649     }
650  
651     return kj::READY_NOW;
652   }
653  
654   kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override {
655     for (auto bytes: pieces) {
656       auto err = lol_html_streaming_sink_write_utf8_chunk(
657           sink, bytes.asChars().begin(), bytes.size(), isHtml);
658       if (err != 0) {
659         return getLastError();
660       }
661     }
662  
663     return kj::READY_NOW;
664   }
665  
666   kj::Promise<void> end() override KJ_WARN_UNUSED_RESULT {
667     // Nothing specific needs to be done to tell lol_html we're done writing the stream
668     return kj::READY_NOW;
669   }
670  
671   void abort(kj::Exception reason) override {
672     // Nothing specific needs to be done. The rewriter will be poisoned in replacerThunkImpl,
673     // and lol_html will bubble up the error through to write.
674   }
675  
676  private:
677   lol_html_streaming_sink_t* sink;
678   bool isHtml;
679 };
680  
681 int Rewriter::replacerThunk(lol_html_streaming_sink_t* sink, void* userData) {
682   auto& registration = *reinterpret_cast<RegisteredReplacer*>(userData);
683   return registration.rewriter.replacerThunkImpl(sink, registration);
684 }
685  
686 int Rewriter::replacerThunkImpl(
687     lol_html_streaming_sink_t* sink, RegisteredReplacer& registeredHandler) {
688   if (isPoisoned()) {
689     // Handlers disabled due to exception.
690     KJ_LOG(ERROR, "poisoned rewriter should not be able to call handlers");
691     return -1;
692   }
693  
694   try {
695     KJ_IF_SOME(exception, kj::runCatchingExceptions([&] {
696       // V8 has a thread local pointer that points to where the stack limit is on this thread which
697       // is tested for overflows when we enter any JS code. However since we're running in a fiber
698       // here, we're in an entirely different stack that V8 doesn't know about, so it gets confused
699       // and may think we've overflowed our stack. evalLater will run thunkPromise on the main stack
700       // to keep V8 from getting confused.
701       auto promise = kj::evalLater([&]() { return replacerThunkPromise(sink, registeredHandler); });
702       promise.wait(KJ_ASSERT_NONNULL(maybeWaitScope));
703     })) {
704       // Exception in handler. We need to abort the streaming parser, but can't do so just yet: we
705       // need to unwind the stack because we're probably still inside a cool_thing_rewriter_write().
706       // We can't unwind with an exception across the Rust/C++ boundary, so instead we'll keep this
707       // exception around and disable all later handlers
708       maybePoison(kj::mv(exception));
709       return -1;
710     }
711   } catch (kj::CanceledException) {
712     // The fiber is being canceled. Same as runCatchingExceptions, we need to abort the parser,
713     // but can't since we're still inside cool_thing_rewriter_write(). This isn't handled by
714     // runCatchingExceptions since CanceledException isn't a kj exception, and we wouldn't want
715     // runCatchingExceptions to handle it anyway. We set canceled to true and once we leave Rust,
716     // we rethrow it to properly cancel the fiber.
717     canceled = true;
718     return -1;
719   }
720   return 0;
721 }
722  
723 kj::Promise<void> Rewriter::replacerThunkPromise(
724     lol_html_streaming_sink_t* sink, RegisteredReplacer& registration) {
725   return ioContext.run([this, sink, &registration](Worker::Lock& lock) -> kj::Promise<void> {
726     jsg::AsyncContextFrame::Scope asyncContextScope(lock, maybeAsyncContext);
727  
728     auto streamSink = kj::heap<ReplacerStreamSink>(sink, registration.isHtml);
729     return ioContext.waitForDeferredProxy(
730         registration.stream->pumpTo(lock, kj::mv(streamSink), true));
731   });
732 }
733  
734 void Rewriter::removeRegisteredReplacer(void* userData) {
735   auto& registration = *reinterpret_cast<RegisteredReplacer*>(userData);
736   KJ_REQUIRE(registration.rewriter.registeredReplacers.erase(userData),
737       "Tried to remove replacer that was not registered");
738 }
739  
740 void Rewriter::onEndTag(lol_html_element_t* element, ElementCallbackFunction&& callback) {
741   auto registeredHandler = Rewriter::RegisteredHandler{*this, kj::mv(callback)};
742   // NOTE: this gets freed in `thunkPromise` above.
743   // TODO(someday): this uses more memory than necessary for implied end tags, which lol-html
744   // doesn't actually call `thunk` on.  LOL HTML drops the handler after it finishes transforming
745   // the current element, but this code will keep it around until the entire HTML document is
746   // transformed. It would be nice to free it directly after the handler is used; unfortunately,
747   // this isn't trivial to do since we have no idea whether there's an end tag or not. The fix for
748   // this probably needs to happen in lol-html; see #110.
749   // WARNING: if we ever start reusing the same Rewriter for multiple documents,
750   // this will cause a memory leak!
751   auto& registeredHandlerPtr = registeredEndTagHandlers.add(kj::heap(kj::mv(registeredHandler)));
752   lol_html_element_clear_end_tag_handlers(element);
753   check(lol_html_element_add_end_tag_handler(
754       element, Rewriter::thunk<EndTag>, registeredHandlerPtr.get()));
755 }
756  
757 void Rewriter::output(const char* buffer, size_t size, void* userdata) {
758   auto& rewriter = *reinterpret_cast<Rewriter*>(userdata);
759   rewriter.outputImpl(kj::asBytes(buffer, size));
760 }
761  
762 void Rewriter::outputImpl(kj::ArrayPtr<const byte> buffer) {
763   if (isPoisoned()) {
764     // Handlers disabled due to exception or running in a destructor.
765     return;
766   }
767  
768   KJ_ASSERT(!flushing);
769   externalMemoryAdjustment.adjust(buffer.size());
770   outputBuffer.addAll(buffer);
771 }
772  
773 // =======================================================================================
774 // HTMLRewriter::Token::ImplBase<CType>
775  
776 template <typename CType>
777 HTMLRewriter::Token::ImplBase<CType>::ImplBase(CType& element, Rewriter& rewriter)
778     : element(element),
779       rewriter(rewriter) {}
780 template <typename CType>
781 HTMLRewriter::Token::ImplBase<CType>::~ImplBase() noexcept(false) {}
782 template <typename CType>
783 template <auto Func, auto StreamingFunc>
784 void HTMLRewriter::Token::ImplBase<CType>::rewriteContentGeneric(
785     Content content, jsg::Optional<ContentOptions> options) {
786   auto isHtml = options.orDefault({}).html.orDefault(false);
787  
788   KJ_SWITCH_ONEOF(content) {
789     KJ_CASE_ONEOF(stringContent, kj::String) {
790       check(Func(&element, stringContent.cStr(), stringContent.size(), isHtml));
791     }
792  
793     KJ_CASE_ONEOF(streamContent, jsg::Ref<ReadableStream>) {
794       auto handler = rewriter.registerReplacer(kj::mv(streamContent), isHtml);
795       check(StreamingFunc(&element, &handler));
796     }
797  
798     KJ_CASE_ONEOF(responseContent, jsg::Ref<Response>) {
799       KJ_IF_SOME(body, responseContent->getBody()) {
800         auto handler = rewriter.registerReplacer(kj::mv(body), isHtml);
801         check(StreamingFunc(&element, &handler));
802       }
803       // Otherwise if no body, there is no replacement to make
804     }
805   }
806 }
807 // =======================================================================================
808 // Element
809  
810 Element::Element(CType& element, Rewriter& rewriter) {
811   impl.emplace(element, rewriter);
812 }
813  
814 kj::String Element::getTagName() {
815   auto tagName = LolString(lol_html_element_tag_name_get(&checkToken(impl).element));
816   return kj::str(tagName.asChars());
817 }
818  
819 void Element::setTagName(kj::String name) {
820   check(lol_html_element_tag_name_set(&checkToken(impl).element, name.cStr(), name.size()));
821 }
822  
823 bool Element::getRemoved() {
824   return lol_html_element_is_removed(&checkToken(impl).element);
825 }
826  
827 kj::StringPtr Element::getNamespaceURI() {
828   // lol-html returns a static C string, no need to handle its lifetime.
829   return lol_html_element_namespace_uri_get(&checkToken(impl).element);
830 }
831  
832 jsg::Ref<Element::AttributesIterator> Element::getAttributes(jsg::Lock& js) {
833   auto& implRef = checkToken(impl);
834  
835   auto iter = LOL_HTML_OWN(attributes_iterator, lol_html_attributes_iterator_get(&implRef.element));
836  
837   auto jsIter = js.alloc<Element::AttributesIterator>(kj::mv(iter));
838   implRef.attributesIterators.add(jsIter.addRef());
839   return kj::mv(jsIter);
840 }
841  
842 kj::Maybe<kj::String> Element::getAttribute(kj::String name) {
843   // NOTE: lol_html_element_get_attribute() returns NULL for both nonexistent attributes and for
844   //   errors, so we can't use check() here.
845   LolString attr(
846       lol_html_element_get_attribute(&checkToken(impl).element, name.cStr(), name.size()));
847   // TODO(perf): We could construct a v8::String directly here, saving a copy.
848   kj::Maybe kjAttr = attr.asKjString();
849   if (kjAttr != kj::none) {
850     return kj::mv(kjAttr);
851   }
852  
853   KJ_IF_SOME(exception, tryGetLastError()) {
854     kj::throwFatalException(kj::mv(exception));
855   }
856  
857   // No error, just doesn't exist.
858   return kj::none;
859 }
860  
861 bool Element::hasAttribute(kj::String name) {
862   return !!check(
863       lol_html_element_has_attribute(&checkToken(impl).element, name.cStr(), name.size()));
864 }
865  
866 jsg::Ref<Element> Element::setAttribute(kj::String name, kj::String value) {
867   auto& implRef = checkToken(impl);
868   check(lol_html_element_set_attribute(
869       &implRef.element, name.cStr(), name.size(), value.cStr(), value.size()));
870  
871   // Mutating attributes may cause lol-html's internal Vec to reallocate, invalidating
872   // any live iterators' pointers. We must invalidate all outstanding iterators.
873   for (auto& iter: implRef.attributesIterators) {
874     iter->invalidate();
875   }
876  
877   return JSG_THIS;
878 }
879  
880 jsg::Ref<Element> Element::removeAttribute(kj::String name) {
881   auto& implRef = checkToken(impl);
882   check(lol_html_element_remove_attribute(&implRef.element, name.cStr(), name.size()));
883  
884   // Removing attributes may shift elements in lol-html's internal Vec (via retain()),
885   // invalidating any live iterators' pointers.
886   for (auto& iter: implRef.attributesIterators) {
887     iter->invalidate();
888   }
889  
890   return JSG_THIS;
891 }
892  
893 namespace {
894 kj::String unwrapContent(Content content) {
895   return kj::mv(JSG_REQUIRE_NONNULL(content.tryGet<kj::String>(), TypeError,
896       "Replacing content in HTML comments using a ReadableStream or Response object is not "
897       "implemented. You must provide a string."));
898 }
899 }  // namespace
900  
901 jsg::Ref<Element> Element::before(Content content, jsg::Optional<ContentOptions> options) {
902   checkToken(impl)
903       .rewriteContentGeneric<lol_html_element_before, lol_html_element_streaming_before>(
904           kj::mv(content), options);
905   return JSG_THIS;
906 }
907  
908 jsg::Ref<Element> Element::after(Content content, jsg::Optional<ContentOptions> options) {
909   checkToken(impl).rewriteContentGeneric<lol_html_element_after, lol_html_element_streaming_after>(
910       kj::mv(content), options);
911   return JSG_THIS;
912 }
913  
914 jsg::Ref<Element> Element::prepend(Content content, jsg::Optional<ContentOptions> options) {
915   checkToken(impl)
916       .rewriteContentGeneric<lol_html_element_prepend, lol_html_element_streaming_prepend>(
917           kj::mv(content), options);
918   return JSG_THIS;
919 }
920  
921 jsg::Ref<Element> Element::append(Content content, jsg::Optional<ContentOptions> options) {
922   checkToken(impl)
923       .rewriteContentGeneric<lol_html_element_append, lol_html_element_streaming_append>(
924           kj::mv(content), options);
925   return JSG_THIS;
926 }
927  
928 jsg::Ref<Element> Element::replace(Content content, jsg::Optional<ContentOptions> options) {
929   checkToken(impl)
930       .rewriteContentGeneric<lol_html_element_replace, lol_html_element_streaming_replace>(
931           kj::mv(content), options);
932   return JSG_THIS;
933 }
934  
935 jsg::Ref<Element> Element::setInnerContent(Content content, jsg::Optional<ContentOptions> options) {
936   checkToken(impl)
937       .rewriteContentGeneric<lol_html_element_set_inner_content,
938           lol_html_element_streaming_set_inner_content>(kj::mv(content), options);
939   return JSG_THIS;
940 }
941  
942 jsg::Ref<Element> Element::remove() {
943   lol_html_element_remove(&checkToken(impl).element);
944   return JSG_THIS;
945 }
946  
947 jsg::Ref<Element> Element::removeAndKeepContent() {
948   lol_html_element_remove_and_keep_content(&checkToken(impl).element);
949   return JSG_THIS;
950 }
951  
952 void Element::onEndTag(ElementCallbackFunction&& callback) {
953   auto& knownImpl = checkToken(impl);
954   knownImpl.rewriter.onEndTag(&knownImpl.element, kj::mv(callback));
955 }
956  
957 EndTag::EndTag(CType& endTag, Rewriter& rewriter) {
958   impl.emplace(endTag, rewriter);
959 }
960  
961 void EndTag::htmlContentScopeEnd() {
962   impl = kj::none;
963 }
964  
965 kj::String EndTag::getName() {
966   auto text = LolString(lol_html_end_tag_name_get(&checkToken(impl).element));
967   return kj::str(text.asChars());
968 }
969  
970 void EndTag::setName(kj::String text) {
971   check(lol_html_end_tag_name_set(&checkToken(impl).element, text.cStr(), text.size()));
972 }
973  
974 jsg::Ref<EndTag> EndTag::before(Content content, jsg::Optional<ContentOptions> options) {
975   checkToken(impl)
976       .rewriteContentGeneric<lol_html_end_tag_before, lol_html_end_tag_streaming_before>(
977           kj::mv(content), kj::mv(options));
978   return JSG_THIS;
979 }
980  
981 jsg::Ref<EndTag> EndTag::after(Content content, jsg::Optional<ContentOptions> options) {
982   checkToken(impl).rewriteContentGeneric<lol_html_end_tag_after, lol_html_end_tag_streaming_after>(
983       kj::mv(content), kj::mv(options));
984   return JSG_THIS;
985 }
986  
987 jsg::Ref<EndTag> EndTag::remove() {
988   lol_html_end_tag_remove(&checkToken(impl).element);
989   return JSG_THIS;
990 }
991  
992 void Element::htmlContentScopeEnd() {
993   impl = kj::none;
994 }
995  
996 Element::Impl::~Impl() noexcept(false) {
997   for (auto& jsIter: attributesIterators) {
998     static_cast<HTMLRewriter::Token&>(*jsIter).htmlContentScopeEnd();
999   }
1000 }
1001  
1002 // =======================================================================================
1003 // Element::AttributesIterator
1004  
1005 Element::AttributesIterator::AttributesIterator(kj::Own<CType> iter): impl(kj::mv(iter)) {}
1006  
1007 jsg::Ref<Element::AttributesIterator> Element::AttributesIterator::self() {
1008   return JSG_THIS;
1009 }
1010  
1011 Element::AttributesIterator::Next Element::AttributesIterator::next() {
1012   // If the element's attributes were modified (via setAttribute/removeAttribute) while this
1013   // iterator was live, the underlying lol-html iterator holds stale pointers into a potentially
1014   // reallocated Vec. Continuing to iterate would be a use-after-free.
1015   JSG_REQUIRE(!mutatedDuringIteration, Error,
1016       "The attributes of this element have been modified during iteration. "
1017       "You must create a new iterator after modifying attributes.");
1018  
1019   // NOTE: lol_html_attribute_t doesn't need to be freed.
1020   auto* attribute = lol_html_attributes_iterator_next(checkToken(impl));
1021   if (attribute == nullptr) {
1022     // End of iteration.
1023     // TODO(someday): Eagerly deallocate. Can't seem to nullify the Own without also nullifying the
1024     //   enclosing Maybe, however.
1025     return {true, kj::none};
1026   }
1027  
1028   auto name = LolString(lol_html_attribute_name_get(attribute));
1029   auto value = LolString(lol_html_attribute_value_get(attribute));
1030  
1031   return {false, kj::arr(kj::str(name.asChars()), kj::str(value.asChars()))};
1032 }
1033  
1034 void Element::AttributesIterator::invalidate() {
1035   mutatedDuringIteration = true;
1036   // Also release the underlying lol-html iterator since it's no longer safe to use.
1037   impl = kj::none;
1038 }
1039  
1040 void Element::AttributesIterator::htmlContentScopeEnd() {
1041   // Clear the mutation flag so that after scope end, the "content token is no longer valid"
1042   // error (from checkToken) takes precedence over the mutation error.
1043   mutatedDuringIteration = false;
1044   impl = kj::none;
1045 }
1046  
1047 // =======================================================================================
1048 // Comment
1049  
1050 Comment::Comment(CType& comment, Rewriter&): impl(comment) {}
1051  
1052 kj::String Comment::getText() {
1053   auto text = LolString(lol_html_comment_text_get(&checkToken(impl)));
1054   return kj::str(text.asChars());
1055 }
1056  
1057 void Comment::setText(kj::String text) {
1058   check(lol_html_comment_text_set(&checkToken(impl), text.cStr(), text.size()));
1059 }
1060  
1061 bool Comment::getRemoved() {
1062   // NOTE: No error checking seems required by this function -- it returns a bool directly.
1063   return lol_html_comment_is_removed(&checkToken(impl));
1064 }
1065  
1066 jsg::Ref<Comment> Comment::before(Content content, jsg::Optional<ContentOptions> options) {
1067   // TODO(someday): If lol-html adds support for streaming replacements for comments, this
1068   // function will need to be updated.
1069   auto stringContent = unwrapContent(kj::mv(content));
1070   check(lol_html_comment_before(&checkToken(impl), stringContent.cStr(), stringContent.size(),
1071       options.orDefault({}).html.orDefault(false)));
1072  
1073   return JSG_THIS;
1074 }
1075  
1076 jsg::Ref<Comment> Comment::after(Content content, jsg::Optional<ContentOptions> options) {
1077   // TODO(someday): If lol-html adds support for streaming replacements for comments, this
1078   // function will need to be updated.
1079   auto stringContent = unwrapContent(kj::mv(content));
1080   check(lol_html_comment_after(&checkToken(impl), stringContent.cStr(), stringContent.size(),
1081       options.orDefault({}).html.orDefault(false)));
1082  
1083   return JSG_THIS;
1084 }
1085  
1086 jsg::Ref<Comment> Comment::replace(Content content, jsg::Optional<ContentOptions> options) {
1087   // TODO(someday): If lol-html adds support for streaming replacements for comments, this
1088   // function will need to be updated.
1089   auto stringContent = unwrapContent(kj::mv(content));
1090   check(lol_html_comment_replace(&checkToken(impl), stringContent.cStr(), stringContent.size(),
1091       options.orDefault({}).html.orDefault(false)));
1092  
1093   return JSG_THIS;
1094 }
1095  
1096 jsg::Ref<Comment> Comment::remove() {
1097   lol_html_comment_remove(&checkToken(impl));
1098  
1099   return JSG_THIS;
1100 }
1101  
1102 void Comment::htmlContentScopeEnd() {
1103   impl = kj::none;
1104 }
1105  
1106 // =======================================================================================
1107 // Text
1108  
1109 Text::Text(CType& text, Rewriter& rewriter) {
1110   impl.emplace(text, rewriter);
1111 }
1112  
1113 kj::String Text::getText() {
1114   auto content = lol_html_text_chunk_content_get(&checkToken(impl).element);
1115   return kj::heapString(content.data, content.len);
1116 }
1117  
1118 bool Text::getLastInTextNode() {
1119   // NOTE: No error checking seems required by this function -- it returns a bool directly.
1120   return lol_html_text_chunk_is_last_in_text_node(&checkToken(impl).element);
1121 }
1122  
1123 bool Text::getRemoved() {
1124   // NOTE: No error checking seems required by this function -- it returns a bool directly.
1125   return lol_html_text_chunk_is_removed(&checkToken(impl).element);
1126 }
1127  
1128 jsg::Ref<Text> Text::before(Content content, jsg::Optional<ContentOptions> options) {
1129   checkToken(impl)
1130       .rewriteContentGeneric<lol_html_text_chunk_before, lol_html_text_chunk_streaming_before>(
1131           kj::mv(content), kj::mv(options));
1132   return JSG_THIS;
1133 }
1134  
1135 jsg::Ref<Text> Text::after(Content content, jsg::Optional<ContentOptions> options) {
1136   checkToken(impl)
1137       .rewriteContentGeneric<lol_html_text_chunk_after, lol_html_text_chunk_streaming_after>(
1138           kj::mv(content), kj::mv(options));
1139   return JSG_THIS;
1140 }
1141  
1142 jsg::Ref<Text> Text::replace(Content content, jsg::Optional<ContentOptions> options) {
1143   checkToken(impl)
1144       .rewriteContentGeneric<lol_html_text_chunk_replace, lol_html_text_chunk_streaming_replace>(
1145           kj::mv(content), kj::mv(options));
1146   return JSG_THIS;
1147 }
1148  
1149 jsg::Ref<Text> Text::remove() {
1150   lol_html_text_chunk_remove(&checkToken(impl).element);
1151  
1152   return JSG_THIS;
1153 }
1154  
1155 void Text::htmlContentScopeEnd() {
1156   impl = kj::none;
1157 }
1158  
1159 // =======================================================================================
1160 // Doctype
1161  
1162 Doctype::Doctype(CType& doctype, Rewriter&): impl(doctype) {}
1163  
1164 kj::Maybe<kj::String> Doctype::getName() {
1165   LolString name(lol_html_doctype_name_get(&checkToken(impl)));
1166   return name.asKjString();
1167 }
1168  
1169 kj::Maybe<kj::String> Doctype::getPublicId() {
1170   LolString publicId(lol_html_doctype_public_id_get(&checkToken(impl)));
1171   return publicId.asKjString();
1172 }
1173  
1174 kj::Maybe<kj::String> Doctype::getSystemId() {
1175   LolString systemId(lol_html_doctype_system_id_get(&checkToken(impl)));
1176   return systemId.asKjString();
1177 }
1178  
1179 void Doctype::htmlContentScopeEnd() {
1180   impl = kj::none;
1181 }
1182  
1183 // =======================================================================================
1184 // DocumentEnd
1185  
1186 DocumentEnd::DocumentEnd(CType& documentEnd, Rewriter&): impl(documentEnd) {}
1187  
1188 jsg::Ref<DocumentEnd> DocumentEnd::append(Content content, jsg::Optional<ContentOptions> options) {
1189   // TODO(someday): If lol-html adds support for streaming replacements for the document end,
1190   // this function will need to be updated.
1191   auto stringContent = unwrapContent(kj::mv(content));
1192   check(lol_html_doc_end_append(&checkToken(impl), stringContent.cStr(), stringContent.size(),
1193       options.orDefault({}).html.orDefault(false)));
1194  
1195   return JSG_THIS;
1196 }
1197  
1198 void DocumentEnd::htmlContentScopeEnd() {
1199   impl = kj::none;
1200 }
1201  
1202 // =======================================================================================
1203 // HTMLRewriter
1204  
1205 struct HTMLRewriter::Impl {
1206   // The list of handlers added to this builder.
1207   kj::Vector<UnregisteredElementOrDocumentHandlers> unregisteredHandlers;
1208   // TODO(perf): It'd be nice to eagerly register handlers on the native builder object. However,
1209   //   currently lol-html rewriters are inextricably linked to the builders which created them,
1210   //   and this has concurrency and reentrancy ramifications: two rewriters built from the same
1211   //   builder require synchronization to access safely, and their callbacks must not use the
1212   //   builder which created them, lest the process deadlock.
1213   //
1214   //   In the meantime, we keep this list of handlers around and "replay" their registration, in
1215   //   order, on the builder object that we create inside of .transform().
1216  
1217   JSG_MEMORY_INFO(HTMLRewriter::Impl) {
1218     for (const auto& handlers: unregisteredHandlers) {
1219       KJ_SWITCH_ONEOF(handlers) {
1220         KJ_CASE_ONEOF(h, UnregisteredElementHandlers) {
1221           tracker.trackField(nullptr, h);
1222         }
1223         KJ_CASE_ONEOF(h, UnregisteredDocumentHandlers) {
1224           tracker.trackField(nullptr, h);
1225         }
1226       }
1227     }
1228   }
1229 };
1230  
1231 HTMLRewriter::HTMLRewriter(): impl(kj::heap<Impl>()) {}
1232 HTMLRewriter::~HTMLRewriter() noexcept(false) {}
1233  
1234 void HTMLRewriter::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
1235   tracker.trackField("impl", impl);
1236 }
1237  
1238 jsg::Ref<HTMLRewriter> HTMLRewriter::constructor(jsg::Lock& js) {
1239   return js.alloc<HTMLRewriter>();
1240 }
1241  
1242 jsg::Ref<HTMLRewriter> HTMLRewriter::on(
1243     kj::String stringSelector, ElementContentHandlers&& handlers) {
1244   kj::Own<lol_html_Selector> selector =
1245       LOL_HTML_OWN(selector, lol_html_selector_parse(stringSelector.cStr(), stringSelector.size()));
1246  
1247   impl->unregisteredHandlers.add(UnregisteredElementHandlers{
1248     kj::mv(selector), kj::mv(handlers.element), kj::mv(handlers.comments), kj::mv(handlers.text)});
1249  
1250   return JSG_THIS;
1251 }
1252  
1253 jsg::Ref<HTMLRewriter> HTMLRewriter::onDocument(DocumentContentHandlers&& handlers) {
1254   impl->unregisteredHandlers.add(UnregisteredDocumentHandlers{kj::mv(handlers.doctype),
1255     kj::mv(handlers.comments), kj::mv(handlers.text), kj::mv(handlers.end)});
1256  
1257   return JSG_THIS;
1258 }
1259  
1260 jsg::Ref<Response> HTMLRewriter::transform(jsg::Lock& js, jsg::Ref<Response> response) {
1261  
1262   JSG_REQUIRE(response->getType() != "error"_kj, TypeError,
1263       "HTMLRewriter cannot transform an error response");
1264  
1265   auto maybeInput = response->getBody();
1266  
1267   if (maybeInput == kj::none) {
1268     // That was easy!
1269     return kj::mv(response);
1270   }
1271  
1272   auto& ioContext = IoContext::current();
1273  
1274   auto pipe = newIdentityPipe();
1275   response = Response::constructor(
1276       js, kj::Maybe(js.alloc<ReadableStream>(ioContext, kj::mv(pipe.in))), kj::mv(response));
1277  
1278   kj::String ownContentType;
1279   kj::String encoding = kj::str("utf-8");
1280   KJ_IF_SOME(contentType,
1281       response->getHeaders(js)->getCommon(js, capnp::CommonHeaderName::CONTENT_TYPE)) {
1282     // TODO(cleanup): readContentTypeParameter can be replaced with using
1283     // workerd/util/mimetype.h directly.
1284     KJ_IF_SOME(charset, readContentTypeParameter(contentType, "charset")) {
1285       ownContentType = kj::mv(contentType);
1286       encoding = kj::mv(charset);
1287     }
1288   }
1289  
1290   auto rewriter = kj::heap<Rewriter>(js, impl->unregisteredHandlers, encoding, kj::mv(pipe.out));
1291  
1292   // NOTE: Avoid throwing any exceptions after initiating the pump below. This makes
1293   //   the input response object disturbed (response.bodyUsed === true), which should only happen
1294   //   after we know that nothing else (like invalid encoding) could cause an exception.
1295  
1296   // Drive and flush the parser asynchronously.
1297   ioContext.addTask(
1298       ioContext
1299           .waitForDeferredProxy(KJ_ASSERT_NONNULL(maybeInput)->pumpTo(js, kj::mv(rewriter), true))
1300           .catch_([](kj::Exception&& e) {
1301     // Errors in pumpTo() are already propagated to the destination stream. We don't want to
1302     // throw them from here since it'll cause an uncaught exception to be reported via taskFailed(),
1303     // which would poison the IoContext even though the application may have handled the error.
1304   }));
1305  
1306   // TODO(soon): EW-2025 Make Rewriter a proper wrapper object and put it in hidden property on the
1307   //   response so the GC can find the handlers which Rewriter co-owns.
1308   return kj::mv(response);
1309 }
1310  
1311 void HTMLRewriter::visitForGc(jsg::GcVisitor& visitor) {
1312   for (auto& handlers: impl->unregisteredHandlers) {
1313     KJ_SWITCH_ONEOF(handlers) {
1314       KJ_CASE_ONEOF(elementHandlers, UnregisteredElementHandlers) {
1315         visitor.visit(elementHandlers);
1316       }
1317       KJ_CASE_ONEOF(documentHandlers, UnregisteredDocumentHandlers) {
1318         visitor.visit(documentHandlers);
1319       }
1320     }
1321   }
1322 }
1323  
1324 }  // namespace workerd::api