From 6053b8ee282c19c1d5f6f4d9f3a55ed436f12d4a Mon Sep 17 00:00:00 2001 From: Joaquim Silva Date: Mon, 16 Oct 2023 18:28:23 +0100 Subject: [PATCH] Fix jsg::Ref dtor segfault in actor-state. --- src/workerd/api/actor-state.c++ | 8 ++--- src/workerd/io/hibernation-manager.c++ | 41 +++++++++++++++++--------- src/workerd/io/hibernation-manager.h | 15 ++++++++-- src/workerd/io/worker.h | 4 +-- 4 files changed, 45 insertions(+), 23 deletions(-) diff --git a/src/workerd/api/actor-state.c++ b/src/workerd/api/actor-state.c++ index cb8a66ad99d..0c4e0964718 100644 --- a/src/workerd/api/actor-state.c++ +++ b/src/workerd/api/actor-state.c++ @@ -880,7 +880,7 @@ void DurableObjectState::setWebSocketAutoResponse( // If there's no request/response pair, we unset any current set auto response configuration. KJ_IF_SOME(manager, a.getHibernationManager()) { // If there's no hibernation manager created yet, there's nothing to do here. - manager.unsetWebSocketAutoResponse(); + manager.setWebSocketAutoResponse(kj::none, kj::none); } return; } @@ -902,15 +902,15 @@ void DurableObjectState::setWebSocketAutoResponse( // If there's no hibernation manager created yet, we should create one and // set its auto response. } - KJ_REQUIRE_NONNULL(a.getHibernationManager()).setWebSocketAutoResponse(kj::mv(reqResp)); + KJ_REQUIRE_NONNULL(a.getHibernationManager()).setWebSocketAutoResponse( + reqResp->getRequest(), reqResp->getResponse()); } kj::Maybe> DurableObjectState::getWebSocketAutoResponse() { auto& a = KJ_REQUIRE_NONNULL(IoContext::current().getActor()); KJ_IF_SOME(manager, a.getHibernationManager()) { // If there's no hibernation manager created yet, there's nothing to do here. - auto r = manager.getWebSocketAutoResponse(); - return r; + return manager.getWebSocketAutoResponse(); } return kj::none; } diff --git a/src/workerd/io/hibernation-manager.c++ b/src/workerd/io/hibernation-manager.c++ index 86adc681ae8..e53bd0ec117 100644 --- a/src/workerd/io/hibernation-manager.c++ +++ b/src/workerd/io/hibernation-manager.c++ @@ -99,20 +99,26 @@ kj::Vector> HibernationManagerImpl::getWebSockets( } void HibernationManagerImpl::setWebSocketAutoResponse( - jsg::Ref reqResp) { - autoResponsePair = kj::mv(reqResp); -} - -void HibernationManagerImpl::unsetWebSocketAutoResponse() { - autoResponsePair = kj::none; + kj::Maybe request, kj::Maybe response) { + KJ_IF_SOME(req, request) { + // If we have a request, we must also have a response. If response is kj::none, we'll throw. + autoResponsePair->request = kj::str(req); + autoResponsePair->response = kj::str(KJ_REQUIRE_NONNULL(response)); + return; + } + // If we don't have a request, we must unset both request and response. + autoResponsePair->request = kj::none; + autoResponsePair->response = kj::none; } kj::Maybe> HibernationManagerImpl::getWebSocketAutoResponse() { - KJ_IF_SOME(ar, autoResponsePair) { - return ar.addRef(); - } else { - return kj::none; + KJ_IF_SOME(req, autoResponsePair->request) { + // When getting the currently set auto-response pair, if we have a request we must have a response + // set. If not, we'll throw. + return api::WebSocketRequestResponsePair::constructor(kj::str(req), + kj::str(KJ_REQUIRE_NONNULL(autoResponsePair->response))); } + return kj::none; } void HibernationManagerImpl::setTimerChannel(TimerChannel& timerChannel) { @@ -188,10 +194,12 @@ kj::Promise HibernationManagerImpl::readLoop(HibernatableWebSocket& hib) { auto skip = false; - KJ_IF_SOME (reqResp, autoResponsePair) { + // If we have a request != kj::none, we can compare it the received message. This also implies + // that we have a response set in autoResponsePair. + KJ_IF_SOME (req, autoResponsePair->request) { KJ_SWITCH_ONEOF(message) { KJ_CASE_ONEOF(text, kj::String) { - if (text == (reqResp)->getRequest()) { + if (text == req) { // If the received message matches the one set for auto-response, we must // short-circuit readLoop, store the current timestamp and and automatically respond // with the expected response. @@ -207,7 +215,11 @@ kj::Promise HibernationManagerImpl::readLoop(HibernatableWebSocket& hib) { // If the actor is not hibernated/If the WebSocket is active, we need to update // autoResponseTimestamp on the active websocket. apiWs->setAutoResponseStatus(hib.autoResponseTimestamp, kj::READY_NOW); - co_await apiWs->sendAutoResponse(kj::str(reqResp->getResponse().asArray()), ws); + // Since we had a request set, we must have and response that's sent back using the + // same websocket here. The sending of response is managed in web-socket to avoid + // possible racing problems with regular websocket messages. + co_await apiWs->sendAutoResponse( + kj::str(KJ_REQUIRE_NONNULL(autoResponsePair->response).asArray()), ws); } KJ_CASE_ONEOF(package, api::WebSocket::HibernationPackage) { if (!package.closedOutgoingConnection) { @@ -215,7 +227,8 @@ kj::Promise HibernationManagerImpl::readLoop(HibernatableWebSocket& hib) { // If we do that, we have to provide it with the promise to avoid races. This can // happen if we have a websocket hibernating, that unhibernates and sends a // message while ws.send() for auto-response is also sending. - auto p = ws.send(reqResp->getResponse().asArray()).fork(); + auto p = ws.send( + KJ_REQUIRE_NONNULL(autoResponsePair->response).asArray()).fork(); hib.autoResponsePromise = p.addBranch(); co_await p; hib.autoResponsePromise = kj::READY_NOW; diff --git a/src/workerd/io/hibernation-manager.h b/src/workerd/io/hibernation-manager.h index 8cf24c7a68b..b99bc9bfb4c 100644 --- a/src/workerd/io/hibernation-manager.h +++ b/src/workerd/io/hibernation-manager.h @@ -40,8 +40,8 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager { // This converts our activeOrPackage from an api::WebSocket to a HibernationPackage. void hibernateWebSockets(Worker::Lock& lock) override; - void setWebSocketAutoResponse(jsg::Ref reqResp) override; - void unsetWebSocketAutoResponse() override; + void setWebSocketAutoResponse(kj::Maybe request, + kj::Maybe response) override; kj::Maybe> getWebSocketAutoResponse() override; void setTimerChannel(TimerChannel& timerChannel) override; @@ -187,6 +187,15 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager { TagCollection(TagCollection&& other) = default; }; + // This structure will hold the request and corresponding response for hibernatable websockets + // auto-response feature. Although we store 2 kj::Maybe strings, if we don't have a request set + // we can't have a response, and vice versa. + // TODO(cleanup): Remove kj::Maybe from request and response strings. + struct AutoRequestResponsePair { + kj::Maybe request = kj::none; + kj::Maybe response = kj::none; + }; + // A hashmap of tags to HibernatableWebSockets associated with the tag. // We use a kj::List so we can quickly remove websockets that have disconnected. // Also note that we box the keys and values such that in the event of a hashmap resizing we don't @@ -222,7 +231,7 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager { }; DisconnectHandler onDisconnect; kj::TaskSet readLoopTasks; - kj::Maybe> autoResponsePair; + kj::Own autoResponsePair = kj::heap(); kj::Maybe timer; }; }; // namespace workerd diff --git a/src/workerd/io/worker.h b/src/workerd/io/worker.h index a355cab7ba6..ace5c03f230 100644 --- a/src/workerd/io/worker.h +++ b/src/workerd/io/worker.h @@ -697,8 +697,8 @@ class Worker::Actor final: public kj::Refcounted { jsg::Lock& js, kj::Maybe tag) = 0; virtual void hibernateWebSockets(Worker::Lock& lock) = 0; - virtual void setWebSocketAutoResponse(jsg::Ref reqResp) = 0; - virtual void unsetWebSocketAutoResponse() = 0; + virtual void setWebSocketAutoResponse(kj::Maybe request, + kj::Maybe response) = 0; virtual kj::Maybe> getWebSocketAutoResponse() = 0; virtual void setTimerChannel(TimerChannel& timerChannel) = 0; };