Fix memory leak on malloc errors

Wiebe Cazemier
1 parent 84e76c32
Showing 2 changed files with 19 additions and 4 deletions
client.cpp
client.h
@@ -11,11 +11,24 @@ Client::Client(int fd, ThreadData_p threadData) :
 {
     int flags = fcntl(fd, F_GETFL);
     fcntl(fd, F_SETFL, flags | O_NONBLOCK);
-    readbuf = (char*)malloc(CLIENT_BUFFER_SIZE);
-    writebuf = (char*)malloc(CLIENT_BUFFER_SIZE);
+    char *readbuf = (char*)malloc(CLIENT_BUFFER_SIZE);
+    char *writebuf = (char*)malloc(CLIENT_BUFFER_SIZE);
     if (readbuf == NULL || writebuf == NULL)
+    {
+        if (readbuf != NULL)
+            free(readbuf);
+        if (writebuf != NULL)
+            free(writebuf);
+
+        readbuf = NULL;
+        writebuf = NULL;
+
         throw std::runtime_error("Malloc error constructing client.");
+    }
+
+    this->readbuf = readbuf;
+    this->writebuf = writebuf;
 }
 Client::~Client()
@@ -65,9 +65,10 @@ class Client
     void growReadBuffer()
     {
         const size_t newBufSize = readBufsize * 2;
-        readbuf = (char*)realloc(readbuf, newBufSize);
+        char *readbuf = (char*)realloc(this->readbuf, newBufSize);
         if (readbuf == NULL)
             throw std::runtime_error("Memory allocation failure in growReadBuffer()");
+        this->readbuf = readbuf;
         readBufsize = newBufSize;
         //std::cout << "New read buf size: " << readBufsize << std::endl;
@@ -92,11 +93,12 @@ class Client
         const size_t grow_by = std::max<size_t>(add_size, writeBufsize*2);
         const size_t newBufSize = writeBufsize + grow_by;
-        writebuf = (char*)realloc(writebuf, newBufSize);
+        char *writebuf = (char*)realloc(this->writebuf, newBufSize);
         if (writebuf == NULL)
             throw std::runtime_error("Memory allocation failure in growWriteBuffer()");
+        this->writebuf = writebuf;
         writeBufsize = newBufSize;
         //std::cout << "New write buf size: " << writeBufsize << std::endl;