test: m_on_cancel called after request finishes

Test from: bitcoin/bitcoin#34782 (comment)

Author: Sjors

include/mp/proxy.h

@@ -73,6 +73,8 @@ struct ProxyContext
     //! Hook called on the worker thread just before loop->sync() in PassField
     //! for Context arguments. Used by tests to inject precise disconnect timing.
     std::function<void()> testing_hook_before_sync;
+    //! Hook called on the worker thread just before returning results.
+    std::function<void()> testing_hook_after_cleanup;
 
     ProxyContext(Connection* connection);
 };

include/mp/type-context.h

@@ -73,6 +73,9 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
     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;
                 if (server.m_context.testing_hook_before_sync) server.m_context.testing_hook_before_sync();
+                // Save testing_hook_after_cleanup to a local because the
+                // server may be freed in the cleanup sync() below.
+                auto testing_hook_after_cleanup = std::move(server.m_context.testing_hook_after_cleanup);
                 ServerContext server_context{server, call_context, req};
                 {
                     // Before invoking the function, store a reference to the
@@ -191,6 +194,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     }
                     // End of scope: if KJ_DEFER was reached, it runs here
                 }
+                if (testing_hook_after_cleanup) testing_hook_after_cleanup();
                 return call_context;
             };
 

test/mp/test/test.cpp

@@ -362,6 +362,34 @@ KJ_TEST("Calling async IPC method, with server disconnect racing the call")
     disconnect_thread.join();
 }
 
+KJ_TEST("Calling async IPC method, with server disconnect after cleanup")
+{
+    // Regression test for bitcoin/bitcoin#34782 (stack-use-after-return where
+    // the m_on_cancel callback accessed stack-local cancel_mutex and
+    // server_context after the invoke lambda's inner scope exited). The fix
+    // clears m_on_cancel in the cleanup loop->sync() so it is null by the
+    // time the scope exits.
+    //
+    // Use testing_hook_after_cleanup to trigger a server disconnect after the
+    // inner scope exits (cancel_mutex destroyed). Without the fix, the
+    // disconnect fires m_on_cancel which accesses the destroyed mutex.
+    TestSetup setup;
+    ProxyClient<messages::FooInterface>* foo = setup.client.get();
+    foo->initThreadMap();
+    setup.server->m_impl->m_fn = [] {};
+
+    setup.server->m_context.testing_hook_after_cleanup = [&] {
+        setup.server_disconnect();
+    };
+
+    try {
+        foo->callFnAsync();
+        KJ_EXPECT(false);
+    } catch (const std::runtime_error& e) {
+        KJ_EXPECT(std::string_view{e.what()} == "IPC client method call interrupted by disconnect.");
+    }
+}
+
 KJ_TEST("Make simultaneous IPC calls on single remote thread")
 {
     TestSetup setup;