Fix memory leak in string_map.c in rcutils (ros2#411)

* Add in a string_map test for failure after certain allocations.

Signed-off-by: Chris Lalancette <clalancette@gmail.com>
Signed-off-by: Nikolai Morin <nikolai.morin@apex.ai>

Author: nnmm

src/string_map.c

@@ -29,10 +29,15 @@ extern "C"
 #include "rcutils/format_string.h"
 #include "rcutils/types/rcutils_ret.h"
 
+typedef struct key_value_pair
+{
+  char * key;
+  char * value;
+} key_value_pair_t;
+
 typedef struct rcutils_string_map_impl_s
 {
-  char ** keys;
-  char ** values;
+  key_value_pair_t * key_value_pairs;
   size_t capacity;
   size_t size;
   rcutils_allocator_t allocator;
@@ -64,8 +69,7 @@ rcutils_string_map_init(
     RCUTILS_SET_ERROR_MSG("failed to allocate memory for string map impl struct");
     return RCUTILS_RET_BAD_ALLOC;
   }
-  string_map->impl->keys = NULL;
-  string_map->impl->values = NULL;
+  string_map->impl->key_value_pairs = NULL;
   string_map->impl->capacity = 0;
   string_map->impl->size = 0;
   string_map->impl->allocator = allocator;
@@ -130,6 +134,17 @@ rcutils_string_map_get_size(const rcutils_string_map_t * string_map, size_t * si
   return RCUTILS_RET_OK;
 }
 
+static void
+__remove_key_and_value_at_index(rcutils_string_map_impl_t * string_map_impl, size_t index)
+{
+  rcutils_allocator_t allocator = string_map_impl->allocator;
+  allocator.deallocate(string_map_impl->key_value_pairs[index].key, allocator.state);
+  string_map_impl->key_value_pairs[index].key = NULL;
+  allocator.deallocate(string_map_impl->key_value_pairs[index].value, allocator.state);
+  string_map_impl->key_value_pairs[index].value = NULL;
+  string_map_impl->size--;
+}
+
 rcutils_ret_t
 rcutils_string_map_reserve(rcutils_string_map_t * string_map, size_t capacity)
 {
@@ -147,46 +162,35 @@ rcutils_string_map_reserve(rcutils_string_map_t * string_map, size_t capacity)
     return RCUTILS_RET_OK;
   } else if (capacity == 0) {
     // if the requested capacity is zero, then make sure the existing keys and values are free'd
-    allocator.deallocate(string_map->impl->keys, allocator.state);
-    string_map->impl->keys = NULL;
-    allocator.deallocate(string_map->impl->values, allocator.state);
-    string_map->impl->values = NULL;
+    // size is known to be 0 here because of the recursive call above.
+    allocator.deallocate(string_map->impl->key_value_pairs, allocator.state);
+    string_map->impl->key_value_pairs = NULL;
     // falls through to normal function end
   } else {
     // if the capacity non-zero and different, use realloc to increase/shrink the size
     // note that realloc when the pointer is NULL is the same as malloc
     // note also that realloc will shrink the space if needed
 
-    // ensure that reallocate won't overflow capacity
-    if (capacity > (SIZE_MAX / sizeof(char *))) {
+    // ensure that reallocate won't overflow SIZE_MAX
+    if (capacity > (SIZE_MAX / sizeof(key_value_pair_t))) {
       RCUTILS_SET_ERROR_MSG("requested capacity for string_map too large");
       return RCUTILS_RET_BAD_ALLOC;
     }
 
-    // resize the keys, assigning the result only if it succeeds
-    char ** new_keys =
-      allocator.reallocate(string_map->impl->keys, capacity * sizeof(char *), allocator.state);
-    if (NULL == new_keys) {
-      RCUTILS_SET_ERROR_MSG("failed to allocate memory for string_map keys");
-      return RCUTILS_RET_BAD_ALLOC;
-    }
-    string_map->impl->keys = new_keys;
-
-    // resize the values, assigning the result only if it succeeds
-    char ** new_values =
-      allocator.reallocate(string_map->impl->values, capacity * sizeof(char *), allocator.state);
-    if (NULL == new_values) {
-      RCUTILS_SET_ERROR_MSG("failed to allocate memory for string_map values");
+    // resize the keys and values, assigning the result only if it succeeds
+    key_value_pair_t * new_key_value_pairs = allocator.reallocate(
+      string_map->impl->key_value_pairs, capacity * sizeof(key_value_pair_t), allocator.state);
+    if (NULL == new_key_value_pairs) {
+      RCUTILS_SET_ERROR_MSG("failed to allocate memory for string_map key-value pairs");
       return RCUTILS_RET_BAD_ALLOC;
     }
-    string_map->impl->values = new_values;
+    string_map->impl->key_value_pairs = new_key_value_pairs;
 
     // zero out the new memory, if there is any (expanded instead of shrunk)
     if (capacity > string_map->impl->capacity) {
-      size_t i = string_map->impl->capacity;
-      for (; i < capacity; ++i) {
-        string_map->impl->keys[i] = NULL;
-        string_map->impl->values[i] = NULL;
+      for (size_t i = string_map->impl->capacity; i < capacity; ++i) {
+        string_map->impl->key_value_pairs[i].key = NULL;
+        string_map->impl->key_value_pairs[i].value = NULL;
       }
     }
     // falls through to normal function end
@@ -195,29 +199,19 @@ rcutils_string_map_reserve(rcutils_string_map_t * string_map, size_t capacity)
   return RCUTILS_RET_OK;
 }
 
-static void
-__remove_key_and_value_at_index(rcutils_string_map_impl_t * string_map_impl, size_t index)
-{
-  rcutils_allocator_t allocator = string_map_impl->allocator;
-  allocator.deallocate(string_map_impl->keys[index], allocator.state);
-  string_map_impl->keys[index] = NULL;
-  allocator.deallocate(string_map_impl->values[index], allocator.state);
-  string_map_impl->values[index] = NULL;
-  string_map_impl->size--;
-}
-
 rcutils_ret_t
 rcutils_string_map_clear(rcutils_string_map_t * string_map)
 {
   RCUTILS_CHECK_ARGUMENT_FOR_NULL(string_map, RCUTILS_RET_INVALID_ARGUMENT);
   RCUTILS_CHECK_FOR_NULL_WITH_MSG(
     string_map->impl, "invalid string map", return RCUTILS_RET_STRING_MAP_INVALID);
-  size_t i = 0;
-  for (; i < string_map->impl->capacity; ++i) {
-    if (string_map->impl->keys[i] != NULL) {
+
+  for (size_t i = 0; i < string_map->impl->capacity; ++i) {
+    if (string_map->impl->key_value_pairs[i].key != NULL) {
       __remove_key_and_value_at_index(string_map->impl, i);
     }
   }
+
   return RCUTILS_RET_OK;
 }
 
@@ -255,14 +249,14 @@ __get_index_of_key_if_exists(
 {
   size_t i = 0;
   for (; i < string_map_impl->capacity; ++i) {
-    if (NULL == string_map_impl->keys[i]) {
+    if (NULL == string_map_impl->key_value_pairs[i].key) {
       continue;
     }
-    size_t cmp_count = strlen(string_map_impl->keys[i]);
+    size_t cmp_count = strlen(string_map_impl->key_value_pairs[i].key);
     if (key_length > cmp_count) {
       cmp_count = key_length;
     }
-    if (strncmp(key, string_map_impl->keys[i], cmp_count) == 0) {
+    if (strncmp(key, string_map_impl->key_value_pairs[i].key, cmp_count) == 0) {
       *index = i;
       return true;
     }
@@ -292,30 +286,30 @@ rcutils_string_map_set_no_resize(
       return RCUTILS_RET_NOT_ENOUGH_SPACE;
     }
     for (key_index = 0; key_index < string_map->impl->capacity; ++key_index) {
-      if (NULL == string_map->impl->keys[key_index]) {
+      if (NULL == string_map->impl->key_value_pairs[key_index].key) {
         break;
       }
     }
     assert(key_index < string_map->impl->capacity);  // defensive, this should not happen
-    string_map->impl->keys[key_index] = rcutils_strdup(key, allocator);
-    if (NULL == string_map->impl->keys[key_index]) {
+    string_map->impl->key_value_pairs[key_index].key = rcutils_strdup(key, allocator);
+    if (NULL == string_map->impl->key_value_pairs[key_index].key) {
       RCUTILS_SET_ERROR_MSG("failed to allocate memory for key");
       return RCUTILS_RET_BAD_ALLOC;
     }
     should_free_key_on_error = true;
   }
   // at this point the key is in the map, waiting for the value to set/overwritten
-  char * original_value = string_map->impl->values[key_index];
+  char * original_value = string_map->impl->key_value_pairs[key_index].value;
   char * new_value = rcutils_strdup(value, allocator);
   if (NULL == new_value) {
-    RCUTILS_SET_ERROR_MSG("failed to allocate memory for key");
+    RCUTILS_SET_ERROR_MSG("failed to allocate memory for value");
     if (should_free_key_on_error) {
-      allocator.deallocate(string_map->impl->keys[key_index], allocator.state);
-      string_map->impl->keys[key_index] = NULL;
+      allocator.deallocate(string_map->impl->key_value_pairs[key_index].key, allocator.state);
+      string_map->impl->key_value_pairs[key_index].key = NULL;
     }
     return RCUTILS_RET_BAD_ALLOC;
   }
-  string_map->impl->values[key_index] = new_value;
+  string_map->impl->key_value_pairs[key_index].value = new_value;
   if (original_value != NULL) {
     // clean up the old value if not NULL
     allocator.deallocate(original_value, allocator.state);
@@ -386,7 +380,7 @@ rcutils_string_map_getn(
   }
   size_t key_index;
   if (__get_index_of_key_if_exists(string_map->impl, key, key_length, &key_index)) {
-    return string_map->impl->values[key_index];
+    return string_map->impl->key_value_pairs[key_index].value;
   }
   return NULL;
 }
@@ -408,7 +402,7 @@ rcutils_string_map_get_next_key(
     bool given_key_found = false;
     size_t i = 0;
     for (; i < string_map->impl->capacity; ++i) {
-      if (string_map->impl->keys[i] == key) {
+      if (string_map->impl->key_value_pairs[i].key == key) {
         given_key_found = true;
         // given key found at index i, start there + 1
         start_index = i + 1;
@@ -422,9 +416,9 @@ rcutils_string_map_get_next_key(
   // iterate through the storage and look for another non-NULL key to return
   size_t i = start_index;
   for (; i < string_map->impl->capacity; ++i) {
-    if (string_map->impl->keys[i] != NULL) {
+    if (string_map->impl->key_value_pairs[i].key != NULL) {
       // next key found, return it
-      return string_map->impl->keys[i];
+      return string_map->impl->key_value_pairs[i].key;
     }
   }
   // next key (or first key) not found

test/test_string_map.cpp

@@ -1939,3 +1939,64 @@ TEST(test_string_map, strange_keys) {
     EXPECT_STREQ("value1", rcutils_string_map_get(&string_map, "key with spaces"));
   }
 }
+
+static int realloc_counter = 0;
+static int realloc_fail_after = -1;
+
+static void *
+reallocate_fail_after(void * pointer, size_t size, void * state)
+{
+  RCUTILS_UNUSED(state);
+
+  if (realloc_fail_after >= 0 && realloc_counter++ >= realloc_fail_after) {
+    return NULL;
+  }
+
+  return realloc(pointer, size);
+}
+
+TEST(test_string_map, partial_allocation_failures)
+{
+  {
+    rcutils_string_map_t map = rcutils_get_zero_initialized_string_map();
+
+    rcutils_allocator_t allocator = rcutils_get_default_allocator();
+    allocator.reallocate = reallocate_fail_after;
+
+    realloc_counter = 0;
+    realloc_fail_after = 0;
+
+    rcutils_ret_t ret;
+
+    // Since we set realloc_fail_after to 0, that means we'll fail immediately
+    // and hence rcutils_string_map_init should return BAD_ALLOC
+    ret = rcutils_string_map_init(&map, 4, allocator);
+    ASSERT_EQ(ret, RCUTILS_RET_BAD_ALLOC);
+  }
+
+  {
+    rcutils_string_map_t map = rcutils_get_zero_initialized_string_map();
+
+    rcutils_allocator_t allocator = rcutils_get_default_allocator();
+    allocator.reallocate = reallocate_fail_after;
+
+    realloc_counter = 0;
+    realloc_fail_after = 1;
+
+    rcutils_ret_t ret;
+
+    // Since we set realloc_fail_after to 1, that means we'll succeed on the
+    // first allocation, so rcutils_string_map_init should return OK
+    ret = rcutils_string_map_init(&map, 4, allocator);
+    ASSERT_EQ(ret, RCUTILS_RET_OK);
+
+    // but reserve should fail
+    ret = rcutils_string_map_reserve(&map, 2);
+    ASSERT_EQ(ret, RCUTILS_RET_BAD_ALLOC);
+
+    // and we should still be able to fini things without accessing
+    // out-of-bounds memory or leaking.
+    ret = rcutils_string_map_fini(&map);
+    ASSERT_EQ(ret, RCUTILS_RET_OK);
+  }
+}