Fix TextEncoderStream surrogate pair handling across chunks (#5680)

Co-authored-by: Erik Corry <[email protected]>

Author: anonrig

src/workerd/api/BUILD.bazel

@@ -436,14 +436,14 @@ wd_cc_library(
     implementation_deps = [
         "//src/workerd/io:features",
         "//src/workerd/util:strings",
-        "@simdutf",
     ],
     visibility = ["//visibility:public"],
     deps = [
         ":util",
         "//src/workerd/io:compatibility-date_capnp",
         "//src/workerd/jsg",
         "@capnp-cpp//src/kj",
+        "@simdutf",
     ],
 )
 

src/workerd/api/streams/encoding.c++

@@ -4,37 +4,105 @@
 
 #include "encoding.h"
 
+#include "simdutf.h"
+
 #include <workerd/api/encoding.h>
 #include <workerd/api/streams/standard.h>
 #include <workerd/io/features.h>
 #include <workerd/jsg/jsg.h>
 
+#include <v8.h>
+
+#include <kj/common.h>
+#include <kj/refcount.h>
+
 namespace workerd::api {
 
+namespace {
+constexpr kj::byte REPLACEMENT_UTF8[] = {0xEF, 0xBF, 0xBD};
+
+struct Holder: public kj::Refcounted {
+  kj::Maybe<char16_t> pending = kj::none;
+};
+}  // namespace
+
+// TextEncoderStream encodes a stream of JavaScript strings into UTF-8 bytes.
+//
+// WHATWG Encoding spec requirement (https://encoding.spec.whatwg.org/#interface-textencoderstream):
+// The encoder must encode unpaired UTF-16 surrogates as replacement characters.
+//
+// simdutf handles this for us, but we have to be careful of surrogate pairs
+//   (high surrogate, followed by low surrogate) split across chunk boundaries.
+//
+// We do this with the pending field:
+//   holder->pending = kj::none    -> No pending high surrogate from previous chunk
+//   holder->pending = char16_t    -> High surrogate waiting for a matching low surrogate
+//
+// Ref: https://github.com/web-platform-tests/wpt/blob/master/encoding/streams/encode-utf8.any.js
 jsg::Ref<TextEncoderStream> TextEncoderStream::constructor(jsg::Lock& js) {
-  auto transformer = TransformStream::constructor(js,
-      Transformer{.transform = jsg::Function<Transformer::TransformAlgorithm>(
-                      [](jsg::Lock& js, auto chunk, auto controller) {
-    auto str = jsg::check(chunk->ToString(js.v8Context()));
-    auto utf8Length = str->Utf8LengthV2(js.v8Isolate);
+  auto state = kj::rc<Holder>();
 
-    // Don't emit empty chunks
-    if (utf8Length == 0) {
-      return js.resolvedPromise();
+  auto transform = [holder = state.addRef()](jsg::Lock& js, v8::Local<v8::Value> chunk,
+                       jsg::Ref<TransformStreamDefaultController> controller) mutable {
+    auto str = jsg::check(chunk->ToString(js.v8Context()));
+    size_t length = str->Length();
+    if (length == 0) return js.resolvedPromise();
+
+    // Allocate buffer: reserve slot 0 for pending surrogate if we have one
+    size_t prefix = (holder->pending == kj::none) ? 0 : 1;
+    size_t end = prefix + length;
+    auto buf = kj::heapArray<char16_t>(end);
+    str->WriteV2(js.v8Isolate, 0, length, reinterpret_cast<uint16_t*>(buf.begin() + prefix));
+
+    KJ_IF_SOME(lead, holder->pending) {
+      buf.begin()[0] = lead;
+      holder->pending = kj::none;
     }
 
-    v8::Local<v8::ArrayBuffer> buffer;
-    JSG_REQUIRE(v8::ArrayBuffer::MaybeNew(js.v8Isolate, utf8Length).ToLocal(&buffer), RangeError,
-        "Cannot allocate space for TextEncoder.encode");
-
-    auto bytes = jsg::asBytes(buffer).releaseAsChars();
-    [[maybe_unused]] auto written = str->WriteUtf8V2(
-        js.v8Isolate, bytes.begin(), bytes.size(), v8::String::WriteFlags::kReplaceInvalidUtf8);
-
-    KJ_DASSERT(written == buffer->ByteLength());
-    controller->enqueue(js, v8::Uint8Array::New(buffer, 0, buffer->ByteLength()));
+    // If chunk ends with high surrogate, save it for next chunk
+    if (end > 0 && U_IS_LEAD(buf[end - 1])) {
+      holder->pending = buf[--end];
+    }
+    if (end == 0) return js.resolvedPromise();
+
+    auto slice = buf.first(end);
+    auto result = simdutf::utf8_length_from_utf16_with_replacement(slice.begin(), slice.size());
+    // Only sanitize if there are surrogates in the buffer - UTF-16 without
+    // surrogates is always well-formed.
+    if (result.error == simdutf::error_code::SURROGATE) {
+      simdutf::to_well_formed_utf16(slice.begin(), slice.size(), slice.begin());
+    }
+    auto utf8Length = result.count;
+    KJ_DASSERT(utf8Length > 0 && utf8Length >= end);
+
+    auto backingStore = js.allocBackingStore(utf8Length, jsg::Lock::AllocOption::UNINITIALIZED);
+    auto dest = kj::ArrayPtr<char>(static_cast<char*>(backingStore->Data()), utf8Length);
+    [[maybe_unused]] auto written =
+        simdutf::convert_utf16_to_utf8(slice.begin(), slice.size(), dest.begin());
+    KJ_DASSERT(written == utf8Length, "simdutf should write exactly utf8Length bytes");
+
+    auto array = v8::Uint8Array::New(
+        v8::ArrayBuffer::New(js.v8Isolate, kj::mv(backingStore)), 0, utf8Length);
+    controller->enqueue(js, jsg::JsUint8Array(array));
     return js.resolvedPromise();
-  })},
+  };
+
+  auto flush = [holder = state.addRef()](
+                   jsg::Lock& js, jsg::Ref<TransformStreamDefaultController> controller) mutable {
+    // If stream ends with orphaned high surrogate, emit replacement character
+    if (holder->pending != kj::none) {
+      auto backingStore = js.allocBackingStore(3, jsg::Lock::AllocOption::UNINITIALIZED);
+      memcpy(backingStore->Data(), REPLACEMENT_UTF8, 3);
+      auto array =
+          v8::Uint8Array::New(v8::ArrayBuffer::New(js.v8Isolate, kj::mv(backingStore)), 0, 3);
+      controller->enqueue(js, jsg::JsUint8Array(array));
+    }
+    return js.resolvedPromise();
+  };
+
+  auto transformer = TransformStream::constructor(js,
+      Transformer{.transform = jsg::Function<Transformer::TransformAlgorithm>(kj::mv(transform)),
+        .flush = jsg::Function<Transformer::FlushAlgorithm>(kj::mv(flush))},
       StreamQueuingStrategy{}, StreamQueuingStrategy{});
 
   return js.alloc<TextEncoderStream>(transformer->getReadable(), transformer->getWritable());

src/wpt/encoding-test.ts

@@ -110,20 +110,7 @@ export default {
     ],
   },
   'streams/encode-bad-chunks.any.js': {},
-  'streams/encode-utf8.any.js': {
-    comment: 'Surrogate pair handling across chunks not yet implemented',
-    expectedFailures: [
-      'a character split between chunks should be correctly encoded',
-      'a character following one split between chunks should be correctly encoded',
-      'an unmatched surrogate at the end of a chunk followed by an astral character in the next chunk should be replaced with the replacement character at the start of the next output chunk',
-      'an unmatched surrogate at the end of a chunk followed by an ascii character in the next chunk should be replaced with the replacement character at the start of the next output chunk',
-      'a non-terminal unpaired leading surrogate should immediately be replaced',
-      'two consecutive astral characters each split down the middle should be correctly reassembled',
-      'two consecutive astral characters each split down the middle with an invalid surrogate in the middle should be correctly encoded',
-      'an unmatched surrogate at the end of a chunk followed by a plane 1 character split into two chunks should result in the encoded plane 1 character appearing in the last output chunk',
-      'a leading surrogate chunk should be carried past empty chunks',
-    ],
-  },
+  'streams/encode-utf8.any.js': {},
   'streams/invalid-realm.window.js': {
     comment: 'Enable when ShadowRealm is supported',
     expectedFailures: [