Add a compat flag for increasing the WS message size limit (#2164)

* Add a compat flag for increasing the WS message size limit

* Move compat flag to end of list

* Reverse negation

* Update messaging

* move maxMessageSize and reading feature flag into startReadLoop

* use constant kj::WebSocket::SUGGESTED_MAX_MESSAGE_SIZE

---------

Co-authored-by: Rahul Sethi <[email protected]>

Author: penalosa

src/workerd/api/web-socket.c++

@@ -476,14 +476,19 @@ WebSocket::Accepted::~Accepted() noexcept(false) {
 }
 
 void WebSocket::startReadLoop(jsg::Lock& js, kj::Maybe<kj::Own<InputGate::CriticalSection>> cs) {
+  size_t maxMessageSize = kj::WebSocket::SUGGESTED_MAX_MESSAGE_SIZE;
+  if (FeatureFlags::get(js).getIncreaseWebsocketMessageSize()) {
+    maxMessageSize = 128u << 20;
+  }
+
   // If the kj::WebSocket happens to be an AbortableWebSocket (see util/abortable.h), then
   // calling readLoop here could throw synchronously if the canceler has already been tripped.
   // Using kj::evalNow() here let's us capture that and handle correctly.
   //
   // We catch exceptions and return Maybe<Exception> instead since we want to handle the exceptions
   // in awaitIo() below, but we don't want the KJ exception converted to JavaScript before we can
   // examine it.
-  kj::Promise<kj::Maybe<kj::Exception>> promise = readLoop(kj::mv(cs));
+  kj::Promise<kj::Maybe<kj::Exception>> promise = readLoop(kj::mv(cs), maxMessageSize);
 
   auto& context = IoContext::current();
 
@@ -942,14 +947,15 @@ kj::Array<kj::StringPtr> WebSocket::getHibernatableTags() {
 }
 
 kj::Promise<kj::Maybe<kj::Exception>> WebSocket::readLoop(
-    kj::Maybe<kj::Own<InputGate::CriticalSection>> cs) {
+    kj::Maybe<kj::Own<InputGate::CriticalSection>> cs,
+    size_t maxMessageSize) {
   try {
     // Note that we'll throw if the websocket has enabled hibernation.
     auto& ws = *KJ_REQUIRE_NONNULL(
         KJ_ASSERT_NONNULL(farNative->state.tryGet<Accepted>()).ws.getIfNotHibernatable());
     auto& context = IoContext::current();
     while (true) {
-      auto message = co_await ws.receive();
+      auto message = co_await ws.receive(maxMessageSize);
 
       auto size = countBytesFromMessage(message);
       KJ_IF_SOME(o, observer) {

src/workerd/api/web-socket.h

@@ -445,6 +445,9 @@ class WebSocket: public EventTarget {
   // `close()`, thereby preventing calls to `send()` even after we wake from hibernation.
   bool closedOutgoingForHib = false;
 
+  // Maximum allowed size for WebSocket messages
+  inline static const size_t SUGGESTED_MAX_MESSAGE_SIZE = 1u << 20;
+
   // Maximum size of a WebSocket attachment.
   inline static const size_t MAX_ATTACHMENT_SIZE = 1024 * 2;
 
@@ -648,7 +651,7 @@ class WebSocket: public EventTarget {
       IoContext& context, OutgoingMessagesMap& outgoingMessages, kj::WebSocket& ws, Native& native,
       AutoResponse& autoResponse, kj::Maybe<kj::Own<WebSocketObserver>>& observer);
 
-  kj::Promise<kj::Maybe<kj::Exception>> readLoop(kj::Maybe<kj::Own<InputGate::CriticalSection>> cs);
+  kj::Promise<kj::Maybe<kj::Exception>> readLoop(kj::Maybe<kj::Own<InputGate::CriticalSection>> cs, size_t maxMessageSize);
 
   void reportError(jsg::Lock& js, kj::Exception&& e);
   void reportError(jsg::Lock& js, jsg::JsRef<jsg::JsValue> err);

src/workerd/io/compatibility-date.capnp

@@ -523,4 +523,11 @@ struct CompatibilityFlags @0x8f8c1b68151b6cef {
   # Enables fetching hosts with a custom port from workers.
   # For orange clouded sites only standard ports are allowed (https://developers.cloudflare.com/fundamentals/reference/network-ports/#network-ports-compatible-with-cloudflares-proxy).
   # For grey clouded sites all ports are allowed.
+
+  increaseWebsocketMessageSize @56 :Bool
+      $compatEnableFlag("increase_websocket_message_size")
+      $experimental;
+  # For local development purposes only, increase the message size limit to 128MB.
+  # This is not expected ever to be made available in production, as large messages are inefficient.
+
 }