Fix error handling when creating clients (mp::ConnectStream)

example/example.cpp

@@ -10,7 +10,6 @@
 #include <fstream>
 #include <future>
 #include <iostream>
-#include <kj/async.h>
 #include <kj/common.h>
 #include <memory>
 #include <mp/proxy.h>

include/mp/proxy-io.h

@@ -592,7 +592,28 @@ ProxyClientBase<Interface, Impl>::ProxyClientBase(typename Interface::Client cli
         });
     }
     });
-    Sub::construct(*this);
+    // If construct() fails, run the cleanup functions before rethrowing,
+    // because ~ProxyClientBase will not run for an object whose constructor
+    // threw, and the connection would otherwise be leaked.
+    try {
+        Sub::construct(*this);
+    } catch (...) {
+        MP_LOG(*m_context.loop, Log::Debug) << "Cleaning up " << CxxTypeName(*this) << " " << this << " after construct() failure";
+        CleanupRun(m_context.cleanup_fns);
+        throw;
+    }
+
+    // If this client owns the connection, delete the connection on disconnect.
+    if (destroy_connection) {
+        m_context.loop->sync([&] {
+            EventLoop& loop = *m_context.loop;
+            Connection* connection = m_context.connection;
+            connection->onDisconnect([&loop, connection] {
+                MP_LOG(loop, Log::Warning) << "IPC client: unexpected network disconnect.";
+                delete connection;
+            });
+        });
+    }
 }
 
 template <typename Interface, typename Impl>
@@ -832,6 +853,10 @@ kj::Promise<T> ProxyServer<Thread>::post(Fn&& fn)
 //! Given a stream, make a new ProxyClient object to send requests over it.
 //! Also create a new Connection object embedded in the client that is freed
 //! when the client is closed.
+//!
+//! If the init interface declares a construct() method, creating the client
+//! calls it, so this function may block making an IPC call and may throw if
+//! the call fails.
 template <typename InitInterface>
 std::unique_ptr<ProxyClient<InitInterface>> ConnectStream(EventLoop& loop, Stream stream)
 {
@@ -840,11 +865,6 @@ std::unique_ptr<ProxyClient<InitInterface>> ConnectStream(EventLoop& loop, Strea
     loop.sync([&] {
         connection = std::make_unique<Connection>(loop, kj::mv(stream));
         init_client = connection->m_rpc_system->bootstrap(ServerVatId().vat_id).castAs<InitInterface>();
-        Connection* connection_ptr = connection.get();
-        connection->onDisconnect([&loop, connection_ptr] {
-            MP_LOG(loop, Log::Warning) << "IPC client: unexpected network disconnect.";
-            delete connection_ptr;
-        });
     });
     return std::make_unique<ProxyClient<InitInterface>>(
         kj::mv(init_client), connection.release(), /* destroy_connection= */ true);

test/CMakeLists.txt

@@ -28,6 +28,7 @@ if(BUILD_TESTING AND TARGET CapnProto::kj-test)
     mp/test/foo.h
     mp/test/listen_tests.cpp
     mp/test/spawn_tests.cpp
+    mp/test/connect_tests.cpp
     mp/test/test.cpp
   )
   include(${PROJECT_SOURCE_DIR}/cmake/TargetCapnpSources.cmake)

test/mp/test/connect_tests.cpp

@@ -0,0 +1,239 @@
+// Copyright (c) The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+#include "unixlistener.h"
+#include <kj/async.h>
+#include <kj/common.h>
+#include <kj/debug.h>
+#include <kj/memory.h>
+#include <kj/test.h>
+#include <mp/proxy.h>
+#include <mp/proxy-io.h>
+#include <mp/test/foo.capnp.h>
+#include <mp/test/foo.capnp.proxy.h>
+#include <sys/socket.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <chrono>
+#include <condition_variable>
+#include <cstring> // IWYU pragma: keep
+#include <functional>
+#include <future>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <string_view>
+#include <thread>
+
+namespace mp {
+namespace test {
+namespace {
+
+constexpr auto FAILURE_TIMEOUT = std::chrono::seconds{30};
+
+//! Default event loop log handler used by tests, throws so the calling code
+//! can assert on errors.
+void DefaultLogHandler(mp::LogMessage log)
+{
+    if (log.level == mp::Log::Raise)
+        throw std::runtime_error(log.message);
+}
+
+class TestSetup
+{
+public:
+    int client_fd;
+    int server_fd;
+
+    mp::EventLoop* loop;
+    std::optional<mp::EventLoopRef> loop_ref;
+    //! Thread variable should be after other struct members so the thread does
+    //! not start until the other members are initialized.
+    std::thread loop_thread;
+
+    TestSetup(mp::LogFn log_handler = DefaultLogHandler)
+        : TestSetup(
+              [](int fds[2]) {
+                  KJ_REQUIRE(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) != -1);
+              },
+              log_handler) {}
+
+    TestSetup(const std::function<void(int[2])>& init_sockets,
+              mp::LogFn log_handler = DefaultLogHandler)
+    {
+        std::promise<mp::EventLoop*> loop_promise;
+        loop_thread = std::thread([&, log_handler] {
+            mp::EventLoop loop("mptest-connect", log_handler);
+            loop_promise.set_value(&loop);
+            loop.loop();
+        });
+        loop = loop_promise.get_future().get();
+        loop_ref.emplace(*loop);
+
+        // Initialize and store sockets
+        int fds[2] = {-1, -1};
+        init_sockets(fds);
+
+        client_fd = fds[0];
+        server_fd = fds[1];
+    }
+
+    ~TestSetup()
+    {
+        loop_ref.reset();
+        loop_thread.join();
+    }
+};
+
+KJ_TEST("ConnectStream connects to a socket serving a valid init interface")
+{
+    TestSetup setup;
+
+    std::thread server_thread([&setup]() {
+        mp::EventLoop server_loop("mptest-valid-server", DefaultLogHandler);
+        std::unique_ptr<FooInit> init = std::make_unique<FooInit>();
+        ServeStream<messages::FooInit>(server_loop, MakeStream(server_loop, setup.server_fd), *init);
+        server_loop.loop();
+    });
+
+    auto init = ConnectStream<messages::FooInit>(*setup.loop, MakeStream(*setup.loop, setup.client_fd));
+
+    init.reset();
+    server_thread.join();
+    KJ_EXPECT(true);
+}
+
+KJ_TEST("ConnectStream throws when the socket is already disconnected")
+{
+    TestSetup setup;
+
+    close(setup.server_fd);
+
+    try {
+        auto init = ConnectStream<messages::FooInit>(*setup.loop, MakeStream(*setup.loop, setup.client_fd));
+
+        KJ_EXPECT(false);
+    } catch (const std::runtime_error& e) {
+        std::string_view reason = e.what();
+
+        KJ_EXPECT(reason == "IPC client method call interrupted by disconnect.");
+    }
+}
+
+KJ_TEST("ConnectStream defers disconnect failure to the first IPC request for interfaces without construct()")
+{
+    TestSetup setup;
+
+    close(setup.server_fd);
+
+    // Without a construct() method no IPC call is made during client
+    // creation, so ConnectStream succeeds even though the peer is gone.
+    auto foo = ConnectStream<messages::FooInterface>(*setup.loop, MakeStream(*setup.loop, setup.client_fd));
+
+    try {
+        foo->add(1, 2);
+        KJ_EXPECT(false);
+    } catch (const std::runtime_error& e) {
+        std::string_view reason = e.what();
+
+        KJ_EXPECT(reason == "IPC client method called after disconnect.");
+    }
+}
+
+KJ_TEST("ConnectStream handles a disconnect when no client calls are made")
+{
+    std::mutex mutex;
+    std::condition_variable cv;
+    bool warned = false;
+
+    TestSetup setup([&](mp::LogMessage log) {
+        if (log.level == mp::Log::Warning && log.message.find("unexpected network disconnect") != std::string::npos) {
+            const std::lock_guard<std::mutex> lock(mutex);
+            warned = true;
+            cv.notify_all();
+        }
+        DefaultLogHandler(log);
+    });
+
+    close(setup.server_fd);
+
+    auto foo = ConnectStream<messages::FooInterface>(*setup.loop, MakeStream(*setup.loop, setup.client_fd));
+
+    // The disconnect handler registered by ProxyClientBase should run and
+    // delete the connection even when no calls are ever made.
+    std::unique_lock<std::mutex> lock(mutex);
+    KJ_EXPECT(cv.wait_for(lock, FAILURE_TIMEOUT, [&] { return warned; }));
+}
+
+KJ_TEST("ConnectStream throws when the socket disconnects after receiving data")
+{
+    TestSetup setup;
+
+    std::thread server_thread([&setup]() {
+        char buf[128];
+
+        ssize_t bytes_received =
+            recv(setup.server_fd, buf, sizeof(buf), 0);
+
+        if (bytes_received > 0) {
+            close(setup.server_fd);
+        }
+    });
+
+    try {
+        auto init = ConnectStream<messages::FooInit>(*setup.loop, MakeStream(*setup.loop, setup.client_fd));
+
+        if (server_thread.joinable()) server_thread.join();
+        KJ_EXPECT(false);
+    } catch (const std::runtime_error& e) {
+        if (server_thread.joinable()) server_thread.join();
+
+        std::string_view reason = e.what();
+        KJ_EXPECT(reason == "IPC client method call interrupted by disconnect.");
+    }
+}
+
+KJ_TEST("ConnectStream throws when a connection accepted from a listener disconnects after receiving data")
+{
+    UnixListener listener;
+
+    TestSetup setup([&listener](int fds[2]) {
+        fds[0] = listener.MakeConnectedSocket(); // client_fd
+        fds[1] = listener.release();             // server_fd
+    });
+
+    std::thread server_thread([&setup]() {
+        char buf[128];
+
+        int connection_fd = accept(setup.server_fd, nullptr, nullptr);
+
+        if (connection_fd >= 0) {
+            ssize_t bytes_received =
+                recv(connection_fd, buf, sizeof(buf), 0);
+
+            if (bytes_received > 0) {
+                close(connection_fd);
+            }
+        }
+        close(setup.server_fd);
+    });
+
+    try {
+        auto init = ConnectStream<messages::FooInit>(*setup.loop, MakeStream(*setup.loop, setup.client_fd));
+
+        if (server_thread.joinable()) server_thread.join();
+        KJ_EXPECT(false);
+    } catch (const std::runtime_error& e) {
+        if (server_thread.joinable()) server_thread.join();
+
+        std::string_view reason = e.what();
+        KJ_EXPECT(reason == "IPC client method call interrupted by disconnect.");
+    }
+}
+
+} // namespace
+} // namespace test
+} // namespace mp

test/mp/test/foo.capnp

@@ -39,6 +39,10 @@ interface FooInterface $Proxy.wrap("mp::test::FooImplementation") {
     passDataPointers @22 (arg :List(Data)) -> (result :List(Data));
 }
 
+interface FooInit $Proxy.wrap("mp::test::FooInit") {
+    construct @0 () -> ();
+}
+
 interface FooCallback $Proxy.wrap("mp::test::FooCallback") {
     destroy @0 (context :Proxy.Context) -> ();
     call @1 (context :Proxy.Context, arg :Int32) -> (result :Int32);

test/mp/test/foo.h

@@ -66,6 +66,10 @@ class ExtendedCallback : public FooCallback
     virtual int callExtended(int arg) = 0;
 };
 
+class FooInit
+{
+};
+
 class FooImplementation
 {
 public: