diff --git a/src/google/protobuf-c/protobuf-c.c b/src/google/protobuf-c/protobuf-c.c index d3e4ae7..4771a13 100644 --- a/src/google/protobuf-c/protobuf-c.c +++ b/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 */ diff --git a/src/test/test-full.proto b/src/test/test-full.proto index 7915b0f..c3ff845 100644 --- a/src/test/test-full.proto +++ b/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; +} diff --git a/src/test/test-generated-code2.c b/src/test/test-generated-code2.c index 5587a85..30d98e2 100644 --- a/src/test/test-generated-code2.c +++ b/src/test/test-generated-code2.c @@ -1,3 +1,4 @@ +#include #include #include #include @@ -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))