2 files changed, 33 insertions, 5 deletions
diff --git a/libsysutils/src/SocketClient.cpp b/libsysutils/src/SocketClient.cpp
index a6aed26..90ca52e 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,25 @@ 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);
+}
+
+bool SocketClient::decRef() {
+    bool deleteSelf = false;
+    pthread_mutex_lock(&mRefCountMutex);
+    mRefCount--;
+    if (mRefCount == 0) {
+        deleteSelf = true;
+    } else if (mRefCount < 0) {
+        SLOGE("SocketClient refcount went negative!");
+    }
+    pthread_mutex_unlock(&mRefCountMutex);
+    if (deleteSelf) {
+        delete this;
+    }
+    return deleteSelf;
+}
diff --git a/libsysutils/src/SocketListener.cpp b/libsysutils/src/SocketListener.cpp
index 611d5fe..fcad624 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;
@@ -213,8 +213,9 @@ void SocketListener::runListener() {
             it = pendingList->begin();
             SocketClient* c = *it;
             pendingList->erase(it);
-            /* Process it, if false is returned, remove and destroy it */
-            if (!onDataAvailable(c)) {
+            /* Process it, if false is returned and our sockets are
+             * connection-based, remove and destroy it */
+            if (!onDataAvailable(c) && mListen) {
                 /* Remove the client from our array */
                 pthread_mutex_lock(&mClientsLock);
                 for (it = mClients->begin(); it != mClients->end(); ++it) {
@@ -225,8 +226,11 @@ void SocketListener::runListener() {
                 }
                 pthread_mutex_unlock(&mClientsLock);
                 /* Destroy the client */
-                close(c->getSocket());
-                delete c;
+                int socket = c->getSocket();
+                if (c->decRef()) {
+                    // Note: 'c' is deleted memory at this point.
+                    close(socket);
+                }
             }
         }
     }