1 files changed, 30 insertions, 8 deletions

diff --git a/libs/binder/RpcSession.cpp b/libs/binder/RpcSession.cpp
index 5fe0b0025f..1c376518f0 100644
--- a/libs/binder/RpcSession.cpp
+++ b/libs/binder/RpcSession.cpp
@@ -132,6 +132,7 @@ bool RpcSession::shutdownAndWait(bool wait) {
     if (wait) {
         LOG_ALWAYS_FATAL_IF(mShutdownListener == nullptr, "Shutdown listener not installed");
         mShutdownListener->waitForShutdown(_l);
+
         LOG_ALWAYS_FATAL_IF(!mThreads.empty(), "Shutdown failed");
     }
 
@@ -259,7 +260,7 @@ status_t RpcSession::readId() {
     return OK;
 }
 
-void RpcSession::WaitForShutdownListener::onSessionLockedAllIncomingThreadsEnded(
+void RpcSession::WaitForShutdownListener::onSessionAllIncomingThreadsEnded(
         const sp<RpcSession>& session) {
     (void)session;
     mShutdown = true;
@@ -291,7 +292,13 @@ RpcSession::PreJoinSetupResult RpcSession::preJoinSetup(base::unique_fd fd) {
     // be able to do nested calls (we can't only read from it)
     sp<RpcConnection> connection = assignIncomingConnectionToThisThread(std::move(fd));
 
-    status_t status = mState->readConnectionInit(connection, sp<RpcSession>::fromExisting(this));
+    status_t status;
+
+    if (connection == nullptr) {
+        status = DEAD_OBJECT;
+    } else {
+        status = mState->readConnectionInit(connection, sp<RpcSession>::fromExisting(this));
+    }
 
     return PreJoinSetupResult{
             .connection = std::move(connection),
@@ -358,6 +365,7 @@ void RpcSession::join(sp<RpcSession>&& session, PreJoinSetupResult&& setupResult
     sp<RpcConnection>& connection = setupResult.connection;
 
     if (setupResult.status == OK) {
+        LOG_ALWAYS_FATAL_IF(!connection, "must have connection if setup succeeded");
         JavaThreadAttacher javaThreadAttacher;
         while (true) {
             status_t status = session->state()->getAndExecuteCommand(connection, session,
@@ -373,9 +381,6 @@ void RpcSession::join(sp<RpcSession>&& session, PreJoinSetupResult&& setupResult
               statusToString(setupResult.status).c_str());
     }
 
-    LOG_ALWAYS_FATAL_IF(!session->removeIncomingConnection(connection),
-                        "bad state: connection object guaranteed to be in list");
-
     sp<RpcSession::EventListener> listener;
     {
         std::lock_guard<std::mutex> _l(session->mMutex);
@@ -387,6 +392,12 @@ void RpcSession::join(sp<RpcSession>&& session, PreJoinSetupResult&& setupResult
         listener = session->mEventListener.promote();
     }
 
+    // done after all cleanup, since session shutdown progresses via callbacks here
+    if (connection != nullptr) {
+        LOG_ALWAYS_FATAL_IF(!session->removeIncomingConnection(connection),
+                            "bad state: connection object guaranteed to be in list");
+    }
+
     session = nullptr;
 
     if (listener != nullptr) {
@@ -577,24 +588,35 @@ bool RpcSession::setForServer(const wp<RpcServer>& server, const wp<EventListene
 
 sp<RpcSession::RpcConnection> RpcSession::assignIncomingConnectionToThisThread(unique_fd fd) {
     std::lock_guard<std::mutex> _l(mMutex);
+
+    // Don't accept any more connections, some have shutdown. Usually this
+    // happens when new connections are still being established as part of a
+    // very short-lived session which shuts down after it already started
+    // accepting new connections.
+    if (mIncomingConnections.size() < mMaxIncomingConnections) {
+        return nullptr;
+    }
+
     sp<RpcConnection> session = sp<RpcConnection>::make();
     session->fd = std::move(fd);
     session->exclusiveTid = gettid();
+
     mIncomingConnections.push_back(session);
+    mMaxIncomingConnections = mIncomingConnections.size();
 
     return session;
 }
 
 bool RpcSession::removeIncomingConnection(const sp<RpcConnection>& connection) {
-    std::lock_guard<std::mutex> _l(mMutex);
+    std::unique_lock<std::mutex> _l(mMutex);
     if (auto it = std::find(mIncomingConnections.begin(), mIncomingConnections.end(), connection);
         it != mIncomingConnections.end()) {
         mIncomingConnections.erase(it);
         if (mIncomingConnections.size() == 0) {
             sp<EventListener> listener = mEventListener.promote();
             if (listener) {
-                listener->onSessionLockedAllIncomingThreadsEnded(
-                        sp<RpcSession>::fromExisting(this));
+                _l.unlock();
+                listener->onSessionAllIncomingThreadsEnded(sp<RpcSession>::fromExisting(this));
             }
         }
         return true;