Merge pull request #2143 from cloudflare/kenton/fix-jsrpc-context-check

Fix JsRpcPromise::Resolved context checking.

src/workerd/api/tests/js-rpc-test.js

@@ -883,7 +883,8 @@ export let crossContextSharingDoesntWork = {
           "Cannot perform I/O on behalf of a different request. I/O objects (such as streams, " +
           "request/response bodies, and others) created in the context of one request handler " +
           "cannot be accessed from a different request's handler. This is a limitation of " +
-          "Cloudflare Workers which allows us to improve overall performance."
+          "Cloudflare Workers which allows us to improve overall performance. " +
+          "(I/O type: JsRpcPromise)"
     });
 
     // Now let's try accessing a JsRpcProperty, where the property is NOT a direct property of a
@@ -907,7 +908,8 @@ export let crossContextSharingDoesntWork = {
           "Cannot perform I/O on behalf of a different request. I/O objects (such as streams, " +
           "request/response bodies, and others) created in the context of one request handler " +
           "cannot be accessed from a different request's handler. This is a limitation of " +
-          "Cloudflare Workers which allows us to improve overall performance."
+          "Cloudflare Workers which allows us to improve overall performance. " +
+          "(I/O type: JsRpcPromise)"
     });
   },
 }

src/workerd/api/worker-rpc.c++

@@ -339,7 +339,7 @@ void JsRpcPromise::resolve(jsg::Lock& js, jsg::JsValue result) {
   if (state.is<Pending>()) {
     state = Resolved {
       .result = jsg::Value(js.v8Isolate, result),
-      .ioCtx = IoContext::current().getWeakRef()
+      .ctxCheck = IoContext::current().addObject(*this),
     };
   } else {
     // We'd better dispose this.
@@ -370,7 +370,9 @@ rpc::JsRpcTarget::Client JsRpcPromise::getClientForOneCall(
       return pending.pipeline->getCallPipeline();
     }
     KJ_CASE_ONEOF(resolved, Resolved) {
-      IoContext::requireCurrentOrThrowJs(*resolved.ioCtx);
+      // Dereference `ctxCheck` just to verify we're running in the correct context. (If not,
+      // this will throw.)
+      *resolved.ctxCheck;
 
       // A value was already returned, and we closed the original RPC pipeline. But the application
       // kept the promise around and is still trying to pipeline on it. What do we do?

src/workerd/api/worker-rpc.h

@@ -213,8 +213,11 @@ class JsRpcPromise: public JsRpcClientProvider {
   struct Resolved {
     jsg::Value result;
 
-    // We only use this to prohibit use from the wrong context.
-    kj::Own<IoContext::WeakRef> ioCtx;
+    // Dummy IoPtr to self, used only to verify that we're running in the correct context.
+    // (Dereferencing from the wrong context would throw an exception.)
+    // Note: Can't use IoContext::WeakRef here because it's not thread-safe (it's only intended to
+    //   be helf from KJ I/O objects, but this is a JSG object).
+    IoPtr<JsRpcPromise> ctxCheck;
   };
   struct Disposed {};