Tests to handle allocation errors gracefully.

Based on patch by Jason Lunz.



git-svn-id: https://protobuf-c.googlecode.com/svn/trunk@173 00440858-1255-0410-a3e6-75ea37f81c3a

src/google/protobuf-c/protobuf-c.c

@@ -27,8 +27,8 @@
 
 /* convenience macros */
 #define TMPALLOC(allocator, size) ((allocator)->tmp_alloc ((allocator)->allocator_data, (size)))
-#define ALLOC(allocator, size) ((allocator)->alloc ((allocator)->allocator_data, (size)))
-#define FREE(allocator, ptr)   ((allocator)->free ((allocator)->allocator_data, (ptr)))
+#define FREE(allocator, ptr)   \
+   do { if ((ptr) != NULL) ((allocator)->free ((allocator)->allocator_data, (ptr))); } while(0)
 #define UNALIGNED_ALLOC(allocator, size) ALLOC (allocator, size) /* placeholder */
 #define STRUCT_MEMBER_P(struct_p, struct_offset)   \
     ((void *) ((uint8_t*) (struct_p) + (struct_offset)))
@@ -39,6 +39,30 @@
 #define TRUE 1
 #define FALSE 0
 
+static void
+alloc_failed_warning (unsigned size, const char *filename, unsigned line)
+{
+  fprintf (stderr,
+           "WARNING: out-of-memory allocating a block of size %u (%s:%u)\n",
+           size, filename, line);
+}
+
+/* Try to allocate memory, running some special code if it fails. */
+#define DO_ALLOC(dst, allocator, size, fail_code)                           \
+{ size_t da__allocation_size = (size);                                      \
+  if (da__allocation_size == 0)                                             \
+    dst = NULL;                                                             \
+  else if ((dst=((allocator)->alloc ((allocator)->allocator_data,           \
+                                     da__allocation_size))) == NULL)        \
+    {                                                                       \
+      alloc_failed_warning (da__allocation_size, __FILE__, __LINE__);       \
+      fail_code;                                                            \
+    }                                                                       \
+}
+#define DO_UNALIGNED_ALLOC  DO_ALLOC            /* placeholder */
+
+
+
 #define ASSERT_IS_ENUM_DESCRIPTOR(desc) \
   assert((desc)->magic == PROTOBUF_C_ENUM_DESCRIPTOR_MAGIC)
 #define ASSERT_IS_MESSAGE_DESCRIPTOR(desc) \
@@ -108,7 +132,7 @@ protobuf_c_buffer_simple_append (ProtobufCBuffer *buffer,
       uint8_t *new_data;
       while (new_alloced < new_len)
         new_alloced += new_alloced;
-      new_data = ALLOC (&protobuf_c_default_allocator, new_alloced);
+      DO_ALLOC (new_data, &protobuf_c_default_allocator, new_alloced, return);
       memcpy (new_data, simp->data, simp->len);
       if (simp->must_free_data)
         FREE (&protobuf_c_default_allocator, simp->data);
@@ -1189,7 +1213,7 @@ parse_required_member (ScannedMember *scanned_member,
             if (*pstr != NULL && *pstr != def)
               FREE (allocator, *pstr);
           }
-        *pstr = ALLOC (allocator, len - pref_len + 1);
+        DO_ALLOC (*pstr, allocator, len - pref_len + 1, return 0);
         memcpy (*pstr, data + pref_len, len - pref_len);
         (*pstr)[len-pref_len] = 0;
         return 1;
@@ -1204,7 +1228,7 @@ parse_required_member (ScannedMember *scanned_member,
         def_bd = scanned_member->field->default_value;
         if (maybe_clear && bd->data != NULL && bd->data != def_bd->data)
           FREE (allocator, bd->data);
-        bd->data = ALLOC (allocator, len - pref_len);
+        DO_ALLOC (bd->data, allocator, len - pref_len, return 0);
         memcpy (bd->data, data + pref_len, len - pref_len);
         bd->len = len - pref_len;
         return 1;
@@ -1280,7 +1304,7 @@ parse_member (ScannedMember *scanned_member,
       ufield->tag = scanned_member->tag;
       ufield->wire_type = scanned_member->wire_type;
       ufield->len = scanned_member->len;
-      ufield->data = UNALIGNED_ALLOC (allocator, scanned_member->len);
+      DO_UNALIGNED_ALLOC (ufield->data, allocator, scanned_member->len, return 0);
       memcpy (ufield->data, scanned_member->data, ufield->len);
       return 1;
     }
@@ -1370,7 +1394,7 @@ protobuf_c_message_unpack         (const ProtobufCMessageDescriptor *desc,
 
   if (allocator == NULL)
     allocator = &protobuf_c_default_allocator;
-  rv = ALLOC (allocator, desc->sizeof_message);
+  DO_ALLOC (rv, allocator, desc->sizeof_message, return NULL);
   scanned_member_slabs[0] = first_member_slab;
 
   memset (rv, 0, desc->sizeof_message);
@@ -1480,7 +1504,7 @@ protobuf_c_message_unpack         (const ProtobufCMessageDescriptor *desc,
           if (allocator->tmp_alloc != NULL)
             scanned_member_slabs[which_slab] = TMPALLOC(allocator, size);
           else
-            scanned_member_slabs[which_slab] = ALLOC(allocator, size);
+            DO_ALLOC (scanned_member_slabs[which_slab], allocator, size, goto error_cleanup);
         }
       scanned_member_slabs[which_slab][in_slab_index++] = tmp;
 
@@ -1503,15 +1527,20 @@ protobuf_c_message_unpack         (const ProtobufCMessageDescriptor *desc,
         size_t *n_ptr = STRUCT_MEMBER_PTR (size_t, rv, field->quantifier_offset);
         if (*n_ptr != 0)
           {
-            STRUCT_MEMBER (void *, rv, field->offset) = ALLOC (allocator, siz * (*n_ptr));
+            unsigned n = *n_ptr;
             *n_ptr = 0;
+            DO_ALLOC (STRUCT_MEMBER (void *, rv, field->offset),
+                      allocator, siz * n,
+                      goto error_cleanup);
           }
       }
 
   /* allocate space for unknown fields */
   if (n_unknown)
     {
-      rv->unknown_fields = ALLOC (allocator, n_unknown * sizeof (ProtobufCMessageUnknownField));
+      DO_ALLOC (rv->unknown_fields,
+                allocator, n_unknown * sizeof (ProtobufCMessageUnknownField),
+                goto error_cleanup);
     }
 
   /* do real parsing */

src/test/test-full.proto

@@ -174,3 +174,10 @@ message DefaultOptionalValues {
   optional string v_string = 7 [default = "hi mom\n"];
   optional bytes v_bytes   = 8 [default = "a \0 character"];
 }
+message AllocValues {
+  optional bytes o_bytes = 1;
+  repeated string r_string = 2;
+  required string a_string = 3;
+  required bytes a_bytes = 4;
+  required DefaultRequiredValues a_mess = 5;
+}

src/test/test-generated-code2.c

+#include <stdbool.h>
 #include <stdlib.h>
 #include <stdio.h>
 #include <string.h>
@@ -1177,6 +1178,86 @@ test_optional_default_values (void)
   foo__default_optional_values__free_unpacked (mess2, NULL);
 }
 
+static struct alloc_data {
+  uint32_t alloc_count;
+  int32_t allocs_left;
+} test_allocator_data;
+
+static void *test_alloc(void *allocator_data, size_t size)
+{
+  struct alloc_data *ad = allocator_data;
+  void *rv = NULL;
+  if (ad->allocs_left-- > 0)
+      rv = malloc (size);
+  /* fprintf (stderr, "alloc %d = %p\n", size, rv); */
+  if (rv)
+    ad->alloc_count++;
+  return rv;
+}
+
+static void test_free (void *allocator_data, void *data)
+{
+  struct alloc_data *ad = allocator_data;
+  /* fprintf (stderr, "free %p\n", data); */
+  free (data);
+  if (data)
+    ad->alloc_count--;
+}
+
+static ProtobufCAllocator test_allocator = {
+  .alloc = test_alloc,
+  .free = test_free,
+  .allocator_data = &test_allocator_data,
+};
+
+#define SETUP_TEST_ALLOC_BUFFER(pbuf, len)					\
+  Foo__DefaultRequiredValues _req = FOO__DEFAULT_REQUIRED_VALUES__INIT;		\
+  Foo__AllocValues _mess = FOO__ALLOC_VALUES__INIT;				\
+  _mess.a_string = "some string";						\
+  _mess.r_string = repeated_strings_2;						\
+  _mess.n_r_string = sizeof(repeated_strings_2) / sizeof(*repeated_strings_2);	\
+  uint8_t bytes[] = "some bytes";						\
+  _mess.a_bytes.len = sizeof(bytes);						\
+  _mess.a_bytes.data = bytes;							\
+  _mess.a_mess = &_req;								\
+  size_t len = foo__alloc_values__get_packed_size (&_mess);			\
+  uint8_t *pbuf = malloc (len);							\
+  assert (pbuf);								\
+  size_t _len2 = foo__alloc_values__pack (&_mess, pbuf);			\
+  assert (len == _len2);
+
+static void
+test_alloc_graceful_cleanup (uint8_t *packed, size_t len, int good_allocs)
+{
+  test_allocator_data.alloc_count = 0;
+  test_allocator_data.allocs_left = good_allocs;
+  Foo__AllocValues *mess;
+  mess = foo__alloc_values__unpack (&test_allocator, len, packed);
+  assert (test_allocator_data.allocs_left < 0 ? !mess : !!mess);
+  if (mess)
+    foo__alloc_values__free_unpacked (mess, &test_allocator);
+  assert (0 == test_allocator_data.alloc_count);
+}
+
+static void
+test_alloc_free_all (void)
+{
+  SETUP_TEST_ALLOC_BUFFER (packed, len);
+  test_alloc_graceful_cleanup (packed, len, INT32_MAX);
+  free (packed);
+}
+
+/* TODO: test alloc failure for slab, unknown fields */
+static void
+test_alloc_fail (void)
+{
+  int i = 0;
+  SETUP_TEST_ALLOC_BUFFER (packed, len);
+  do test_alloc_graceful_cleanup (packed, len, i++);
+  while (test_allocator_data.allocs_left < 0);
+  free (packed);
+}
+
 /* === simple testing framework === */
 
 typedef void (*TestFunc) (void);
@@ -1258,6 +1339,9 @@ static Test tests[] =
 
   { "test required default values", test_required_default_values },
   { "test optional default values", test_optional_default_values },
+
+  { "test free unpacked", test_alloc_free_all },
+  { "test alloc failure", test_alloc_fail },
 };
 #define n_tests (sizeof(tests)/sizeof(Test))