From 8a5c0059acb82d1400048c175e0b8294c0ad89e2 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++ | 7 ++++--- src/workerd/io/hibernation-manager.c++ | 23 ++++++++++++++--------- src/workerd/io/hibernation-manager.h | 11 ++++++++--- src/workerd/io/worker.h | 4 ++-- 4 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/workerd/api/actor-state.c++ b/src/workerd/api/actor-state.c++ index cb8a66ad99d..0d9454bf283 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,7 +902,8 @@ 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() { @@ -910,7 +911,7 @@ kj::Maybe> DurableObjectState::getWe 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 kj::mv(r); } return kj::none; } diff --git a/src/workerd/io/hibernation-manager.c++ b/src/workerd/io/hibernation-manager.c++ index 86adc681ae8..c0220863a52 100644 --- a/src/workerd/io/hibernation-manager.c++ +++ b/src/workerd/io/hibernation-manager.c++ @@ -99,17 +99,22 @@ kj::Vector> HibernationManagerImpl::getWebSockets( } void HibernationManagerImpl::setWebSocketAutoResponse( - jsg::Ref reqResp) { - autoResponsePair = kj::mv(reqResp); -} - -void HibernationManagerImpl::unsetWebSocketAutoResponse() { + kj::Maybe request, kj::Maybe response) { + KJ_IF_SOME(req, request) { + KJ_IF_SOME(resp, response){ + auto autoRR = kj::heap(); + autoRR->response = kj::mv(req); + autoRR->request = kj::mv(resp); + autoResponsePair = kj::mv(autoRR); + return; + } + } autoResponsePair = kj::none; } kj::Maybe> HibernationManagerImpl::getWebSocketAutoResponse() { KJ_IF_SOME(ar, autoResponsePair) { - return ar.addRef(); + return api::WebSocketRequestResponsePair::constructor(kj::str(ar.request), kj::str(ar.response)); } else { return kj::none; } @@ -191,7 +196,7 @@ kj::Promise HibernationManagerImpl::readLoop(HibernatableWebSocket& hib) { KJ_IF_SOME (reqResp, autoResponsePair) { KJ_SWITCH_ONEOF(message) { KJ_CASE_ONEOF(text, kj::String) { - if (text == (reqResp)->getRequest()) { + if (text == reqResp.request) { // 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 +212,7 @@ 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); + co_await apiWs->sendAutoResponse(kj::str(reqResp.response.asArray()), ws); } KJ_CASE_ONEOF(package, api::WebSocket::HibernationPackage) { if (!package.closedOutgoingConnection) { @@ -215,7 +220,7 @@ 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(reqResp.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..2964bcf5fed 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,11 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager { TagCollection(TagCollection&& other) = default; }; + struct AutoRequestResponsePair { + kj::StringPtr request; + kj::StringPtr response; + }; + // 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 +227,7 @@ class HibernationManagerImpl final : public Worker::Actor::HibernationManager { }; DisconnectHandler onDisconnect; kj::TaskSet readLoopTasks; - kj::Maybe> autoResponsePair; + kj::Maybe autoResponsePair; 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; };