CDRIVER-5669 improve string function handling of large strings (#1696)

* add `bson_string_ensure_space` to check addition and casts.

* add private `bson_next_power_of_two_u32` to avoid casts.

* document length expectations

* add test for maximum capacity of `bson_string_t`

Skip test by default. Large allocations fail on TSan tasks, and are slow on ASan tasks. Test may be used to verify behavior locally.

---------

Co-authored-by: Ezra Chung <[email protected]>
Co-authored-by: Roberto C. Sánchez <[email protected]>

Author: 3 people

CONTRIBUTING.md

@@ -318,6 +318,10 @@ To test a load balanced deployment, set the following environment variables:
 * `SINGLE_MONGOS_LB_URI=<string>` to a MongoDB URI with a host of a load balancer fronting one mongos.
 * `MULTI_MONGOS_LB_URI=<string>` to a MongoDB URI with a host of a load balancer fronting multiple mongos processes.
 
+To run test cases with large allocations, set:
+
+* `MONGOC_TEST_LARGE_ALLOCATIONS=on` This may result in sudden test suite termination due to allocation failure. Use with caution.
+
 All tests should pass before submitting a patch.
 
 ## Configuring the test runner

src/libbson/doc/bson_string_append.rst

@@ -22,3 +22,4 @@ Description
 
 Appends the ASCII or UTF-8 encoded string ``str`` to ``string``. This is not suitable for embedding NULLs in strings.
 
+.. warning:: This function will abort if the length of the resulting string (including the NULL terminator) would exceed ``UINT32_MAX``.

src/libbson/doc/bson_string_append_c.rst

@@ -22,3 +22,4 @@ Description
 
 Appends ``c`` to the string builder ``string``.
 
+.. warning:: This function will abort if the length of the resulting string (including the NULL terminator) would exceed ``UINT32_MAX``.

src/libbson/doc/bson_string_append_printf.rst

@@ -23,3 +23,4 @@ Description
 
 Like bson_string_append() but formats a printf style string and then appends that to ``string``.
 
+.. warning:: This function will abort if the length of the resulting string (including the NULL terminator) would exceed ``UINT32_MAX``.

src/libbson/doc/bson_string_append_unichar.rst

@@ -22,3 +22,4 @@ Description
 
 Appends a unicode character to string. The character will be encoded into its multi-byte UTF-8 representation.
 
+.. warning:: This function will abort if the length of the resulting string (including the NULL terminator) would exceed ``UINT32_MAX``.

src/libbson/doc/bson_string_new.rst

@@ -21,6 +21,8 @@ Description
 
 Creates a new string builder, which uses power-of-two growth of buffers. Use the various bson_string_append*() functions to append to the string.
 
+.. warning:: This function will abort if the length of the resulting string (including the NULL terminator) would exceed ``UINT32_MAX``.
+
 Returns
 -------
 

src/libbson/src/bson/bson-string.c

@@ -31,6 +31,48 @@
 #include <string.h>
 #endif
 
+// `bson_next_power_of_two_u32` returns 0 on overflow.
+static BSON_INLINE uint32_t
+bson_next_power_of_two_u32 (uint32_t v)
+{
+   BSON_ASSERT (v > 0);
+
+   // https://graphics.stanford.edu/%7Eseander/bithacks.html#RoundUpPowerOf2
+   v--;
+   v |= v >> 1;
+   v |= v >> 2;
+   v |= v >> 4;
+   v |= v >> 8;
+   v |= v >> 16;
+   v++;
+
+   return v;
+}
+
+// `bson_string_ensure_space` ensures `string` has enough room for `needed` + a null terminator.
+static void
+bson_string_ensure_space (bson_string_t *string, uint32_t needed)
+{
+   BSON_ASSERT_PARAM (string);
+   BSON_ASSERT (needed <= UINT32_MAX - 1u);
+   needed += 1u; // Add one for trailing NULL byte.
+   if (string->alloc >= needed) {
+      return;
+   }
+   // Get the next largest power of 2 if possible.
+   uint32_t alloc = bson_next_power_of_two_u32 (needed);
+   if (alloc == 0) {
+      // Overflowed: saturate at UINT32_MAX.
+      alloc = UINT32_MAX;
+   }
+   if (!string->str) {
+      string->str = bson_malloc (alloc);
+   } else {
+      string->str = bson_realloc (string->str, alloc);
+   }
+   string->alloc = alloc;
+}
+
 /*
  *--------------------------------------------------------------------------
  *
@@ -62,33 +104,18 @@ bson_string_t *
 bson_string_new (const char *str) /* IN */
 {
    bson_string_t *ret;
-   size_t len_sz;
 
    ret = bson_malloc0 (sizeof *ret);
+   const size_t len_sz = str == NULL ? 0u : strlen (str);
+   BSON_ASSERT (bson_in_range_unsigned (uint32_t, len_sz));
+   const uint32_t len_u32 = (uint32_t) len_sz;
+   bson_string_ensure_space (ret, len_u32);
    if (str) {
-      len_sz = strlen (str);
-      BSON_ASSERT (len_sz <= UINT32_MAX);
-      ret->len = (uint32_t) len_sz;
-   } else {
-      ret->len = 0;
-   }
-   ret->alloc = ret->len + 1;
-
-   if (!bson_is_power_of_two (ret->alloc)) {
-      len_sz = bson_next_power_of_two ((size_t) ret->alloc);
-      BSON_ASSERT (len_sz <= UINT32_MAX);
-      ret->alloc = (uint32_t) len_sz;
-   }
-
-   BSON_ASSERT (ret->alloc >= ret->len + 1);
-
-   ret->str = bson_malloc (ret->alloc);
-
-   if (str) {
-      memcpy (ret->str, str, ret->len);
+      memcpy (ret->str, str, len_sz);
    }
 
-   ret->str[ret->len] = '\0';
+   ret->str[len_u32] = '\0';
+   ret->len = len_u32;
 
    return ret;
 }
@@ -135,31 +162,18 @@ void
 bson_string_append (bson_string_t *string, /* IN */
                     const char *str)       /* IN */
 {
-   uint32_t len;
-   size_t len_sz;
-
    BSON_ASSERT (string);
    BSON_ASSERT (str);
 
-   len_sz = strlen (str);
+   const size_t len_sz = strlen (str);
    BSON_ASSERT (bson_in_range_unsigned (uint32_t, len_sz));
-   len = (uint32_t) len_sz;
-
-   if ((string->alloc - string->len - 1) < len) {
-      BSON_ASSERT (string->alloc <= UINT32_MAX - len);
-      string->alloc += len;
-      if (!bson_is_power_of_two (string->alloc)) {
-         len_sz = bson_next_power_of_two ((size_t) string->alloc);
-         BSON_ASSERT (len_sz <= UINT32_MAX);
-         string->alloc = (uint32_t) len_sz;
-      }
-      BSON_ASSERT (string->alloc >= string->len + len);
-      string->str = bson_realloc (string->str, string->alloc);
-   }
-
-   memcpy (string->str + string->len, str, len);
-   string->len += len;
-   string->str[string->len] = '\0';
+   const uint32_t len_u32 = (uint32_t) len_sz;
+   BSON_ASSERT (len_u32 <= UINT32_MAX - string->len);
+   const uint32_t new_len = len_u32 + string->len;
+   bson_string_ensure_space (string, new_len);
+   memcpy (string->str + string->len, str, len_sz);
+   string->str[new_len] = '\0';
+   string->len = new_len;
 }
 
 

src/libbson/tests/test-string.c

@@ -18,6 +18,7 @@
 #include <bson/bson.h>
 
 #include "TestSuite.h"
+#include "test-libmongoc.h"
 
 
 static void
@@ -303,6 +304,68 @@ test_bson_strcasecmp (void)
    BSON_ASSERT (bson_strcasecmp ("FoZ", "foo") > 0);
 }
 
+static void
+test_bson_string_capacity (void *unused)
+{
+   BSON_UNUSED (unused);
+
+   char *large_str = bson_malloc (UINT32_MAX);
+   memset (large_str, 's', UINT32_MAX); // Do not NULL terminate. Each test case sets NULL byte.
+
+   // Test the largest possible string that can be constructed.
+   {
+      large_str[UINT32_MAX - 1u] = '\0'; // Set size.
+      bson_string_t *str = bson_string_new (large_str);
+      bson_string_free (str, true);
+      large_str[UINT32_MAX - 1u] = 's'; // Restore.
+   }
+
+   // Test appending with `bson_string_append` to get maximum size.
+   {
+      large_str[UINT32_MAX - 1u] = '\0'; // Set size.
+      bson_string_t *str = bson_string_new ("");
+      bson_string_append (str, large_str);
+      bson_string_free (str, true);
+      large_str[UINT32_MAX - 1u] = 's'; // Restore.
+   }
+
+   // Test appending with `bson_string_append_c` to get maximum size.
+   {
+      large_str[UINT32_MAX - 2u] = '\0'; // Set size.
+      bson_string_t *str = bson_string_new (large_str);
+      bson_string_append_c (str, 'c');
+      bson_string_free (str, true);
+      large_str[UINT32_MAX - 2u] = 's'; // Restore.
+   }
+
+   // Test appending with `bson_string_append_printf` to get maximum size.
+   {
+      large_str[UINT32_MAX - 2u] = '\0'; // Set size.
+      bson_string_t *str = bson_string_new (large_str);
+      bson_string_append_printf (str, "c");
+      bson_string_free (str, true);
+      large_str[UINT32_MAX - 2u] = 's'; // Restore.
+   }
+
+   // Test appending with single characters.
+   {
+      large_str[UINT32_MAX - 2u] = '\0'; // Set size.
+      bson_string_t *str = bson_string_new (large_str);
+      bson_string_append_unichar (str, (bson_unichar_t) 's');
+      bson_string_free (str, true);
+      large_str[UINT32_MAX - 2u] = 's'; // Restore.
+   }
+
+   bson_free (large_str);
+}
+
+static int
+skip_if_no_large_allocations (void)
+{
+   // Skip tests requiring large allocations.
+   // Large allocations were observed to fail when run with TSan, and are time consuming with ASan.
+   return test_framework_getenv_bool ("MONGOC_TEST_LARGE_ALLOCATIONS");
+}
 
 void
 test_string_install (TestSuite *suite)
@@ -320,4 +383,6 @@ test_string_install (TestSuite *suite)
    TestSuite_Add (suite, "/bson/string/snprintf", test_bson_snprintf);
    TestSuite_Add (suite, "/bson/string/strnlen", test_bson_strnlen);
    TestSuite_Add (suite, "/bson/string/strcasecmp", test_bson_strcasecmp);
+   TestSuite_AddFull (
+      suite, "/bson/string/capacity", test_bson_string_capacity, NULL, NULL, skip_if_no_large_allocations);
 }