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)