Prevent crash on unclean disconnect if abandoned IPC call returns interface pointer

This is a fix for bitcoin/bitcoin#34250 which reports
that bitcoin node crashes if a rust stratrumv2 mining client calls
BlockTemplate.waitNext() and disconnects without waiting for a response from
the call, and if mempool fees increased so the call returns a non-null
interface BlockTemplate pointer.

The node would crash in this case while trying to call MakeProxyServer on the
returned BlockTemplate pointer, which would fail because MakeProxyServer would
try to use a reference to the Connection object that had been deleted as a
result of the disconnect.

The fix works by:

- Adding a Connection::m_canceler member variable and using it to cancel any
  IPC response promises that are pending when the connection is destroyed.

- Updating ProxyServer<Thread>::post to use promise.attach() as
  described https://capnproto.org/cxxrpc.html#cancellation to detect
  cancellation and set a ServerContext::request_canceled variable.

- Updating ServerCall to check the ServerContext::request_canceled status after
  any C++ server method returns, and throw an exception if it is set.

- Updating type-context.h PassField() function to deal with the exception
  by catching and logging it.

Author: ryanofsky

include/mp/proxy-io.h

@@ -48,6 +48,19 @@ struct ServerInvokeContext : InvokeContext
     ProxyServer& proxy_server;
     CallContext& call_context;
     int req;
+    //! For IPC methods that execute asynchronously, not on the event-loop
+    //! thread: lock preventing the event-loop thread from freeing the params or
+    //! results structs if the request is canceled while the worker thread is
+    //! reading params (`call_context.getParams()`) or writing results
+    //! (`call_context.getResults()`).
+    Lock* cancel_lock{nullptr};
+    //! For IPC methods that execute asynchronously, not on the event-loop
+    //! thread, this is set to true if the IPC call was canceled by the client
+    //! or canceled by a disconnection. If the call runs on the event-loop
+    //! thread, it can't be canceled. This should be accessed with cancel_lock
+    //! held if it is not null, since in the asynchronous case it is accessed
+    //! from multiple threads.
+    bool request_canceled{false};
 
     ServerInvokeContext(ProxyServer& proxy_server, CallContext& call_context, int req)
         : InvokeContext{*proxy_server.m_context.connection}, proxy_server{proxy_server}, call_context{call_context}, req{req}
@@ -421,18 +434,21 @@ class Connection
     template <typename F>
     void onDisconnect(F&& f)
     {
-        // Add disconnect handler to local TaskSet to ensure it is cancelled and
+        // Add disconnect handler to local TaskSet to ensure it is canceled and
         // will never run after connection object is destroyed. But when disconnect
         // handler fires, do not call the function f right away, instead add it
         // to the EventLoop TaskSet to avoid "Promise callback destroyed itself"
-        // error in cases where f deletes this Connection object.
+        // error in the typical case where f deletes this Connection object.
         m_on_disconnect.add(m_network.onDisconnect().then(
             [f = std::forward<F>(f), this]() mutable { m_loop->m_task_set->add(kj::evalLater(kj::mv(f))); }));
     }
 
     EventLoopRef m_loop;
     kj::Own<kj::AsyncIoStream> m_stream;
     LoggingErrorHandler m_error_handler{*m_loop};
+    //! TaskSet used to cancel the m_network.onDisconnect() handler for remote
+    //! disconnections, if the connection is closed locally first by deleting
+    //! this Connection object.
     kj::TaskSet m_on_disconnect{m_error_handler};
     ::capnp::TwoPartyVatNetwork m_network;
     std::optional<::capnp::RpcSystem<::capnp::rpc::twoparty::VatId>> m_rpc_system;
@@ -445,6 +461,11 @@ class Connection
     //! ThreadMap.makeThread) used to service requests to clients.
     ::capnp::CapabilityServerSet<Thread> m_threads;
 
+    //! Canceler for canceling promises that we want to discard when the
+    //! connection is destroyed. This is used to interrupt method calls that are
+    //! still executing at time of disconnection.
+    kj::Canceler m_canceler;
+
     //! Cleanup functions to run if connection is broken unexpectedly.  List
     //! will be empty if all ProxyClient are destroyed cleanly before the
     //! connection is destroyed.
@@ -696,9 +717,17 @@ template<typename T, typename Fn>
 kj::Promise<T> ProxyServer<Thread>::post(Fn&& fn)
 {
     auto ready = kj::newPromiseAndFulfiller<void>(); // Signaled when waiter is ready to post again.
-    auto ret = m_thread_ready.then([this, fn = std::forward<Fn>(fn), ready_fulfiller = kj::mv(ready.fulfiller)]() mutable {
+    auto cancel_monitor_ptr = kj::heap<CancelMonitor>();
+    CancelMonitor& cancel_monitor = *cancel_monitor_ptr;
+    // Keep a reference to the ProxyServer<Thread> instance by assigning it to
+    // the self variable. ProxyServer instances are reference-counted and if the
+    // client drops its reference, this variable keeps the instance alive until
+    // the thread finishes executing. The self variable needs to be destroyed on
+    // the event loop thread so it is freed in a sync() call below.
+    auto self = thisCap();
+    auto ret = m_thread_ready.then([this, self = std::move(self), fn = std::forward<Fn>(fn), ready_fulfiller = kj::mv(ready.fulfiller), cancel_monitor_ptr = kj::mv(cancel_monitor_ptr)]() mutable {
         auto result = kj::newPromiseAndFulfiller<T>(); // Signaled when fn() is called, with its return value.
-        bool posted = m_thread_context.waiter->post([this, fn = std::forward<Fn>(fn), ready_fulfiller = kj::mv(ready_fulfiller), result_fulfiller = kj::mv(result.fulfiller)]() mutable {
+        bool posted = m_thread_context.waiter->post([this, self = std::move(self), fn = std::forward<Fn>(fn), ready_fulfiller = kj::mv(ready_fulfiller), result_fulfiller = kj::mv(result.fulfiller), cancel_monitor_ptr = kj::mv(cancel_monitor_ptr)]() mutable {
             // Fulfill ready.promise now, as soon as the Waiter starts executing
             // this lambda, so the next ProxyServer<Thread>::post() call can
             // immediately call waiter->post(). It is important to do this
@@ -712,8 +741,15 @@ kj::Promise<T> ProxyServer<Thread>::post(Fn&& fn)
                 ready_fulfiller = nullptr;
             });
             std::optional<T> result_value;
-            kj::Maybe<kj::Exception> exception{kj::runCatchingExceptions([&]{ result_value.emplace(fn()); })};
-            m_loop->sync([&result_value, &exception, result_fulfiller = kj::mv(result_fulfiller)]() mutable {
+            kj::Maybe<kj::Exception> exception{kj::runCatchingExceptions([&]{ result_value.emplace(fn(*cancel_monitor_ptr)); })};
+            m_loop->sync([this, &result_value, &exception, self = kj::mv(self), result_fulfiller = kj::mv(result_fulfiller), cancel_monitor_ptr = kj::mv(cancel_monitor_ptr)]() mutable {
+                // Destroy CancelMonitor here before fulfilling or rejecting the
+                // promise so it doesn't get triggered when the promise is
+                // destroyed.
+                cancel_monitor_ptr = nullptr;
+                // Send results to the fulfiller. Technically it would be ok to
+                // skip this if promise was canceled, but it's simpler to just
+                // do it unconditionally.
                 KJ_IF_MAYBE(e, exception) {
                     assert(!result_value);
                     result_fulfiller->reject(kj::mv(*e));
@@ -723,6 +759,11 @@ kj::Promise<T> ProxyServer<Thread>::post(Fn&& fn)
                     result_value.reset();
                 }
                 result_fulfiller = nullptr;
+                // Use evalLater to destroy the ProxyServer<Thread> self
+                // reference, if it is the last reference, because the
+                // ProxyServer<Thread> destructor needs to join the thread,
+                // which can't happen until this sync() block has exited.
+                m_loop->m_task_set->add(kj::evalLater([self = kj::mv(self)] {}));
             });
         });
         // Assert that calling Waiter::post did not fail. It could only return
@@ -731,7 +772,7 @@ kj::Promise<T> ProxyServer<Thread>::post(Fn&& fn)
         // signaled, so this should never happen.
         assert(posted);
         return kj::mv(result.promise);
-    });
+    }).attach(kj::heap<CancelProbe>(cancel_monitor));
     m_thread_ready = kj::mv(ready.promise);
     return ret;
 }

include/mp/proxy-types.h

@@ -445,9 +445,36 @@ struct ServerCall
     template <typename ServerContext, typename... Args>
     decltype(auto) invoke(ServerContext& server_context, TypeList<>, Args&&... args) const
     {
-        return ProxyServerMethodTraits<typename decltype(server_context.call_context.getParams())::Reads>::invoke(
-            server_context,
-            std::forward<Args>(args)...);
+        // If cancel_lock is set, release it while executing the method, and
+        // reacquire it afterwards. The lock is needed to prevent params and
+        // response structs from being deleted by the event loop thread if the
+        // request is canceled, so it is only needed before and after method
+        // execution. It is important to release the lock during execution
+        // because the method can take arbitrarily long to return and the event
+        // loop will need the lock itself in on_cancel if the call is canceled.
+        if (server_context.cancel_lock) server_context.cancel_lock->m_lock.unlock();
+        return TryFinally(
+            [&]() -> decltype(auto) {
+                return ProxyServerMethodTraits<
+                    typename decltype(server_context.call_context.getParams())::Reads
+                >::invoke(server_context, std::forward<Args>(args)...);
+            },
+            [&] {
+                if (server_context.cancel_lock) server_context.cancel_lock->m_lock.lock();
+                // If the IPC request was canceled, throw InterruptException
+                // because there is no point continuing and trying to fill the
+                // call_context.getResults() struct. It's also important to stop
+                // executing because the connection may have been destroyed as
+                // described in https://github.com/bitcoin/bitcoin/issues/34250
+                // and there could be invalid references to the destroyed
+                // Connection object if this continued.
+                // If the IPC method itself threw an exception, the
+                // InterruptException thrown below will take precedence over it.
+                // Since the call has been canceled that exception can't be
+                // returned to the caller, so it needs to be discarded like
+                // other result values.
+                if (server_context.request_canceled) throw InterruptException{"canceled"};
+            });
     }
 };
 

include/mp/proxy.h

@@ -153,6 +153,7 @@ struct ProxyServerBase : public virtual Interface_::Server
     ProxyServerBase(std::shared_ptr<Impl> impl, Connection& connection);
     virtual ~ProxyServerBase();
     void invokeDestroy();
+    using Interface_::Server::thisCap;
 
     /**
      * Implementation pointer that may or may not be owned and deleted when this

include/mp/type-context.h

@@ -8,6 +8,8 @@
 #include <mp/proxy-io.h>
 #include <mp/util.h>
 
+#include <kj/string.h>
+
 namespace mp {
 template <typename Output>
 void CustomBuildField(TypeList<>,
@@ -63,7 +65,14 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
     Context::Reader context_arg = Accessor::get(params);
     auto& server = server_context.proxy_server;
     int req = server_context.req;
-    auto invoke = [call_context = kj::mv(server_context.call_context), &server, req, fn, args...]() mutable {
+    // Keep a reference to the ProxyServer instance by assigning it to the self
+    // variable. ProxyServer instances are reference-counted and if the client
+    // drops its reference and the IPC call is canceled, this variable keeps the
+    // instance alive until the method finishes executing. The self variable
+    // needs to be destroyed on the event loop thread so it is freed in a sync()
+    // call below.
+    auto self = server.thisCap();
+    auto invoke = [self = kj::mv(self), call_context = kj::mv(server_context.call_context), &server, req, fn, args...](CancelMonitor& cancel_monitor) mutable {
                 MP_LOG(*server.m_context.loop, Log::Debug) << "IPC server executing request #" << req;
                 const auto& params = call_context.getParams();
                 Context::Reader context_arg = Accessor::get(params);
@@ -89,8 +98,35 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     auto& thread_context = g_thread_context;
                     auto& request_threads = thread_context.request_threads;
                     ConnThread request_thread;
-                    bool inserted;
+                    bool inserted{false};
+                    Mutex cancel_mutex;
+                    Lock cancel_lock{cancel_mutex};
+                    server_context.cancel_lock = &cancel_lock;
                     server.m_context.loop->sync([&] {
+                        // Detect request being canceled before it executes.
+                        if (cancel_monitor.m_canceled) {
+                            server_context.request_canceled = true;
+                            return;
+                        }
+                        // Detect request being canceled while it executes.
+                        assert(!cancel_monitor.m_on_cancel);
+                        cancel_monitor.m_on_cancel = [&server, &server_context, &cancel_mutex, req]() {
+                            MP_LOG(*server.m_context.loop, Log::Info) << "IPC server request #" << req << " canceled while executing.";
+                            // Lock cancel_mutex here to block the event loop
+                            // thread and prevent it from deleting the request's
+                            // params and response structs while the execution
+                            // thread is accessing them. Because this lock is
+                            // released before the event loop thread does delete
+                            // the structs, the mutex does not provide any
+                            // protection from the event loop deleting the
+                            // structs _before_ the execution thread acquires
+                            // it. So in addition to locking the mutex, the
+                            // execution thread always checks request_canceled
+                            // as well before accessing the structs.
+                            Lock{cancel_mutex};
+                            server_context.request_canceled = true;
+                        };
+                        // Update requests_threads map if not canceled.
                         std::tie(request_thread, inserted) = SetThread(
                             GuardedRef{thread_context.waiter->m_mutex, request_threads}, server.m_context.connection,
                             [&] { return context_arg.getCallbackThread(); });
@@ -102,13 +138,23 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     // recursive call (IPC call calling back to the caller which
                     // makes another IPC call), so avoid modifying the map.
                     const bool erase_thread{inserted};
-                    KJ_DEFER(if (erase_thread) {
+                    KJ_DEFER(
+                        // Release the cancel lock before calling loop->sync and
+                        // waiting for the event loop thread, because if a
+                        // cancellation happened, it needs to run the on_cancel
+                        // callback above. It's safe to release cancel_lock at
+                        // this point because the fn.invoke() call below will be
+                        // finished and no longer accessing the params or
+                        // results structs.
+                        cancel_lock.m_lock.unlock();
                         // Erase the request_threads entry on the event loop
                         // thread with loop->sync(), so if the connection is
                         // broken there is not a race between this thread and
                         // the disconnect handler trying to destroy the thread
                         // client object.
                         server.m_context.loop->sync([&] {
+                            auto self_dispose{kj::mv(self)};
+                            if (erase_thread) {
                             // Look up the thread again without using existing
                             // iterator since entry may no longer be there after
                             // a disconnect. Destroy node after releasing
@@ -120,9 +166,22 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                                 Lock lock(thread_context.waiter->m_mutex);
                                 removed = request_threads.extract(server.m_context.connection);
                             }
+                            }
                         });
-                    });
-                    fn.invoke(server_context, args...);
+                    );
+                    if (server_context.request_canceled) {
+                        MP_LOG(*server.m_context.loop, Log::Info) << "IPC server request #" << req << " canceled before it could be executed";
+                    } else KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]{
+                        try {
+                            fn.invoke(server_context, args...);
+                        } catch (const InterruptException& e) {
+                            MP_LOG(*server.m_context.loop, Log::Info) << "IPC server request #" << req << " interrupted (" << e.what() << ")";
+                        }
+                    })) {
+                        MP_LOG(*server.m_context.loop, Log::Error) << "IPC server request #" << req << " uncaught exception (" << kj::str(*exception).cStr() << ")";
+                        throw kj::mv(*exception);
+                    }
+                    // End of scope: if KJ_DEFER was reached, it runs here
                 }
                 return call_context;
             };
@@ -131,7 +190,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
     // be a local Thread::Server object, but it needs to be looked up
     // asynchronously with getLocalServer().
     auto thread_client = context_arg.getThread();
-    return server.m_context.connection->m_threads.getLocalServer(thread_client)
+    auto result = server.m_context.connection->m_threads.getLocalServer(thread_client)
         .then([&server, invoke = kj::mv(invoke), req](const kj::Maybe<Thread::Server&>& perhaps) mutable {
             // Assuming the thread object is found, pass it a pointer to the
             // `invoke` lambda above which will invoke the function on that
@@ -147,6 +206,12 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                 throw std::runtime_error("invalid thread handle");
             }
         });
+    // Use connection m_canceler object to cancel the result promise if the
+    // connection is destroyed. (By default Cap'n Proto does not cancel requests
+    // on disconnect, since it's possible clients might want to make requests
+    // and immediately disconnect without waiting for results, but not want the
+    // the requests to be canceled.)
+    return server.m_context.connection->m_canceler.wrap(kj::mv(result));
 }
 } // namespace mp