1 files changed, 68 insertions, 16 deletions

diff --git a/libsysutils/src/SocketClient.cpp b/libsysutils/src/SocketClient.cpp
index 8e5f154..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);
@@ -32,14 +34,24 @@ SocketClient::SocketClient(int socket)
 
 int SocketClient::sendMsg(int code, const char *msg, bool addErrno) {
     char *buf;
+    const char* arg;
+    const char* fmt;
+    char tmp[1];
+    int  len;
 
     if (addErrno) {
-        buf = (char *) alloca(strlen(msg) + strlen(strerror(errno)) + 8);
-        sprintf(buf, "%.3d %s (%s)", code, msg, strerror(errno));
+        fmt = "%.3d %s (%s)";
+        arg = strerror(errno);
     } else {
-        buf = (char *) alloca(strlen(msg) + strlen("XXX "));
-        sprintf(buf, "%.3d %s", code, msg);
+        fmt = "%.3d %s";
+        arg = NULL;
     }
+    /* Measure length of required buffer */
+    len = snprintf(tmp, sizeof tmp, fmt, code, msg, arg);
+    /* Allocate in the stack, then write to it */
+    buf = (char*)alloca(len+1);
+    snprintf(buf, len+1, fmt, code, msg, arg);
+    /* Send the zero-terminated message */
     return sendMsg(buf);
 }
 
@@ -50,25 +62,65 @@ int SocketClient::sendMsg(const char *msg) {
     }
 
     // Send the message including null character
+    if (sendData(msg, strlen(msg) + 1) != 0) {
+        SLOGW("Unable to send msg '%s'", msg);
+        return -1;
+    }
+    return 0;
+}
+
+int SocketClient::sendData(const void* data, int len) {
     int rc = 0;
-    const char *p = msg;
-    int brtw = strlen(msg) + 1;
+    const char *p = (const char*) data;
+    int brtw = len;
+
+    if (len == 0) {
+        return 0;
+    }
 
     pthread_mutex_lock(&mWriteMutex);
-    while(brtw) {
-        if ((rc = write(mSocket,p, brtw)) < 0) {
-            SLOGW("Unable to send msg '%s' (%s)", msg, strerror(errno));
-            pthread_mutex_unlock(&mWriteMutex);
-            return -1;
-        } else if (!rc) {
+    while (brtw > 0) {
+        rc = write(mSocket, p, brtw);
+        if (rc > 0) {
+            p += rc;
+            brtw -= rc;
+            continue;
+        }
+
+        if (rc < 0 && errno == EINTR)
+            continue;
+
+        pthread_mutex_unlock(&mWriteMutex);
+        if (rc == 0) {
             SLOGW("0 length write :(");
             errno = EIO;
-            pthread_mutex_unlock(&mWriteMutex);
-            return -1;
+        } else {
+            SLOGW("write error (%s)", strerror(errno));
         }
-        p += rc;
-        brtw -= rc;
+        return -1;
     }
     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;
+}