diff --git a/capnp/helpers/helpers.pxd b/capnp/helpers/helpers.pxd index 577506e..bdca969 100644 --- a/capnp/helpers/helpers.pxd +++ b/capnp/helpers/helpers.pxd @@ -25,6 +25,7 @@ cdef extern from "../helpers/capabilityHelper.h": VoidPromise convert_to_voidpromise(PyPromise&) cdef extern from "../helpers/rpcHelper.h": + Capability.Client restoreHelper(RpcSystem&) Capability.Client restoreHelper(RpcSystem&, MessageBuilder&) Capability.Client restoreHelper(RpcSystem&, MessageReader&) Capability.Client restoreHelper(RpcSystem&, AnyPointer.Reader&) diff --git a/capnp/helpers/rpcHelper.h b/capnp/helpers/rpcHelper.h index b8ff922..7f89011 100644 --- a/capnp/helpers/rpcHelper.h +++ b/capnp/helpers/rpcHelper.h @@ -1,6 +1,7 @@ #pragma once #include "capnp/dynamic.h" +#include #include "capnp/rpc-twoparty.h" #include "Python.h" #include "capabilityHelper.h" @@ -33,29 +34,42 @@ private: PyObject * py_restorer; }; -capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::MessageBuilder & objectId) { capnp::MallocMessageBuilder hostIdMessage(8); +capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::MessageBuilder & objectId) { + capnp::MallocMessageBuilder hostIdMessage(8); auto hostId = hostIdMessage.initRoot(); hostId.setSide(capnp::rpc::twoparty::Side::SERVER); return client.restore(hostId, objectId.getRoot()); } - -capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::MessageReader & objectId) { capnp::MallocMessageBuilder hostIdMessage(8); +capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::MessageReader & objectId) { + capnp::MallocMessageBuilder hostIdMessage(8); auto hostId = hostIdMessage.initRoot(); hostId.setSide(capnp::rpc::twoparty::Side::SERVER); return client.restore(hostId, objectId.getRoot()); } -capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::AnyPointer::Reader & objectId) { capnp::MallocMessageBuilder hostIdMessage(8); + +capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::AnyPointer::Reader & objectId) { + capnp::MallocMessageBuilder hostIdMessage(8); auto hostId = hostIdMessage.initRoot(); hostId.setSide(capnp::rpc::twoparty::Side::SERVER); return client.restore(hostId, objectId); } -capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::AnyPointer::Builder & objectId) { capnp::MallocMessageBuilder hostIdMessage(8); + +capnp::Capability::Client restoreHelper(capnp::RpcSystem& client, capnp::AnyPointer::Builder & objectId) { + capnp::MallocMessageBuilder hostIdMessage(8); auto hostId = hostIdMessage.initRoot(); hostId.setSide(capnp::rpc::twoparty::Side::SERVER); return client.restore(hostId, objectId); } +capnp::Capability::Client restoreHelper(capnp::RpcSystem& client) { + capnp::MallocMessageBuilder message; + capnp::rpc::SturdyRef::Builder ref = message.getRoot(); + auto hostId = ref.getHostId().initAs(); + hostId.setSide(capnp::rpc::twoparty::Side::SERVER); + return client.restore(hostId, ref.getObjectId()); +} + template capnp::RpcSystem makeRpcClientWithRestorer( diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index 185de1a..7ba44cf 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -1856,7 +1856,9 @@ cdef class TwoPartyClient: cdef _DynamicObjectBuilder object_builder cdef _DynamicObjectReader object_reader - if type(objectId) is _DynamicObjectBuilder: + if objectId is None: + return _CapabilityClient()._init(helpers.restoreHelper(deref(self.thisptr)), self) + elif type(objectId) is _DynamicObjectBuilder: object_builder = objectId return _CapabilityClient()._init(helpers.restoreHelper(deref(self.thisptr), deref(object_builder.thisptr)), self) elif type(objectId) is _DynamicObjectReader: