Allow for realloc to be intercepted too

Previous change to allow interposition of user-provided malloc/free didn't cover use of realloc().
daviddrysdale · Nov 24, 2015 · 492a1cc · 492a1cc
1 parent 2c5f595
commit 492a1cc
Show file tree

Hide file tree

Showing 10 changed files with 45 additions and 17 deletions.
diff --git a/ares.h b/ares.h
@@ -298,7 +298,8 @@ CARES_EXTERN int ares_library_init(int flags);
 
 CARES_EXTERN int ares_library_init_mem(int flags,
                                        void *(*amalloc)(size_t size),
-                                       void (*afree)(void *ptr));
+                                       void (*afree)(void *ptr),
+                                       void *(*arealloc)(void *ptr, size_t size));
 
 CARES_EXTERN int ares_library_initialized(void);
 

diff --git a/ares__read_line.c b/ares__read_line.c
@@ -59,7 +59,7 @@ int ares__read_line(FILE *fp, char **buf, size_t *bufsize)
         continue;
 
       /* Allocate more space. */
-      newbuf = realloc(*buf, *bufsize * 2);
+      newbuf = ares_realloc(*buf, *bufsize * 2);
       if (!newbuf)
         {
           ares_free(*buf);

diff --git a/ares_init.c b/ares_init.c
@@ -868,7 +868,7 @@ static int get_DNS_NetworkParams(char **outptr)
   if ((res != ERROR_BUFFER_OVERFLOW) && (res != ERROR_SUCCESS))
     goto done;
 
-  newfi = realloc(fi, size);
+  newfi = ares_realloc(fi, size);
   if (!newfi)
     goto done;
 
@@ -967,7 +967,7 @@ static int get_DNS_AdaptersAddresses(char **outptr)
   {
     if (Bufsz < ReqBufsz)
     {
-      newipaa = realloc(ipaa, ReqBufsz);
+      newipaa = ares_realloc(ipaa, ReqBufsz);
       if (!newipaa)
         goto done;
       Bufsz = ReqBufsz;
@@ -1399,7 +1399,7 @@ static int init_by_defaults(ares_channel channel)
         char *p;
         len *= 2;
         lenv *= 2;
-        p = realloc(hostname, len);
+        p = ares_realloc(hostname, len);
         if(!p) {
           rc = ARES_ENOMEM;
           goto error;
@@ -1561,8 +1561,8 @@ static int config_nameserver(struct server_state **servers, int *nservers,
         continue;
 
       /* Resize servers state array. */
-      newserv = realloc(*servers, (*nservers + 1) *
-                        sizeof(struct server_state));
+      newserv = ares_realloc(*servers, (*nservers + 1) *
+                             sizeof(struct server_state));
       if (!newserv)
         return ARES_ENOMEM;
 
@@ -1839,7 +1839,7 @@ static int sortlist_alloc(struct apattern **sortlist, int *nsort,
                           struct apattern *pat)
 {
   struct apattern *newsort;
-  newsort = realloc(*sortlist, (*nsort + 1) * sizeof(struct apattern));
+  newsort = ares_realloc(*sortlist, (*nsort + 1) * sizeof(struct apattern));
   if (!newsort)
     return 0;
   newsort[*nsort] = *pat;

diff --git a/ares_library_init.c b/ares_library_init.c
@@ -36,6 +36,7 @@ static int          ares_init_flags;
 
 /* library-private global vars with visibility across the whole library */
 void *(*ares_malloc)(size_t size) = malloc;
+void *(*ares_realloc)(void *ptr, size_t size) = realloc;
 void (*ares_free)(void *ptr) = free;
 
 #ifdef USE_WINSOCK
@@ -125,10 +126,15 @@ int ares_library_init(int flags)
 
 int ares_library_init_mem(int flags,
                           void *(*amalloc)(size_t size),
-                          void (*afree)(void *ptr))
+                          void (*afree)(void *ptr),
+                          void *(*arealloc)(void *ptr, size_t size))
 {
-  ares_malloc = amalloc;
-  ares_free = afree;
+  if (amalloc)
+    ares_malloc = amalloc;
+  if (arealloc)
+    ares_realloc = arealloc;
+  if (afree)
+    ares_free = afree;
   return ares_library_init(flags);
 }
 

diff --git a/ares_parse_ptr_reply.c b/ares_parse_ptr_reply.c
@@ -136,7 +136,7 @@ int ares_parse_ptr_reply(const unsigned char *abuf, int alen, const void *addr,
           if (aliascnt >= alias_alloc) {
             char **ptr;
             alias_alloc *= 2;
-            ptr = realloc(aliases, alias_alloc * sizeof(char *));
+            ptr = ares_realloc(aliases, alias_alloc * sizeof(char *));
             if(!ptr) {
               ares_free(rr_name);
               status = ARES_ENOMEM;

diff --git a/ares_private.h b/ares_private.h
@@ -306,6 +306,7 @@ struct ares_channeldata {
 
 /* Memory management functions */
 extern void *(*ares_malloc)(size_t size);
+extern void *(*ares_realloc)(void *ptr, size_t size);
 extern void (*ares_free)(void *ptr);
 
 /* return true if now is exactly check time or later */

diff --git a/ares_process.c b/ares_process.c
@@ -570,7 +570,7 @@ static void process_answer(ares_channel channel, unsigned char *abuf,
           query->tcpbuf[0] = (unsigned char)((qlen >> 8) & 0xff);
           query->tcpbuf[1] = (unsigned char)(qlen & 0xff);
           DNS_HEADER_SET_ARCOUNT(query->tcpbuf + 2, 0);
-          query->tcpbuf = realloc(query->tcpbuf, query->tcplen);
+          query->tcpbuf = ares_realloc(query->tcpbuf, query->tcplen);
           ares__send_query(channel, query, now);
           return;
       }

diff --git a/test/ares-test-init.cc b/test/ares-test-init.cc
@@ -125,7 +125,8 @@ TEST_F(LibraryTest, FailChannelInit) {
   EXPECT_EQ(ARES_SUCCESS,
             ares_library_init_mem(ARES_LIB_INIT_ALL,
                                   &LibraryTest::amalloc,
-                                  &LibraryTest::afree));
+                                  &LibraryTest::afree,
+                                  &LibraryTest::arealloc));
   SetAllocFail(1);
   ares_channel channel = nullptr;
   EXPECT_EQ(ARES_ENOMEM, ares_init(&channel));

diff --git a/test/ares-test.cc b/test/ares-test.cc
@@ -74,22 +74,38 @@ void LibraryTest::ClearFails() {
   size_fails_.clear();
 }
 
+
 // static
-void* LibraryTest::amalloc(size_t size) {
+bool LibraryTest::ShouldAllocFail(size_t size) {
   bool fail = (fails_ & 0x01);
   fails_ >>= 1;
   if (size_fails_[size] > 0) {
     size_fails_[size]--;
     fail = true;
   }
-  if (fail) {
+  return fail;
+}
+
+// static
+void* LibraryTest::amalloc(size_t size) {
+  if (ShouldAllocFail(size)) {
     if (verbose) std::cerr << "Failing malloc(" << size << ") request" << std::endl;
     return nullptr;
   } else {
     return malloc(size);
   }
 }
 
+// static
+void* LibraryTest::arealloc(void *ptr, size_t size) {
+  if (ShouldAllocFail(size)) {
+    if (verbose) std::cerr << "Failing realloc(" << ptr << ", " << size << ") request" << std::endl;
+    return nullptr;
+  } else {
+    return realloc(ptr, size);
+  }
+}
+
 // static
 void LibraryTest::afree(void *ptr) {
   free(ptr);

diff --git a/test/ares-test.h b/test/ares-test.h
@@ -37,7 +37,8 @@ class LibraryTest : public ::testing::Test {
     EXPECT_EQ(ARES_SUCCESS,
               ares_library_init_mem(ARES_LIB_INIT_ALL,
                                     &LibraryTest::amalloc,
-                                    &LibraryTest::afree));
+                                    &LibraryTest::afree,
+                                    &LibraryTest::arealloc));
   }
   ~LibraryTest() {
     ares_library_cleanup();
@@ -52,8 +53,10 @@ class LibraryTest : public ::testing::Test {
   static void ClearFails();
 
   static void *amalloc(size_t size);
+  static void* arealloc(void *ptr, size_t size);
   static void afree(void *ptr);
  private:
+  static bool ShouldAllocFail(size_t size);
   static unsigned long fails_;
   static std::map<size_t, int> size_fails_;
 };