// Copyright (c) 2017-2022 Cloudflare, Inc.
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
//     https://opensource.org/licenses/Apache-2.0

#include "html-rewriter.h"

#include "util.h"

#include <workerd/api/streams/common.h>
#include <workerd/api/streams/identity-transform-stream.h>
#include <workerd/io/features.h>
#include <workerd/io/io-context.h>

#include <lol_html.h>

struct lol_html_HtmlRewriter {};
struct lol_html_HtmlRewriterBuilder {};
struct lol_html_AttributesIterator {};
struct lol_html_Selector {};
// TODO(cleanup): These are defined internally in lol-html, but kj::Own<T> needs to check whether
//   or not T is polymorphic, so here's a dummy definition.

namespace workerd::api {

namespace {

// =======================================================================================
// RAII helpers for lol-html

// RAII helper for lol-html types which are managed by pointers and have straightforward _free()
// functions.
template <typename T, void (*lolhtmlFree)(T*)>
class LolHtmlDisposer: public kj::Disposer {
 public:
  static const LolHtmlDisposer INSTANCE;

 protected:
  void disposeImpl(void* pointer) const override {
    lolhtmlFree(reinterpret_cast<T*>(pointer));
  }
};

template <typename T, void (*lolhtmlFree)(T*)>
const LolHtmlDisposer<T, lolhtmlFree> LolHtmlDisposer<T, lolhtmlFree>::INSTANCE;

#define LOL_HTML_OWN(name, ...)                                                                    \
  ({                                                                                               \
    using T = lol_html_##name##_t;                                                                 \
    constexpr auto* lolhtmlFree = lol_html_##name##_free;                                          \
    kj::Own<T>(&check(__VA_ARGS__), LolHtmlDisposer<T, lolhtmlFree>::INSTANCE);                    \
  })

// RAII helper for lol_html_str_t.
//
// We cannot use a kj::Own<T> because lol_html_str_t is a struct, not a pointer, so instead we
// have this LolString RAII wrapper.
//
// Use `kj::str(LolString.asChars())` to allocate your own copy of a LolString.
class LolString {
 public:
  explicit LolString(lol_html_str_t s): chars(s.data, s.len) {}
  ~LolString() noexcept(false) {
    lol_html_str_free({chars.begin(), chars.size()});
  }
  KJ_DISALLOW_COPY(LolString);

  kj::ArrayPtr<const char> asChars() const {
    return chars;
  }

  kj::Maybe<kj::String> asKjString() {
    if (chars.begin() != nullptr) {
      return kj::str(chars);
    } else {
      return kj::none;
    }
  }

 private:
  kj::ArrayPtr<const char> chars;
};

// =======================================================================================
// Error checking for lol-html

kj::Maybe<kj::Exception> tryGetLastError() {
  auto maybeErrorString = lol_html_take_last_error();
  if (maybeErrorString.data == nullptr) {
    return kj::none;
  }
  auto errorString = LolString(maybeErrorString);
  return kj::Exception(kj::Exception::Type::FAILED, __FILE__, __LINE__,
      kj::str(JSG_EXCEPTION(TypeError) ": Parser error: ", errorString.asChars()));
}

void discardLastError() {
  auto drop = LolString(lol_html_take_last_error());
}

kj::Exception getLastError() {
  return KJ_REQUIRE_NONNULL(tryGetLastError(),
      "lol-html reported error through return value, but lol_html_take_last_error() is null");
}

int check(int rc) {
  if (rc == -1) {
    kj::throwFatalException(getLastError());
  }
  return rc;
}

template <typename T>
[[nodiscard]] T& check(T* ptr) {
  // Nodiscard, because this function typically checks references that must later be freed.
  if (ptr == nullptr) {
    kj::throwFatalException(getLastError());
  }
  return *ptr;
}

// Helper function to determine if a content token is still valid. Each content token has an
// implementation object inside a Maybe -- when HTMLRewriter::TokenScope (defined below)
// gets destroyed, that Maybe gets nullified, and the content token becomes a dead, useless,
// JavaScript object occupying space, waiting to get garbage collected.
//
// In other words, if you try to access a content token (Element, Text, etc.) outside of a
// content handler, you're gonna get this exception.
template <typename T>
decltype(auto) checkToken(kj::Maybe<T>& impl) {
  return JSG_REQUIRE_NONNULL(impl, TypeError,
      "This content token is no longer valid. Content tokens are only valid "
      "during the execution of the relevant content handler.");
}

}  // namespace

// =======================================================================================
// HTMLRewriter::TokenScope

class HTMLRewriter::TokenScope {
 public:
  template <typename T>
  explicit TokenScope(jsg::Ref<T>& value): contentToken(value.addRef()) {}
  ~TokenScope() noexcept(false) {
    KJ_IF_SOME(token, contentToken) {
      token->htmlContentScopeEnd();
    }
  }
  TokenScope(TokenScope&& o): contentToken(kj::mv(o.contentToken)) {
    o.contentToken = kj::none;
  }
  KJ_DISALLOW_COPY(TokenScope);

 private:
  kj::Maybe<jsg::Ref<HTMLRewriter::Token>> contentToken;
};

namespace {

// =======================================================================================
// Rewriter
using ElementCallbackFunction = HTMLRewriter::ElementCallbackFunction;

struct UnregisteredElementHandlers {
  kj::Own<lol_html_Selector> selector;

  // The actual handler functions. We store them as jsg::Values for compatibility with GcVisitor.

  jsg::Optional<ElementCallbackFunction> element;
  jsg::Optional<ElementCallbackFunction> comments;
  jsg::Optional<ElementCallbackFunction> text;

  void visitForGc(jsg::GcVisitor& visitor) {
    visitor.visit(element, comments, text);
  }

  JSG_MEMORY_INFO(UnregisteredElementHandlers) {
    tracker.trackField("element", element);
    tracker.trackField("comments", comments);
    tracker.trackField("text", text);
  }
};

struct UnregisteredDocumentHandlers {

  // The actual handler functions. We store them as jsg::Values for compatibility with GcVisitor.

  jsg::Optional<ElementCallbackFunction> doctype;
  jsg::Optional<ElementCallbackFunction> comments;
  jsg::Optional<ElementCallbackFunction> text;
  jsg::Optional<ElementCallbackFunction> end;

  // The `this` object used to call the handler functions.

  void visitForGc(jsg::GcVisitor& visitor) {
    visitor.visit(doctype, comments, text, end);
  }

  JSG_MEMORY_INFO(UnregisteredDocumentHandlers) {
    tracker.trackField("doctype", doctype);
    tracker.trackField("comments", comments);
    tracker.trackField("text", text);
    tracker.trackField("end", end);
  }
};

using UnregisteredElementOrDocumentHandlers =
    kj::OneOf<UnregisteredElementHandlers, UnregisteredDocumentHandlers>;

}  // namespace

// Wrapper around an actual rewriter (streaming parser).
class Rewriter final: public WritableStreamSink {
 public:
  explicit Rewriter(jsg::Lock& js,
      kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
      kj::ArrayPtr<const char> encoding,
      kj::Own<WritableStreamSink> inner);
  KJ_DISALLOW_COPY_AND_MOVE(Rewriter);

  // WritableStreamSink implementation. The input body pumpTo() operation calls these.
  kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override;
  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override;
  kj::Promise<void> end() override;
  void abort(kj::Exception reason) override;

  // Implementation for `Element::onEndTag` to avoid exposing private details of Rewriter.
  void onEndTag(lol_html_element_t* element, ElementCallbackFunction&& callback);

  lol_html_streaming_handler_t registerReplacer(jsg::Ref<ReadableStream> content, bool isHtml);

  ~Rewriter() {
    KJ_ASSERT(registeredReplacers.size() == 0, "Some replacers were leaked by lol-html");
  }

 private:
  // Wait for the write promise (if any) produced by our `output()` callback, then, if there is a
  // stored exception, abort the wrapped WritableStreamSink with it, then return the exception.
  // Otherwise, just return.
  kj::Promise<void> finishWrite();

  kj::Promise<void> flushWrite();

  static kj::Own<lol_html_HtmlRewriter> buildRewriter(jsg::Lock& js,
      kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
      kj::ArrayPtr<const char> encoding,
      Rewriter& rewriterWrapper);

  static void output(const char* buffer, size_t size, void* userdata);
  void outputImpl(kj::ArrayPtr<const byte> buffer);

  void tryHandleCancellation(int rc) {
    if (canceled) {
      canceled = false;

      // We canceled this, which means we used LOL_HTML_STOP. That means we have an error sitting
      // in the error buffer in the lol-html C API. Let's make sure our return code is -1 and get
      // rid of that error value to make sure nobody picks it up later on accident and thinks an
      // error occurred.
      KJ_ASSERT(rc == -1);
      discardLastError();

      throw kj::CanceledException{};
    }
  }

  friend class ::workerd::api::HTMLRewriter;

  // Keeps track of streams currently being used as replacement content for tokens.
  // lol-html will invoke replacerThunk with a pointer to the RegisteredReplacer to be used.
  class RegisteredReplacer {
   public:
    Rewriter& rewriter;
    bool isHtml;
    jsg::Ref<ReadableStream> stream;
  };

  struct RegisteredHandler {
    // A back-reference to the rewriter which owns this particular registered handler.
    Rewriter& rewriter;

    ElementCallbackFunction callback;
  };

  kj::Vector<kj::Own<RegisteredHandler>> registeredHandlers;
  // TODO(perf): Don't store Owns. We need to pass stable pointers as the userdata parameter to
  //   lol_html_rewriter_builder_add_*_content_handlers(), but don't have a really easy way to
  //   know precisely how many handlers we're going to register beforehand, so we need a vector. But
  //   vectors can grow, moving their objects around, invalidating pointers into their storage.

  // This is separate from `registeredHandlers` so we can delete them more eagerly when EndTags are
  // destroyed, and not have to look through all other handlers.
  kj::Vector<kj::Own<RegisteredHandler>> registeredEndTagHandlers;
  // TODO(perf) Don't store Owns, same as `registeredHandlers` above.

  template <typename T, typename CType = T::CType>
  static lol_html_rewriter_directive_t thunk(CType* content, void* userdata);
  template <typename T, typename CType = T::CType>
  lol_html_rewriter_directive_t thunkImpl(CType* content, RegisteredHandler& registration);
  template <typename T, typename CType = T::CType>
  kj::Promise<void> thunkPromise(CType* content, RegisteredHandler& registration);

  // Eagerly free this handler. Should only be called if we're confident the handler will never be
  // used again.
  void removeEndTagHandler(RegisteredHandler& registration);

  // Field stores a list of readable streams that are being used by lol-html for replacements
  kj::HashMap<void*, kj::Own<RegisteredReplacer>> registeredReplacers;

  static int replacerThunk(lol_html_streaming_sink_t* sink, void* userData);
  kj::Promise<void> replacerThunkPromise(
      lol_html_streaming_sink_t* sink, RegisteredReplacer& registration);
  int replacerThunkImpl(lol_html_streaming_sink_t* sink, RegisteredReplacer& registration);
  static void removeRegisteredReplacer(void* userData);

  // Must be constructed AFTER the registered handler vector, since the function which constructs
  // this (buildRewriter()) modifies that vector.
  kj::Own<lol_html_HtmlRewriter> rewriter;

  // Stores data written by lol-html, which will be periodically flushed to inner.
  kj::Vector<kj::byte> outputBuffer;

  // Used to ensure memory usage is reported to V8 and cannot grow arbritrarily
  jsg::ExternalMemoryAdjustment externalMemoryAdjustment;

  // True if we are currently flushing from outputBuffer to inner (in flushWrite)
  bool flushing = false;

  // The destination for the output from lol-html
  kj::Own<WritableStreamSink> inner;

  kj::Maybe<kj::Exception> maybeException;

  IoContext& ioContext;

  kj::Maybe<kj::WaitScope&> maybeWaitScope;

  bool canceled = false;

  kj::Maybe<jsg::Ref<jsg::AsyncContextFrame>> maybeAsyncContext;

  bool isPoisoned() {
    // If a call to `lol-html` returned an error or propagated a user error from a handler
    // (LOL_HTML_STOP for instance); we consider its instance as poisoned. Future calls to
    // `lol_html_rewriter_write` and `lol_html_rewriter_end` will probably throw.
    return maybeException != kj::none;
  }

  void maybePoison(kj::Exception exception) {
    // Ignore this error if maybeException is already populated -- this error is probably just a
    // secondary effect.
    if (maybeException == kj::none) {
      maybeException = kj::mv(exception);
    }
  }
};

kj::Own<lol_html_HtmlRewriter> Rewriter::buildRewriter(jsg::Lock& js,
    kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
    kj::ArrayPtr<const char> encoding,
    Rewriter& rewriter) {
  auto builder = LOL_HTML_OWN(rewriter_builder, lol_html_rewriter_builder_new());

  auto registerCallback = [&](ElementCallbackFunction& callback) {
    auto registeredHandler = RegisteredHandler{rewriter, callback.addRef(js)};
    return rewriter.registeredHandlers.add(kj::heap(kj::mv(registeredHandler))).get();
  };

  for (auto& handlers: unregisteredHandlers) {
    KJ_SWITCH_ONEOF(handlers) {
      KJ_CASE_ONEOF(elementHandlers, UnregisteredElementHandlers) {
        auto element = elementHandlers.element.map(registerCallback);
        auto comments = elementHandlers.comments.map(registerCallback);
        auto text = elementHandlers.text.map(registerCallback);

        check(lol_html_rewriter_builder_add_element_content_handlers(builder,
            elementHandlers.selector, element == kj::none ? nullptr : &Rewriter::thunk<Element>,
            element.orDefault(nullptr), comments == kj::none ? nullptr : &Rewriter::thunk<Comment>,
            comments.orDefault(nullptr), text == kj::none ? nullptr : &Rewriter::thunk<Text>,
            text.orDefault(nullptr)));
      }
      KJ_CASE_ONEOF(documentHandlers, UnregisteredDocumentHandlers) {
        auto doctype = documentHandlers.doctype.map(registerCallback);
        auto comments = documentHandlers.comments.map(registerCallback);
        auto text = documentHandlers.text.map(registerCallback);
        auto end = documentHandlers.end.map(registerCallback);

        // Adding document content handlers cannot fail, so no need for check().
        lol_html_rewriter_builder_add_document_content_handlers(builder,
            doctype == kj::none ? nullptr : &Rewriter::thunk<Doctype>, doctype.orDefault(nullptr),
            comments == kj::none ? nullptr : &Rewriter::thunk<Comment>, comments.orDefault(nullptr),
            text == kj::none ? nullptr : &Rewriter::thunk<Text>, text.orDefault(nullptr),
            end == kj::none ? nullptr : &Rewriter::thunk<DocumentEnd>, end.orDefault(nullptr));
      }
    }
  }

  // `strict` mode will bail out from tokenization process in cases when
  // there is no way to determine correct parsing context. Recommended
  // setting for safety reasons.
  bool isStrict = true;

  // Configure a maximum memory limit that `lol-html` is allowed to use and
  // preallocate some memory for its internal buffer.
  lol_html_memory_settings_t memorySettings = {
    .preallocated_parsing_buffer_size = 1024, .max_allowed_memory_usage = 3 * 1024 * 1024};

  if (FeatureFlags::get(js).getEsiIncludeIsVoidTag()) {
    return LOL_HTML_OWN(rewriter,
        unstable_lol_html_rewriter_build_with_esi_tags(builder, encoding.begin(), encoding.size(),
            memorySettings, &Rewriter::output, &rewriter, isStrict));

  } else {
    return LOL_HTML_OWN(rewriter,
        lol_html_rewriter_build(builder, encoding.begin(), encoding.size(), memorySettings,
            &Rewriter::output, &rewriter, isStrict));
  }
}

Rewriter::Rewriter(jsg::Lock& js,
    kj::ArrayPtr<UnregisteredElementOrDocumentHandlers> unregisteredHandlers,
    kj::ArrayPtr<const char> encoding,
    kj::Own<WritableStreamSink> inner)
    : rewriter(buildRewriter(js, unregisteredHandlers, encoding, *this)),
      externalMemoryAdjustment(js.getExternalMemoryAdjustment()),
      inner(kj::mv(inner)),
      ioContext(IoContext::current()),
      maybeAsyncContext(jsg::AsyncContextFrame::currentRef(js)) {}

namespace {

// The stack size floor enforced by kj. We could go lower,
// but it'd always be increased to this anyway.
const size_t FIBER_STACK_SIZE = 1024 * 64;

const kj::FiberPool& getFiberPool() {
  const static kj::FiberPool FIBER_POOL(FIBER_STACK_SIZE);
  return FIBER_POOL;
}

}  // namespace

kj::Promise<void> Rewriter::write(kj::ArrayPtr<const byte> buffer) {
  KJ_ASSERT(maybeWaitScope == kj::none);
  // Defer fiber creation until the event loop runs. If this promise is dropped synchronously
  // (e.g. by Canceler::cancel() during PumpToReader destruction), no fiber is created, avoiding
  // a KJ assertion failure when destroying an unfired fiber. Once the event loop processes this,
  // the fiber is created and immediately fires (armDepthFirst), so cancellation works normally.
  return kj::evalLater([this, buffer]() {
    return getFiberPool().startFiber([this, buffer](kj::WaitScope& scope) {
      maybeWaitScope = scope;
      if (!isPoisoned()) {
        // Cannot use `check()` because `finishWrite()` implements the error path.
        auto rc = lol_html_rewriter_write(rewriter, buffer.asChars().begin(), buffer.size());
        tryHandleCancellation(rc);
        if (rc == -1) {
          maybePoison(getLastError());
        }
      }
      return finishWrite();
    });
  });
}

kj::Promise<void> Rewriter::write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) {
  KJ_ASSERT(maybeWaitScope == kj::none);
  return kj::evalLater([this, pieces]() {
    return getFiberPool().startFiber([this, pieces](kj::WaitScope& scope) {
      maybeWaitScope = scope;
      if (!isPoisoned()) {
        for (auto bytes: pieces) {
          auto chars = bytes.asChars();
          // Cannot use `check()` because `finishWrite()` implements the error path.
          auto rc = lol_html_rewriter_write(rewriter, chars.begin(), chars.size());
          tryHandleCancellation(rc);
          if (rc == -1) {
            maybePoison(getLastError());
            // A handler threw an exception; stop calling `lol_html_rewriter_write()`.
            break;
          }
        }
      }
      return finishWrite();
    });
  });
}

kj::Promise<void> Rewriter::end() {
  KJ_ASSERT(maybeWaitScope == kj::none);
  return kj::evalLater([this]() {
    return getFiberPool().startFiber([this](kj::WaitScope& scope) {
      maybeWaitScope = scope;
      if (!isPoisoned()) {
        // Cannot use `check()` because `finishWrite()` implements the error path.
        auto rc = lol_html_rewriter_end(rewriter);
        tryHandleCancellation(rc);
        if (rc == -1) {
          maybePoison(getLastError());
        }
      }
      return finishWrite().then([this]() { return inner->end(); });
    });
  });
}

void Rewriter::abort(kj::Exception reason) {
  // End the rewriter and forward the error to the wrapped output stream.
  maybeException = reason.clone();

  inner->abort(kj::mv(reason));
}

kj::Promise<void> Rewriter::finishWrite() {
  maybeWaitScope = kj::none;
  return flushWrite();
}

kj::Promise<void> Rewriter::flushWrite() {
  KJ_ASSERT(!flushing);

  if (!outputBuffer.empty()) {
    KJ_DEFER({
      externalMemoryAdjustment.set(0);
      outputBuffer.clear();
      flushing = false;
    });

    flushing = true;
    co_await inner->write(outputBuffer);
  }

  KJ_IF_SOME(exception, maybeException) {
    inner->abort(exception.clone());
    kj::throwFatalException(exception.clone());
  }
}
template <typename T, typename CType>
lol_html_rewriter_directive_t Rewriter::thunk(CType* content, void* userdata) {
  auto& registration = *reinterpret_cast<RegisteredHandler*>(userdata);
  return registration.rewriter.thunkImpl<T>(content, registration);
}

template <typename T, typename CType>
lol_html_rewriter_directive_t Rewriter::thunkImpl(
    CType* content, RegisteredHandler& registeredHandler) {
  if (isPoisoned()) {
    // Handlers disabled due to exception.
    KJ_LOG(ERROR, "poisoned rewriter should not be able to call handlers");
    return LOL_HTML_STOP;
  }

  try {
    KJ_IF_SOME(exception, kj::runCatchingExceptions([&] {
      // V8 has a thread local pointer that points to where the stack limit is on this thread which
      // is tested for overflows when we enter any JS code. However since we're running in a fiber
      // here, we're in an entirely different stack that V8 doesn't know about, so it gets confused
      // and may think we've overflowed our stack. evalLater will run thunkPromise on the main stack
      // to keep V8 from getting confused.
      auto promise = kj::evalLater([&]() { return thunkPromise<T>(content, registeredHandler); });
      promise.wait(KJ_ASSERT_NONNULL(maybeWaitScope));
      flushWrite().wait(KJ_ASSERT_NONNULL(maybeWaitScope));
    })) {
      // Exception in handler. We need to abort the streaming parser, but can't do so just yet: we
      // need to unwind the stack because we're probably still inside a cool_thing_rewriter_write().
      // We can't unwind with an exception across the Rust/C++ boundary, so instead we'll keep this
      // exception around and disable all later handlers.
      maybePoison(kj::mv(exception));
      return LOL_HTML_STOP;
    }
  } catch (kj::CanceledException) {
    // The fiber is being canceled. Same as runCatchingExceptions, we need to abort the parser,
    // but can't since we're still inside cool_thing_rewriter_write(). This isn't handled by
    // runCatchingExceptions since CanceledException isn't a kj exception, and we wouldn't want
    // runCatchingExceptions to handle it anyway. We set canceled to true and once we leave Rust,
    // we rethrow it to properly cancel the fiber.
    canceled = true;
    return LOL_HTML_STOP;
  }
  return LOL_HTML_CONTINUE;
}

void Rewriter::removeEndTagHandler(RegisteredHandler& handler) {
  auto size = registeredEndTagHandlers.size();
  for (auto counter = size; counter != 0; --counter) {
    auto idx = counter - 1;
    if (registeredEndTagHandlers[idx].get() == &handler) {
      // equivalent of `Vec::swap_remove` in Rust
      if (counter != size) {
        registeredEndTagHandlers[idx] = kj::mv(registeredEndTagHandlers[size - 1]);
      }
      registeredEndTagHandlers.removeLast();
      break;
    }
  }
}

template <typename T, typename CType>
kj::Promise<void> Rewriter::thunkPromise(CType* content, RegisteredHandler& registeredHandler) {
  return ioContext.run(
      [this, content, &registeredHandler](Worker::Lock& lock) -> kj::Promise<void> {
    // We enter the AsyncContextFrame that was current when the Rewriter was created
    // (when transform() was called). If someone wants, instead, to use the context
    // that was current when on(...) is called, the ElementHandler can use AsyncResource
    // (or eventually the standard AsyncContext once that lands).
    jsg::Lock& js = lock;
    jsg::AsyncContextFrame::Scope asyncContextScope(js, maybeAsyncContext);
    auto jsContent = js.alloc<T>(*content, *this);
    auto scope = HTMLRewriter::TokenScope(jsContent);
    auto value = registeredHandler.callback(js, kj::mv(jsContent));

    if constexpr (kj::isSameType<T, EndTag>()) {
      // TODO(someday): We can't unconditionally pop the top of `registeredEndTagHandlers`,
      //   because that depends on https://github.com/cloudflare/lol-html/issues/110
      //   being resolved. For now we let handles to end tag handlers tags live for the duration of
      //   the response transformation, but eagerly release ones that we can.
      //   In particular, note that `thunkPromise` is never called for implied end tags.
      removeEndTagHandler(registeredHandler);
    }

    return value.attach(kj::mv(scope));
  });
}

lol_html_streaming_handler_t Rewriter::registerReplacer(
    jsg::Ref<ReadableStream> content, bool isHtml) {
  auto replacer = kj::heap<RegisteredReplacer>(*this, isHtml, kj::mv(content));
  auto userData = replacer.get();
  registeredReplacers.insert(userData, kj::mv(replacer));

  return {
    .user_data = userData,
    .write_all_callback = Rewriter::replacerThunk,
    .drop_callback = Rewriter::removeRegisteredReplacer,
  };
}

// Adapter that allows pumping a ReadableStream to a pre-established lol_html
// streaming sink, named `sink`. Writes of arbitrary bytes and sizes are allowed,
// but the content must be valid UTF-8 or lol_html will reject it.
class ReplacerStreamSink final: public WritableStreamSink {
 public:
  ReplacerStreamSink(lol_html_streaming_sink_t* sink, bool isHtml): sink(sink), isHtml(isHtml) {}

  kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override KJ_WARN_UNUSED_RESULT {
    auto err = lol_html_streaming_sink_write_utf8_chunk(
        sink, buffer.asChars().begin(), buffer.size(), isHtml);
    if (err != 0) {
      return getLastError();
    }

    return kj::READY_NOW;
  }

  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override {
    for (auto bytes: pieces) {
      auto err = lol_html_streaming_sink_write_utf8_chunk(
          sink, bytes.asChars().begin(), bytes.size(), isHtml);
      if (err != 0) {
        return getLastError();
      }
    }

    return kj::READY_NOW;
  }

  kj::Promise<void> end() override KJ_WARN_UNUSED_RESULT {
    // Nothing specific needs to be done to tell lol_html we're done writing the stream
    return kj::READY_NOW;
  }

  void abort(kj::Exception reason) override {
    // Nothing specific needs to be done. The rewriter will be poisoned in replacerThunkImpl,
    // and lol_html will bubble up the error through to write.
  }

 private:
  lol_html_streaming_sink_t* sink;
  bool isHtml;
};

int Rewriter::replacerThunk(lol_html_streaming_sink_t* sink, void* userData) {
  auto& registration = *reinterpret_cast<RegisteredReplacer*>(userData);
  return registration.rewriter.replacerThunkImpl(sink, registration);
}

int Rewriter::replacerThunkImpl(
    lol_html_streaming_sink_t* sink, RegisteredReplacer& registeredHandler) {
  if (isPoisoned()) {
    // Handlers disabled due to exception.
    KJ_LOG(ERROR, "poisoned rewriter should not be able to call handlers");
    return -1;
  }

  try {
    KJ_IF_SOME(exception, kj::runCatchingExceptions([&] {
      // V8 has a thread local pointer that points to where the stack limit is on this thread which
      // is tested for overflows when we enter any JS code. However since we're running in a fiber
      // here, we're in an entirely different stack that V8 doesn't know about, so it gets confused
      // and may think we've overflowed our stack. evalLater will run thunkPromise on the main stack
      // to keep V8 from getting confused.
      auto promise = kj::evalLater([&]() { return replacerThunkPromise(sink, registeredHandler); });
      promise.wait(KJ_ASSERT_NONNULL(maybeWaitScope));
    })) {
      // Exception in handler. We need to abort the streaming parser, but can't do so just yet: we
      // need to unwind the stack because we're probably still inside a cool_thing_rewriter_write().
      // We can't unwind with an exception across the Rust/C++ boundary, so instead we'll keep this
      // exception around and disable all later handlers
      maybePoison(kj::mv(exception));
      return -1;
    }
  } catch (kj::CanceledException) {
    // The fiber is being canceled. Same as runCatchingExceptions, we need to abort the parser,
    // but can't since we're still inside cool_thing_rewriter_write(). This isn't handled by
    // runCatchingExceptions since CanceledException isn't a kj exception, and we wouldn't want
    // runCatchingExceptions to handle it anyway. We set canceled to true and once we leave Rust,
    // we rethrow it to properly cancel the fiber.
    canceled = true;
    return -1;
  }
  return 0;
}

kj::Promise<void> Rewriter::replacerThunkPromise(
    lol_html_streaming_sink_t* sink, RegisteredReplacer& registration) {
  return ioContext.run([this, sink, &registration](Worker::Lock& lock) -> kj::Promise<void> {
    jsg::AsyncContextFrame::Scope asyncContextScope(lock, maybeAsyncContext);

    auto streamSink = kj::heap<ReplacerStreamSink>(sink, registration.isHtml);
    return ioContext.waitForDeferredProxy(
        registration.stream->pumpTo(lock, kj::mv(streamSink), true));
  });
}

void Rewriter::removeRegisteredReplacer(void* userData) {
  auto& registration = *reinterpret_cast<RegisteredReplacer*>(userData);
  KJ_REQUIRE(registration.rewriter.registeredReplacers.erase(userData),
      "Tried to remove replacer that was not registered");
}

void Rewriter::onEndTag(lol_html_element_t* element, ElementCallbackFunction&& callback) {
  auto registeredHandler = Rewriter::RegisteredHandler{*this, kj::mv(callback)};
  // NOTE: this gets freed in `thunkPromise` above.
  // TODO(someday): this uses more memory than necessary for implied end tags, which lol-html
  // doesn't actually call `thunk` on.  LOL HTML drops the handler after it finishes transforming
  // the current element, but this code will keep it around until the entire HTML document is
  // transformed. It would be nice to free it directly after the handler is used; unfortunately,
  // this isn't trivial to do since we have no idea whether there's an end tag or not. The fix for
  // this probably needs to happen in lol-html; see #110.
  // WARNING: if we ever start reusing the same Rewriter for multiple documents,
  // this will cause a memory leak!
  auto& registeredHandlerPtr = registeredEndTagHandlers.add(kj::heap(kj::mv(registeredHandler)));
  lol_html_element_clear_end_tag_handlers(element);
  check(lol_html_element_add_end_tag_handler(
      element, Rewriter::thunk<EndTag>, registeredHandlerPtr.get()));
}

void Rewriter::output(const char* buffer, size_t size, void* userdata) {
  auto& rewriter = *reinterpret_cast<Rewriter*>(userdata);
  rewriter.outputImpl(kj::asBytes(buffer, size));
}

void Rewriter::outputImpl(kj::ArrayPtr<const byte> buffer) {
  if (isPoisoned()) {
    // Handlers disabled due to exception or running in a destructor.
    return;
  }

  KJ_ASSERT(!flushing);
  externalMemoryAdjustment.adjust(buffer.size());
  outputBuffer.addAll(buffer);
}

// =======================================================================================
// HTMLRewriter::Token::ImplBase<CType>

template <typename CType>
HTMLRewriter::Token::ImplBase<CType>::ImplBase(CType& element, Rewriter& rewriter)
    : element(element),
      rewriter(rewriter) {}
template <typename CType>
HTMLRewriter::Token::ImplBase<CType>::~ImplBase() noexcept(false) {}
template <typename CType>
template <auto Func, auto StreamingFunc>
void HTMLRewriter::Token::ImplBase<CType>::rewriteContentGeneric(
    Content content, jsg::Optional<ContentOptions> options) {
  auto isHtml = options.orDefault({}).html.orDefault(false);

  KJ_SWITCH_ONEOF(content) {
    KJ_CASE_ONEOF(stringContent, kj::String) {
      check(Func(&element, stringContent.cStr(), stringContent.size(), isHtml));
    }

    KJ_CASE_ONEOF(streamContent, jsg::Ref<ReadableStream>) {
      auto handler = rewriter.registerReplacer(kj::mv(streamContent), isHtml);
      check(StreamingFunc(&element, &handler));
    }

    KJ_CASE_ONEOF(responseContent, jsg::Ref<Response>) {
      KJ_IF_SOME(body, responseContent->getBody()) {
        auto handler = rewriter.registerReplacer(kj::mv(body), isHtml);
        check(StreamingFunc(&element, &handler));
      }
      // Otherwise if no body, there is no replacement to make
    }
  }
}
// =======================================================================================
// Element

Element::Element(CType& element, Rewriter& rewriter) {
  impl.emplace(element, rewriter);
}

kj::String Element::getTagName() {
  auto tagName = LolString(lol_html_element_tag_name_get(&checkToken(impl).element));
  return kj::str(tagName.asChars());
}

void Element::setTagName(kj::String name) {
  check(lol_html_element_tag_name_set(&checkToken(impl).element, name.cStr(), name.size()));
}

bool Element::getRemoved() {
  return lol_html_element_is_removed(&checkToken(impl).element);
}

kj::StringPtr Element::getNamespaceURI() {
  // lol-html returns a static C string, no need to handle its lifetime.
  return lol_html_element_namespace_uri_get(&checkToken(impl).element);
}

jsg::Ref<Element::AttributesIterator> Element::getAttributes(jsg::Lock& js) {
  auto& implRef = checkToken(impl);

  auto iter = LOL_HTML_OWN(attributes_iterator, lol_html_attributes_iterator_get(&implRef.element));

  auto jsIter = js.alloc<Element::AttributesIterator>(kj::mv(iter));
  implRef.attributesIterators.add(jsIter.addRef());
  return kj::mv(jsIter);
}

kj::Maybe<kj::String> Element::getAttribute(kj::String name) {
  // NOTE: lol_html_element_get_attribute() returns NULL for both nonexistent attributes and for
  //   errors, so we can't use check() here.
  LolString attr(
      lol_html_element_get_attribute(&checkToken(impl).element, name.cStr(), name.size()));
  // TODO(perf): We could construct a v8::String directly here, saving a copy.
  kj::Maybe kjAttr = attr.asKjString();
  if (kjAttr != kj::none) {
    return kj::mv(kjAttr);
  }

  KJ_IF_SOME(exception, tryGetLastError()) {
    kj::throwFatalException(kj::mv(exception));
  }

  // No error, just doesn't exist.
  return kj::none;
}

bool Element::hasAttribute(kj::String name) {
  return !!check(
      lol_html_element_has_attribute(&checkToken(impl).element, name.cStr(), name.size()));
}

jsg::Ref<Element> Element::setAttribute(kj::String name, kj::String value) {
  auto& implRef = checkToken(impl);
  check(lol_html_element_set_attribute(
      &implRef.element, name.cStr(), name.size(), value.cStr(), value.size()));

  // Mutating attributes may cause lol-html's internal Vec to reallocate, invalidating
  // any live iterators' pointers. We must invalidate all outstanding iterators.
  for (auto& iter: implRef.attributesIterators) {
    iter->invalidate();
  }

  return JSG_THIS;
}

jsg::Ref<Element> Element::removeAttribute(kj::String name) {
  auto& implRef = checkToken(impl);
  check(lol_html_element_remove_attribute(&implRef.element, name.cStr(), name.size()));

  // Removing attributes may shift elements in lol-html's internal Vec (via retain()),
  // invalidating any live iterators' pointers.
  for (auto& iter: implRef.attributesIterators) {
    iter->invalidate();
  }

  return JSG_THIS;
}

namespace {
kj::String unwrapContent(Content content) {
  return kj::mv(JSG_REQUIRE_NONNULL(content.tryGet<kj::String>(), TypeError,
      "Replacing content in HTML comments using a ReadableStream or Response object is not "
      "implemented. You must provide a string."));
}
}  // namespace

jsg::Ref<Element> Element::before(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_element_before, lol_html_element_streaming_before>(
          kj::mv(content), options);
  return JSG_THIS;
}

jsg::Ref<Element> Element::after(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl).rewriteContentGeneric<lol_html_element_after, lol_html_element_streaming_after>(
      kj::mv(content), options);
  return JSG_THIS;
}

jsg::Ref<Element> Element::prepend(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_element_prepend, lol_html_element_streaming_prepend>(
          kj::mv(content), options);
  return JSG_THIS;
}

jsg::Ref<Element> Element::append(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_element_append, lol_html_element_streaming_append>(
          kj::mv(content), options);
  return JSG_THIS;
}

jsg::Ref<Element> Element::replace(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_element_replace, lol_html_element_streaming_replace>(
          kj::mv(content), options);
  return JSG_THIS;
}

jsg::Ref<Element> Element::setInnerContent(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_element_set_inner_content,
          lol_html_element_streaming_set_inner_content>(kj::mv(content), options);
  return JSG_THIS;
}

jsg::Ref<Element> Element::remove() {
  lol_html_element_remove(&checkToken(impl).element);
  return JSG_THIS;
}

jsg::Ref<Element> Element::removeAndKeepContent() {
  lol_html_element_remove_and_keep_content(&checkToken(impl).element);
  return JSG_THIS;
}

void Element::onEndTag(ElementCallbackFunction&& callback) {
  auto& knownImpl = checkToken(impl);
  knownImpl.rewriter.onEndTag(&knownImpl.element, kj::mv(callback));
}

EndTag::EndTag(CType& endTag, Rewriter& rewriter) {
  impl.emplace(endTag, rewriter);
}

void EndTag::htmlContentScopeEnd() {
  impl = kj::none;
}

kj::String EndTag::getName() {
  auto text = LolString(lol_html_end_tag_name_get(&checkToken(impl).element));
  return kj::str(text.asChars());
}

void EndTag::setName(kj::String text) {
  check(lol_html_end_tag_name_set(&checkToken(impl).element, text.cStr(), text.size()));
}

jsg::Ref<EndTag> EndTag::before(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_end_tag_before, lol_html_end_tag_streaming_before>(
          kj::mv(content), kj::mv(options));
  return JSG_THIS;
}

jsg::Ref<EndTag> EndTag::after(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl).rewriteContentGeneric<lol_html_end_tag_after, lol_html_end_tag_streaming_after>(
      kj::mv(content), kj::mv(options));
  return JSG_THIS;
}

jsg::Ref<EndTag> EndTag::remove() {
  lol_html_end_tag_remove(&checkToken(impl).element);
  return JSG_THIS;
}

void Element::htmlContentScopeEnd() {
  impl = kj::none;
}

Element::Impl::~Impl() noexcept(false) {
  for (auto& jsIter: attributesIterators) {
    static_cast<HTMLRewriter::Token&>(*jsIter).htmlContentScopeEnd();
  }
}

// =======================================================================================
// Element::AttributesIterator

Element::AttributesIterator::AttributesIterator(kj::Own<CType> iter): impl(kj::mv(iter)) {}

jsg::Ref<Element::AttributesIterator> Element::AttributesIterator::self() {
  return JSG_THIS;
}

Element::AttributesIterator::Next Element::AttributesIterator::next() {
  // If the element's attributes were modified (via setAttribute/removeAttribute) while this
  // iterator was live, the underlying lol-html iterator holds stale pointers into a potentially
  // reallocated Vec. Continuing to iterate would be a use-after-free.
  JSG_REQUIRE(!mutatedDuringIteration, Error,
      "The attributes of this element have been modified during iteration. "
      "You must create a new iterator after modifying attributes.");

  // NOTE: lol_html_attribute_t doesn't need to be freed.
  auto* attribute = lol_html_attributes_iterator_next(checkToken(impl));
  if (attribute == nullptr) {
    // End of iteration.
    // TODO(someday): Eagerly deallocate. Can't seem to nullify the Own without also nullifying the
    //   enclosing Maybe, however.
    return {true, kj::none};
  }

  auto name = LolString(lol_html_attribute_name_get(attribute));
  auto value = LolString(lol_html_attribute_value_get(attribute));

  return {false, kj::arr(kj::str(name.asChars()), kj::str(value.asChars()))};
}

void Element::AttributesIterator::invalidate() {
  mutatedDuringIteration = true;
  // Also release the underlying lol-html iterator since it's no longer safe to use.
  impl = kj::none;
}

void Element::AttributesIterator::htmlContentScopeEnd() {
  // Clear the mutation flag so that after scope end, the "content token is no longer valid"
  // error (from checkToken) takes precedence over the mutation error.
  mutatedDuringIteration = false;
  impl = kj::none;
}

// =======================================================================================
// Comment

Comment::Comment(CType& comment, Rewriter&): impl(comment) {}

kj::String Comment::getText() {
  auto text = LolString(lol_html_comment_text_get(&checkToken(impl)));
  return kj::str(text.asChars());
}

void Comment::setText(kj::String text) {
  check(lol_html_comment_text_set(&checkToken(impl), text.cStr(), text.size()));
}

bool Comment::getRemoved() {
  // NOTE: No error checking seems required by this function -- it returns a bool directly.
  return lol_html_comment_is_removed(&checkToken(impl));
}

jsg::Ref<Comment> Comment::before(Content content, jsg::Optional<ContentOptions> options) {
  // TODO(someday): If lol-html adds support for streaming replacements for comments, this
  // function will need to be updated.
  auto stringContent = unwrapContent(kj::mv(content));
  check(lol_html_comment_before(&checkToken(impl), stringContent.cStr(), stringContent.size(),
      options.orDefault({}).html.orDefault(false)));

  return JSG_THIS;
}

jsg::Ref<Comment> Comment::after(Content content, jsg::Optional<ContentOptions> options) {
  // TODO(someday): If lol-html adds support for streaming replacements for comments, this
  // function will need to be updated.
  auto stringContent = unwrapContent(kj::mv(content));
  check(lol_html_comment_after(&checkToken(impl), stringContent.cStr(), stringContent.size(),
      options.orDefault({}).html.orDefault(false)));

  return JSG_THIS;
}

jsg::Ref<Comment> Comment::replace(Content content, jsg::Optional<ContentOptions> options) {
  // TODO(someday): If lol-html adds support for streaming replacements for comments, this
  // function will need to be updated.
  auto stringContent = unwrapContent(kj::mv(content));
  check(lol_html_comment_replace(&checkToken(impl), stringContent.cStr(), stringContent.size(),
      options.orDefault({}).html.orDefault(false)));

  return JSG_THIS;
}

jsg::Ref<Comment> Comment::remove() {
  lol_html_comment_remove(&checkToken(impl));

  return JSG_THIS;
}

void Comment::htmlContentScopeEnd() {
  impl = kj::none;
}

// =======================================================================================
// Text

Text::Text(CType& text, Rewriter& rewriter) {
  impl.emplace(text, rewriter);
}

kj::String Text::getText() {
  auto content = lol_html_text_chunk_content_get(&checkToken(impl).element);
  return kj::heapString(content.data, content.len);
}

bool Text::getLastInTextNode() {
  // NOTE: No error checking seems required by this function -- it returns a bool directly.
  return lol_html_text_chunk_is_last_in_text_node(&checkToken(impl).element);
}

bool Text::getRemoved() {
  // NOTE: No error checking seems required by this function -- it returns a bool directly.
  return lol_html_text_chunk_is_removed(&checkToken(impl).element);
}

jsg::Ref<Text> Text::before(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_text_chunk_before, lol_html_text_chunk_streaming_before>(
          kj::mv(content), kj::mv(options));
  return JSG_THIS;
}

jsg::Ref<Text> Text::after(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_text_chunk_after, lol_html_text_chunk_streaming_after>(
          kj::mv(content), kj::mv(options));
  return JSG_THIS;
}

jsg::Ref<Text> Text::replace(Content content, jsg::Optional<ContentOptions> options) {
  checkToken(impl)
      .rewriteContentGeneric<lol_html_text_chunk_replace, lol_html_text_chunk_streaming_replace>(
          kj::mv(content), kj::mv(options));
  return JSG_THIS;
}

jsg::Ref<Text> Text::remove() {
  lol_html_text_chunk_remove(&checkToken(impl).element);

  return JSG_THIS;
}

void Text::htmlContentScopeEnd() {
  impl = kj::none;
}

// =======================================================================================
// Doctype

Doctype::Doctype(CType& doctype, Rewriter&): impl(doctype) {}

kj::Maybe<kj::String> Doctype::getName() {
  LolString name(lol_html_doctype_name_get(&checkToken(impl)));
  return name.asKjString();
}

kj::Maybe<kj::String> Doctype::getPublicId() {
  LolString publicId(lol_html_doctype_public_id_get(&checkToken(impl)));
  return publicId.asKjString();
}

kj::Maybe<kj::String> Doctype::getSystemId() {
  LolString systemId(lol_html_doctype_system_id_get(&checkToken(impl)));
  return systemId.asKjString();
}

void Doctype::htmlContentScopeEnd() {
  impl = kj::none;
}

// =======================================================================================
// DocumentEnd

DocumentEnd::DocumentEnd(CType& documentEnd, Rewriter&): impl(documentEnd) {}

jsg::Ref<DocumentEnd> DocumentEnd::append(Content content, jsg::Optional<ContentOptions> options) {
  // TODO(someday): If lol-html adds support for streaming replacements for the document end,
  // this function will need to be updated.
  auto stringContent = unwrapContent(kj::mv(content));
  check(lol_html_doc_end_append(&checkToken(impl), stringContent.cStr(), stringContent.size(),
      options.orDefault({}).html.orDefault(false)));

  return JSG_THIS;
}

void DocumentEnd::htmlContentScopeEnd() {
  impl = kj::none;
}

// =======================================================================================
// HTMLRewriter

struct HTMLRewriter::Impl {
  // The list of handlers added to this builder.
  kj::Vector<UnregisteredElementOrDocumentHandlers> unregisteredHandlers;
  // TODO(perf): It'd be nice to eagerly register handlers on the native builder object. However,
  //   currently lol-html rewriters are inextricably linked to the builders which created them,
  //   and this has concurrency and reentrancy ramifications: two rewriters built from the same
  //   builder require synchronization to access safely, and their callbacks must not use the
  //   builder which created them, lest the process deadlock.
  //
  //   In the meantime, we keep this list of handlers around and "replay" their registration, in
  //   order, on the builder object that we create inside of .transform().

  JSG_MEMORY_INFO(HTMLRewriter::Impl) {
    for (const auto& handlers: unregisteredHandlers) {
      KJ_SWITCH_ONEOF(handlers) {
        KJ_CASE_ONEOF(h, UnregisteredElementHandlers) {
          tracker.trackField(nullptr, h);
        }
        KJ_CASE_ONEOF(h, UnregisteredDocumentHandlers) {
          tracker.trackField(nullptr, h);
        }
      }
    }
  }
};

HTMLRewriter::HTMLRewriter(): impl(kj::heap<Impl>()) {}
HTMLRewriter::~HTMLRewriter() noexcept(false) {}

void HTMLRewriter::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
  tracker.trackField("impl", impl);
}

jsg::Ref<HTMLRewriter> HTMLRewriter::constructor(jsg::Lock& js) {
  return js.alloc<HTMLRewriter>();
}

jsg::Ref<HTMLRewriter> HTMLRewriter::on(
    kj::String stringSelector, ElementContentHandlers&& handlers) {
  kj::Own<lol_html_Selector> selector =
      LOL_HTML_OWN(selector, lol_html_selector_parse(stringSelector.cStr(), stringSelector.size()));

  impl->unregisteredHandlers.add(UnregisteredElementHandlers{
    kj::mv(selector), kj::mv(handlers.element), kj::mv(handlers.comments), kj::mv(handlers.text)});

  return JSG_THIS;
}

jsg::Ref<HTMLRewriter> HTMLRewriter::onDocument(DocumentContentHandlers&& handlers) {
  impl->unregisteredHandlers.add(UnregisteredDocumentHandlers{kj::mv(handlers.doctype),
    kj::mv(handlers.comments), kj::mv(handlers.text), kj::mv(handlers.end)});

  return JSG_THIS;
}

jsg::Ref<Response> HTMLRewriter::transform(jsg::Lock& js, jsg::Ref<Response> response) {

  JSG_REQUIRE(response->getType() != "error"_kj, TypeError,
      "HTMLRewriter cannot transform an error response");

  auto maybeInput = response->getBody();

  if (maybeInput == kj::none) {
    // That was easy!
    return kj::mv(response);
  }

  auto& ioContext = IoContext::current();

  auto pipe = newIdentityPipe();
  response = Response::constructor(
      js, kj::Maybe(js.alloc<ReadableStream>(ioContext, kj::mv(pipe.in))), kj::mv(response));

  kj::String ownContentType;
  kj::String encoding = kj::str("utf-8");
  KJ_IF_SOME(contentType,
      response->getHeaders(js)->getCommon(js, capnp::CommonHeaderName::CONTENT_TYPE)) {
    // TODO(cleanup): readContentTypeParameter can be replaced with using
    // workerd/util/mimetype.h directly.
    KJ_IF_SOME(charset, readContentTypeParameter(contentType, "charset")) {
      ownContentType = kj::mv(contentType);
      encoding = kj::mv(charset);
    }
  }

  auto rewriter = kj::heap<Rewriter>(js, impl->unregisteredHandlers, encoding, kj::mv(pipe.out));

  // NOTE: Avoid throwing any exceptions after initiating the pump below. This makes
  //   the input response object disturbed (response.bodyUsed === true), which should only happen
  //   after we know that nothing else (like invalid encoding) could cause an exception.

  // Drive and flush the parser asynchronously.
  ioContext.addTask(
      ioContext
          .waitForDeferredProxy(KJ_ASSERT_NONNULL(maybeInput)->pumpTo(js, kj::mv(rewriter), true))
          .catch_([](kj::Exception&& e) {
    // Errors in pumpTo() are already propagated to the destination stream. We don't want to
    // throw them from here since it'll cause an uncaught exception to be reported via taskFailed(),
    // which would poison the IoContext even though the application may have handled the error.
  }));

  // TODO(soon): EW-2025 Make Rewriter a proper wrapper object and put it in hidden property on the
  //   response so the GC can find the handlers which Rewriter co-owns.
  return kj::mv(response);
}

void HTMLRewriter::visitForGc(jsg::GcVisitor& visitor) {
  for (auto& handlers: impl->unregisteredHandlers) {
    KJ_SWITCH_ONEOF(handlers) {
      KJ_CASE_ONEOF(elementHandlers, UnregisteredElementHandlers) {
        visitor.visit(elementHandlers);
      }
      KJ_CASE_ONEOF(documentHandlers, UnregisteredDocumentHandlers) {
        visitor.visit(documentHandlers);
      }
    }
  }
}

}  // namespace workerd::api