shutdown bugfix: Prevent segfault in server if connection is broken during long function call

include/mp/proxy-types.h

@@ -122,6 +122,7 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                 const auto& params = call_context.getParams();
                 Context::Reader context_arg = Accessor::get(params);
                 ServerContext server_context{server, call_context, req};
+                bool disconnected{false};
                 {
                     // Before invoking the function, store a reference to the
                     // callbackThread provided by the client in the
@@ -152,13 +153,36 @@ auto PassField(Priority<1>, TypeList<>, ServerContext& server_context, const Fn&
                     // inserted, one already existed, meaning this must be a
                     // recursive call (IPC call calling back to the caller which
                     // makes another IPC call), so avoid modifying the map.
-                    auto erase_thread{inserted ? request_thread : request_threads.end()};
-                    KJ_DEFER(if (erase_thread != request_threads.end()) {
+                    const bool erase_thread{inserted};
+                    KJ_DEFER({
                         std::unique_lock<std::mutex> lock(thread_context.waiter->m_mutex);
-                        request_threads.erase(erase_thread);
+                        // Call erase here with a Connection* argument instead
+                        // of an iterator argument, because the `request_thread`
+                        // iterator may be invalid if the connection is closed
+                        // during this function call. More specifically, the
+                        // iterator may be invalid because SetThread adds a
+                        // cleanup callback to the Connection destructor that
+                        // erases the thread from the map, and also because the
+                        // ProxyServer<Thread> destructor calls
+                        // request_threads.clear().
+                        if (erase_thread) {
+                            disconnected = !request_threads.erase(server.m_context.connection);
+                        } else {
+                            disconnected = !request_threads.count(server.m_context.connection);
+                        }
                     });
                     fn.invoke(server_context, args...);
                 }
+                if (disconnected) {
+                    // If disconnected is true, the Connection object was
+                    // destroyed during the method call. Deal with this by
+                    // returning without ever fulfilling the promise, which will
+                    // cause the ProxyServer object to leak. This is not ideal,
+                    // but fixing the leak will require nontrivial code changes
+                    // because there is a lot of code assuming ProxyServer
+                    // objects are destroyed before Connection objects.
+                    return;
+                }
                 KJ_IF_MAYBE(exception, kj::runCatchingExceptions([&]() {
                     server.m_context.connection->m_loop.sync([&] {
                         auto fulfiller_dispose = kj::mv(fulfiller);

src/mp/proxy.cpp

@@ -277,6 +277,11 @@ std::tuple<ConnThread, bool> SetThread(ConnThreads& threads, std::mutex& mutex,
         std::piecewise_construct, std::forward_as_tuple(connection),
         std::forward_as_tuple(make_thread(), connection, /* destroy_connection= */ false)).first;
     thread->second.setCleanup([&threads, &mutex, thread] {
+        // Note: it is safe to use the `thread` iterator in this cleanup
+        // function, because the iterator would only be invalid if the map entry
+        // was removed, and if the map entry is removed the ProxyClient<Thread>
+        // destructor unregisters the cleanup.
+
         // Connection is being destroyed before thread client is, so reset
         // thread client m_cleanup member so thread client destructor does not
         // try unregister this callback after connection is destroyed.