From 998d373de43c726ea37db00ef3cef676ff9a0da3 Mon Sep 17 00:00:00 2001 From: Rahul Sethi <5822355+RamIdeas@users.noreply.github.com> Date: Fri, 2 Aug 2024 13:11:04 +0100 Subject: [PATCH] move maxMessageSize and reading feature flag into startReadLoop --- src/workerd/api/web-socket.c++ | 16 +++++++++------- src/workerd/api/web-socket.h | 3 --- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/workerd/api/web-socket.c++ b/src/workerd/api/web-socket.c++ index 92e337a0d240..4d3e7c9d7f12 100644 --- a/src/workerd/api/web-socket.c++ +++ b/src/workerd/api/web-socket.c++ @@ -259,10 +259,6 @@ jsg::Ref WebSocket::constructor( headers.unset(kj::HttpHeaderId::SEC_WEBSOCKET_EXTENSIONS); } - if (FeatureFlags::get(js).getIncreaseWebsocketMessageSize()) { - ws->maxMessageSize = 128u << 20; - } - auto prom = ([](auto& context, auto connUrl, auto headers, auto client) -> kj::Promise { auto response = co_await client->openWebSocket(connUrl, headers); @@ -454,6 +450,11 @@ WebSocket::Accepted::~Accepted() noexcept(false) { } void WebSocket::startReadLoop(jsg::Lock& js, kj::Maybe> cs) { + size_t maxMessageSize = 1u << 20; + 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. @@ -461,7 +462,7 @@ void WebSocket::startReadLoop(jsg::Lock& js, kj::Maybe 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> promise = readLoop(kj::mv(cs)); + kj::Promise> promise = readLoop(kj::mv(cs), maxMessageSize); auto& context = IoContext::current(); @@ -908,14 +909,15 @@ kj::Array WebSocket::getHibernatableTags() { } kj::Promise> WebSocket::readLoop( - kj::Maybe> cs) { + kj::Maybe> 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()).ws.getIfNotHibernatable()); auto& context = IoContext::current(); while (true) { - auto message = co_await ws.receive(ws.maxMessageSize); + auto message = co_await ws.receive(maxMessageSize); context.getLimitEnforcer().topUpActor(); KJ_IF_SOME(a, context.getActor()) { diff --git a/src/workerd/api/web-socket.h b/src/workerd/api/web-socket.h index 70d956289195..1fc0ff4a00c5 100644 --- a/src/workerd/api/web-socket.h +++ b/src/workerd/api/web-socket.h @@ -425,9 +425,6 @@ 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 - size_t maxMessageSize = 1u << 20; - // Maximum size of a WebSocket attachment. inline static const size_t MAX_ATTACHMENT_SIZE = 1024 * 2;