Fix problems with TwoPartyServer

Fixes #61

It turns out I messed up the Server initialization code for the case
where a string is passed in as the address. The tests only cover the
cases where a raw socket is passed in. This will be rectified in a
following commit.

capnp/helpers/helpers.pxd

@@ -32,7 +32,8 @@ cdef extern from "capnp/helpers/rpcHelper.h":
     Capability.Client restoreHelper(RpcSystem&, AnyPointer.Builder&)
     Capability.Client bootstrapHelper(RpcSystem&)
     RpcSystem makeRpcClientWithRestorer(TwoPartyVatNetwork&, PyRestorer&)
-    PyPromise connectServer(TaskSet &, PyRestorer &, AsyncIoContext *, StringPtr)
+    PyPromise connectServerRestorer(TaskSet &, PyRestorer &, AsyncIoContext *, StringPtr)
+    PyPromise connectServer(TaskSet &, Capability.Client, AsyncIoContext *, StringPtr)
 
 cdef extern from "capnp/helpers/serialize.h":
     ByteArray messageToPackedBytes(MessageBuilder &, size_t wordCount)

capnp/helpers/rpcHelper.h

@@ -89,12 +89,12 @@ capnp::RpcSystem<SturdyRefHostId> makeRpcClientWithRestorer(
   return RpcSystem<SturdyRefHostId>(network, restorer);
 }
 
-struct ServerContext {
+struct ServerContextRestorer {
   kj::Own<kj::AsyncIoStream> stream;
   capnp::TwoPartyVatNetwork network;
   capnp::RpcSystem<capnp::rpc::twoparty::SturdyRefHostId> rpcSystem;
 
-  ServerContext(kj::Own<kj::AsyncIoStream>&& stream, capnp::SturdyRefRestorer<capnp::AnyPointer>& restorer)
+  ServerContextRestorer(kj::Own<kj::AsyncIoStream>&& stream, capnp::SturdyRefRestorer<capnp::AnyPointer>& restorer)
       : stream(kj::mv(stream)),
         network(*this->stream, capnp::rpc::twoparty::Side::SERVER),
         rpcSystem(makeRpcServer(network, restorer)) {}
@@ -106,14 +106,14 @@ class ErrorHandler : public kj::TaskSet::ErrorHandler {
   }
 };
 
-void acceptLoop(kj::TaskSet & tasks, PyRestorer & restorer, kj::Own<kj::ConnectionReceiver>&& listener) {
+void acceptLoopRestorer(kj::TaskSet & tasks, PyRestorer & restorer, kj::Own<kj::ConnectionReceiver>&& listener) {
   auto ptr = listener.get();
   tasks.add(ptr->accept().then(kj::mvCapture(kj::mv(listener),
       [&](kj::Own<kj::ConnectionReceiver>&& listener,
              kj::Own<kj::AsyncIoStream>&& connection) {
-    acceptLoop(tasks, restorer, kj::mv(listener));
+    acceptLoopRestorer(tasks, restorer, kj::mv(listener));
 
-    auto server = kj::heap<ServerContext>(kj::mv(connection), restorer);
+    auto server = kj::heap<ServerContextRestorer>(kj::mv(connection), restorer);
 
     // Arrange to destroy the server context when all references are gone, or when the
     // EzRpcServer is destroyed (which will destroy the TaskSet).
@@ -121,7 +121,7 @@ void acceptLoop(kj::TaskSet & tasks, PyRestorer & restorer, kj::Own<kj::Connecti
   })));
 }
 
-kj::Promise<PyObject *> connectServer(kj::TaskSet & tasks, PyRestorer & restorer, kj::AsyncIoContext * context, kj::StringPtr bindAddress) {
+kj::Promise<PyObject *> connectServerRestorer(kj::TaskSet & tasks, PyRestorer & restorer, kj::AsyncIoContext * context, kj::StringPtr bindAddress) {
     auto paf = kj::newPromiseAndFulfiller<unsigned int>();
     auto portPromise = paf.promise.fork();
 
@@ -131,7 +131,50 @@ kj::Promise<PyObject *> connectServer(kj::TaskSet & tasks, PyRestorer & restorer
                  kj::Own<kj::NetworkAddress>&& addr) {
       auto listener = addr->listen();
       portFulfiller->fulfill(listener->getPort());
-      acceptLoop(tasks, restorer, kj::mv(listener));
+      acceptLoopRestorer(tasks, restorer, kj::mv(listener));
+    })));
+
+    return portPromise.addBranch().then([&](unsigned int port) { return PyLong_FromUnsignedLong(port); });
+}
+
+
+struct ServerContext {
+  kj::Own<kj::AsyncIoStream> stream;
+  capnp::TwoPartyVatNetwork network;
+  capnp::RpcSystem<capnp::rpc::twoparty::SturdyRefHostId> rpcSystem;
+
+  ServerContext(kj::Own<kj::AsyncIoStream>&& stream, capnp::Capability::Client client)
+      : stream(kj::mv(stream)),
+        network(*this->stream, capnp::rpc::twoparty::Side::SERVER),
+        rpcSystem(makeRpcServer(network, client)) {}
+};
+
+void acceptLoop(kj::TaskSet & tasks, capnp::Capability::Client client, kj::Own<kj::ConnectionReceiver>&& listener) {
+  auto ptr = listener.get();
+  tasks.add(ptr->accept().then(kj::mvCapture(kj::mv(listener),
+      [&, client](kj::Own<kj::ConnectionReceiver>&& listener,
+             kj::Own<kj::AsyncIoStream>&& connection) mutable {
+    acceptLoop(tasks, client, kj::mv(listener));
+
+    auto server = kj::heap<ServerContext>(kj::mv(connection), client);
+
+    // Arrange to destroy the server context when all references are gone, or when the
+    // EzRpcServer is destroyed (which will destroy the TaskSet).
+    tasks.add(server->network.onDisconnect().attach(kj::mv(server)));
+  })));
+}
+
+kj::Promise<PyObject *> connectServer(kj::TaskSet & tasks, capnp::Capability::Client client, kj::AsyncIoContext * context, kj::StringPtr bindAddress) {
+    auto paf = kj::newPromiseAndFulfiller<unsigned int>();
+    auto portPromise = paf.promise.fork();
+
+    tasks.add(context->provider->getNetwork().parseAddress(bindAddress)
+        .then(kj::mvCapture(paf.fulfiller,
+          [&, client](kj::Own<kj::PromiseFulfiller<unsigned int>>&& portFulfiller,
+                 kj::Own<kj::NetworkAddress>&& addr) mutable {
+      auto listener = addr->listen();
+      portFulfiller->fulfill(listener->getPort());
+      acceptLoop(tasks, client, kj::mv(listener));
     })));
 
     return portPromise.addBranch().then([&](unsigned int port) { return PyLong_FromUnsignedLong(port); });

capnp/lib/capnp.pyx

@@ -2279,7 +2279,7 @@ cdef class TwoPartyServer:
         self._bootstrap = None
 
         if isinstance(socket, basestring):
-            self._connect(socket)
+            self._connect(socket, restorer, bootstrap)
         else:
             self._orig_stream = socket
             self._stream = _FdAsyncIoStream(socket.fileno())
@@ -2302,11 +2302,19 @@ cdef class TwoPartyServer:
             Py_INCREF(self._network)
             self._disconnect_promise = self.on_disconnect().then(self._decref)
 
-    cpdef _connect(self, host_string):
+    cpdef _connect(self, host_string, restorer, bootstrap):
+        cdef _InterfaceSchema schema
         cdef _EventLoop loop = C_DEFAULT_EVENT_LOOP_GETTER()
         cdef capnp.StringPtr temp_string = capnp.StringPtr(<char*>host_string, len(host_string))
         self._task_set = new capnp.TaskSet(self._error_handler)
-        self.port_promise = Promise()._init(helpers.connectServer(deref(self._task_set), deref(self._restorer.thisptr), loop.thisptr, temp_string))
+        if restorer:
+            self._restorer = _convert_restorer(restorer)
+            self.port_promise = Promise()._init(helpers.connectServerRestorer(deref(self._task_set), deref(self._restorer.thisptr), loop.thisptr, temp_string))
+        else:
+            self._bootstrap = bootstrap
+            Py_INCREF(self._bootstrap)
+            schema = bootstrap.schema
+            self.port_promise = Promise()._init(helpers.connectServer(deref(self._task_set), helpers.server_to_client(schema.thisptr, <PyObject *>bootstrap), loop.thisptr, temp_string))
 
     def _decref(self):
         Py_DECREF(self._bootstrap)

examples/calculator_client.py

@@ -39,7 +39,7 @@ def main(host):
     # takes a struct or AnyPointer as an argument), and then cast the returned
     # capability to it's proper type. This casting is due to capabilities not
     # having a reference to their schema
-    calculator = client.ez_restore('calculator').cast_as(calculator_capnp.Calculator)
+    calculator = client.bootstrap().cast_as(calculator_capnp.Calculator)
 
     '''Make a request that just evaluates the literal value 123.
 

examples/calculator_server.py

@@ -130,15 +130,10 @@ given address/port ADDRESS may be '*' to bind to all local addresses.\
     return parser.parse_args()
 
 
-def restore(ref):
-    assert ref.as_text() == 'calculator'
-    return CalculatorImpl()
-
-
 def main():
     address = parse_args().address
 
-    server = capnp.TwoPartyServer(address, restore)
+    server = capnp.TwoPartyServer(address, bootstrap=CalculatorImpl())
     server.run_forever()
 
 if __name__ == '__main__':

test/test_rpc_calculator.py

@@ -12,7 +12,7 @@ import calculator_server
 def test_calculator():
     read, write = socket.socketpair(socket.AF_UNIX)
 
-    server = capnp.TwoPartyServer(write, calculator_server.restore)
+    server = capnp.TwoPartyServer(write, bootstrap=calculator_server.CalculatorImpl())
     calculator_client.main(read)
 
 
@@ -29,7 +29,7 @@ def test_calculator_gc():
     evaluate_impl_orig = calculator_server.evaluate_impl
     calculator_server.evaluate_impl = new_evaluate_impl(evaluate_impl_orig)
 
-    server = capnp.TwoPartyServer(write, calculator_server.restore)
+    server = capnp.TwoPartyServer(write, bootstrap=calculator_server.CalculatorImpl())
     calculator_client.main(read)
 
     calculator_server.evaluate_impl = evaluate_impl_orig