SocketClient: add optional reference counting

Needed to fix a race in netd.

Bug: 3438459
Change-Id: Icd7f5f035510235f733a25c0621479d3e644b152

3 files changed, 34 insertions, 2 deletions

diff --git a/include/sysutils/SocketClient.h b/include/sysutils/SocketClient.h
index 2fcc331..ef6dd9b 100644
--- a/include/sysutils/SocketClient.h
+++ b/include/sysutils/SocketClient.h
@@ -19,6 +19,10 @@ class SocketClient {
     /* Peer group ID */
     gid_t mGid;
 
+    /* Reference count (starts at 1) */
+    pthread_mutex_t mRefCountMutex;
+    int mRefCount;
+
 public:
     SocketClient(int sock);
     virtual ~SocketClient() {}
@@ -34,6 +38,13 @@ public:
 
     // Sending binary data:
     int sendData(const void *data, int len);
+
+    // Optional reference counting.  Reference count starts at 1.  If
+    // it's decremented to 0, it deletes itself.
+    // SocketListener creates a SocketClient (at refcount 1) and calls
+    // decRef() when it's done with the client.
+    void incRef();
+    void decRef();
 };
 
 typedef android::List<SocketClient *> SocketClientCollection;
diff --git a/libsysutils/src/SocketClient.cpp b/libsysutils/src/SocketClient.cpp
index a6aed26..6d4dff4 100644
--- a/libsysutils/src/SocketClient.cpp
+++ b/libsysutils/src/SocketClient.cpp
@@ -15,8 +15,10 @@ SocketClient::SocketClient(int socket)
         , mPid(-1)
         , mUid(-1)
         , mGid(-1)
+        , mRefCount(1)
 {
     pthread_mutex_init(&mWriteMutex, NULL);
+    pthread_mutex_init(&mRefCountMutex, NULL);
 
     struct ucred creds;
     socklen_t szCreds = sizeof(creds);
@@ -100,3 +102,22 @@ int SocketClient::sendData(const void* data, int len) {
     pthread_mutex_unlock(&mWriteMutex);
     return 0;
 }
+
+void SocketClient::incRef() {
+  pthread_mutex_lock(&mRefCountMutex);
+  mRefCount++;
+  pthread_mutex_unlock(&mRefCountMutex);
+}
+
+void SocketClient::decRef() {
+  bool deleteSelf = false;
+  pthread_mutex_lock(&mRefCountMutex);
+  mRefCount--;
+  if (mRefCount == 0) {
+      deleteSelf = true;
+  }
+  pthread_mutex_unlock(&mRefCountMutex);
+  if (deleteSelf) {
+      delete this;
+  }
+}
diff --git a/libsysutils/src/SocketListener.cpp b/libsysutils/src/SocketListener.cpp
index 611d5fe..dde9b55 100644
--- a/libsysutils/src/SocketListener.cpp
+++ b/libsysutils/src/SocketListener.cpp
@@ -55,7 +55,7 @@ SocketListener::~SocketListener() {
     }
     SocketClientCollection::iterator it;
     for (it = mClients->begin(); it != mClients->end();) {
-        delete (*it);
+        (*it)->decRef();
         it = mClients->erase(it);
     }
     delete mClients;
@@ -226,7 +226,7 @@ void SocketListener::runListener() {
                 pthread_mutex_unlock(&mClientsLock);
                 /* Destroy the client */
                 close(c->getSocket());
-                delete c;
+                c->decRef();
             }
         }
     }