From c893b9269f0830fce9af5ca5ad21a10c0f787c0a Mon Sep 17 00:00:00 2001 From: Mike Dirolf Date: Wed, 2 Dec 2009 15:22:16 -0500 Subject: [PATCH] move buffer stuff out to buffer.c/h and be safer about allocs --- ext/cbson/buffer.c | 134 ++++++++++++++++++++++++++ ext/cbson/buffer.h | 52 ++++++++++ ext/cbson/cbson.c | 229 +++++++++++++++++++-------------------------- 3 files changed, 283 insertions(+), 132 deletions(-) create mode 100644 ext/cbson/buffer.c create mode 100644 ext/cbson/buffer.h diff --git a/ext/cbson/buffer.c b/ext/cbson/buffer.c new file mode 100644 index 0000000..2ef381f --- /dev/null +++ b/ext/cbson/buffer.c @@ -0,0 +1,134 @@ +/* + * Copyright 2009 10gen, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "buffer.h" + +#define INITIAL_BUFFER_SIZE 256 + +struct buffer { + char* buffer; + int size; + int position; +}; + +/* Allocate and return a new buffer. + * Return NULL on allocation failure. */ +buffer_t buffer_new(void) { + buffer_t buffer; + buffer = (buffer_t)malloc(sizeof(struct buffer)); + if (buffer == NULL) { + return NULL; + } + + buffer->size = INITIAL_BUFFER_SIZE; + buffer->position = 0; + buffer->buffer = (char*)malloc(sizeof(char) * INITIAL_BUFFER_SIZE); + if (buffer->buffer == NULL) { + free(buffer); + return NULL; + } + + return buffer; +} + +/* Free the memory allocated for `buffer`. + * Return non-zero on failure. */ +int buffer_free(buffer_t buffer) { + if (buffer == NULL) { + return 1; + } + free(buffer->buffer); + free(buffer); + return 0; +} + +/* Grow `buffer` to at least `min_length`. + * Return non-zero on allocation failure. */ +static int buffer_grow(buffer_t buffer, int min_length) { + int size = buffer->size; + char* old_buffer = buffer->buffer; + if (size >= min_length) { + return 0; + } + while (size < min_length) { + size *= 2; + } + buffer->buffer = (char*)realloc(buffer->buffer, sizeof(char) * size); + if (buffer->buffer == NULL) { + free(old_buffer); + free(buffer); + return 1; + } + buffer->size = size; +} + +/* Assure that `buffer` has at least `size` free bytes (and grow if needed). + * Return non-zero on allocation failure. */ +static int buffer_assure_space(buffer_t buffer, int size) { + if (buffer->position + size <= buffer->size) { + return 0; + } + return buffer_grow(buffer, buffer->position + size); +} + +/* Save `size` bytes from the current position in `buffer` (and grow if needed). + * Return offset for writing, or -1 on allocation failure. */ +buffer_position buffer_save_space(buffer_t buffer, int size) { + int position = buffer->position; + if (buffer_assure_space(buffer, size) != 0) { + return -1; + } + buffer->position += size; + return position; +} + +/* Write `size` bytes from `data` to `buffer` (and grow if needed). + * Return non-zero on allocation failure. */ +int buffer_write(buffer_t buffer, const char* data, int size) { + if (buffer_assure_space(buffer, size) != 0) { + return 1; + } + + memcpy(buffer->buffer + buffer->position, data, size); + buffer->position += size; + return 0; +} + +/* Write `size` bytes from `data` to `buffer` at position `position`. + * Does not change the internal position of `buffer`. + * Return non-zero if buffer isn't large enough for write. */ +int buffer_write_at_position(buffer_t buffer, buffer_position position, + const char* data, int size) { + if (position + size > buffer->size) { + buffer_free(buffer); + return 1; + } + + memcpy(buffer->buffer + position, data, size); + return 0; +} + + +int buffer_get_position(buffer_t buffer) { + return buffer->position; +} + +char* buffer_get_buffer(buffer_t buffer) { + return buffer->buffer; +} diff --git a/ext/cbson/buffer.h b/ext/cbson/buffer.h new file mode 100644 index 0000000..6901494 --- /dev/null +++ b/ext/cbson/buffer.h @@ -0,0 +1,52 @@ +/* + * Copyright 2009 10gen, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef BUFFER_H +#define BUFFER_H + +/* A buffer */ +typedef struct buffer* buffer_t; +/* A position in the buffer */ +typedef int buffer_position; + +/* Allocate and return a new buffer. + * Return NULL on allocation failure. */ +buffer_t buffer_new(void); + +/* Free the memory allocated for `buffer`. + * Return non-zero on failure. */ +int buffer_free(buffer_t buffer); + +/* Save `size` bytes from the current position in `buffer` (and grow if needed). + * Return offset for writing, or -1 on allocation failure. */ +buffer_position buffer_save_space(buffer_t buffer, int size); + +/* Write `size` bytes from `data` to `buffer` (and grow if needed). + * Return non-zero on allocation failure. */ +int buffer_write(buffer_t buffer, const char* data, int size); + +/* Write `size` bytes from `data` to `buffer` at position `position`. + * Does not change the internal position of `buffer`. + * Return non-zero if buffer isn't large enough for write. */ +int buffer_write_at_position(buffer_t buffer, buffer_position position, const char* data, int size); + +/* Getters for the internals of a buffer_t. + * Should try to avoid using these as much as possible + * since they break the abstraction. */ +buffer_position buffer_get_position(buffer_t buffer); +char* buffer_get_buffer(buffer_t buffer); + +#endif diff --git a/ext/cbson/cbson.c b/ext/cbson/cbson.c index b4240b3..9c46450 100644 --- a/ext/cbson/cbson.c +++ b/ext/cbson/cbson.c @@ -36,14 +36,23 @@ #include "regex.h" #endif -#include "version.h" - -#include +#include #include #include #include -#define INITIAL_BUFFER_SIZE 256 +#include "version.h" +#include "buffer.h" + +#define SAFE_WRITE(buffer, data, size) \ + if (buffer_write((buffer), (data), (size)) != 0) \ + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c") + +#define SAFE_WRITE_AT_POS(buffer, position, data, size) \ + if (buffer_write_at_position((buffer), (position), (data), (size)) != 0) \ + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c") + + #define MAX_HOSTNAME_LENGTH 256 static VALUE Binary; @@ -93,12 +102,6 @@ static VALUE DigestMD5; #define RREGEXP_SRC_LEN(r) RREGEXP(r)->len #endif -typedef struct { - char* buffer; - int size; - int position; -} bson_buffer; - static char zero = 0; static char one = 1; @@ -106,78 +109,22 @@ static int cmp_char(const void* a, const void* b) { return *(char*)a - *(char*)b; } -static void write_doc(bson_buffer* buffer, VALUE hash, VALUE check_keys); +static void write_doc(buffer_t buffer, VALUE hash, VALUE check_keys); static int write_element(VALUE key, VALUE value, VALUE extra); static VALUE elements_to_hash(const char* buffer, int max); -static bson_buffer* buffer_new(void) { - bson_buffer* buffer; - buffer = ALLOC(bson_buffer); - assert(buffer); - - buffer->size = INITIAL_BUFFER_SIZE; - buffer->position = 0; - buffer->buffer = ALLOC_N(char, INITIAL_BUFFER_SIZE); - assert(buffer->buffer); - - return buffer; -} - -static void buffer_free(bson_buffer* buffer) { - assert(buffer); - assert(buffer->buffer); - - free(buffer->buffer); - free(buffer); -} - -static void buffer_resize(bson_buffer* buffer, int min_length) { - int size = buffer->size; - if (size >= min_length) { - return; - } - while (size < min_length) { - size *= 2; - } - buffer->buffer = REALLOC_N(buffer->buffer, char, size); - assert(buffer->buffer); - buffer->size = size; -} - -static void buffer_assure_space(bson_buffer* buffer, int size) { - if (buffer->position + size <= buffer->size) { - return; - } - buffer_resize(buffer, buffer->position + size); -} - -/* returns offset for writing */ -static int buffer_save_bytes(bson_buffer* buffer, int size) { - int position = buffer->position; - buffer_assure_space(buffer, size); - buffer->position += size; - return position; -} - -static void buffer_write_bytes(bson_buffer* buffer, const char* bytes, int size) { - buffer_assure_space(buffer, size); - - memcpy(buffer->buffer + buffer->position, bytes, size); - buffer->position += size; -} - -static VALUE pack_extra(bson_buffer* buffer, VALUE check_keys) { +static VALUE pack_extra(buffer_t buffer, VALUE check_keys) { return rb_ary_new3(2, LL2NUM((long long)buffer), check_keys); } -static void write_name_and_type(bson_buffer* buffer, VALUE name, char type) { - buffer_write_bytes(buffer, &type, 1); - buffer_write_bytes(buffer, RSTRING_PTR(name), RSTRING_LEN(name)); - buffer_write_bytes(buffer, &zero, 1); +static void write_name_and_type(buffer_t buffer, VALUE name, char type) { + SAFE_WRITE(buffer, &type, 1); + SAFE_WRITE(buffer, RSTRING_PTR(name), RSTRING_LEN(name)); + SAFE_WRITE(buffer, &zero, 1); } static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow_id) { - bson_buffer* buffer = (bson_buffer*)NUM2LL(rb_ary_entry(extra, 0)); + buffer_t buffer = (buffer_t)NUM2LL(rb_ary_entry(extra, 0)); VALUE check_keys = rb_ary_entry(extra, 1); if (TYPE(key) == T_SYMBOL) { @@ -218,32 +165,32 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow long long ll_value; write_name_and_type(buffer, key, 0x12); ll_value = NUM2LL(value); - buffer_write_bytes(buffer, (char*)&ll_value, 8); + SAFE_WRITE(buffer, (char*)&ll_value, 8); } else { int int_value; write_name_and_type(buffer, key, 0x10); int_value = NUM2LL(value); - buffer_write_bytes(buffer, (char*)&int_value, 4); + SAFE_WRITE(buffer, (char*)&int_value, 4); } break; } case T_TRUE: { write_name_and_type(buffer, key, 0x08); - buffer_write_bytes(buffer, &one, 1); + SAFE_WRITE(buffer, &one, 1); break; } case T_FALSE: { write_name_and_type(buffer, key, 0x08); - buffer_write_bytes(buffer, &zero, 1); + SAFE_WRITE(buffer, &zero, 1); break; } case T_FLOAT: { double d = NUM2DBL(value); write_name_and_type(buffer, key, 0x01); - buffer_write_bytes(buffer, (char*)&d, 8); + SAFE_WRITE(buffer, (char*)&d, 8); break; } case T_NIL: @@ -259,14 +206,18 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow } case T_ARRAY: { - int start_position, length_location, items, i, obj_length; + buffer_position length_location, start_position, obj_length; + int items, i; VALUE* values; write_name_and_type(buffer, key, 0x04); - start_position = buffer->position; + start_position = buffer_get_position(buffer); // save space for length - length_location = buffer_save_bytes(buffer, 4); + length_location = buffer_save_space(buffer, 4); + if (length_location == -1) { + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c"); + } items = RARRAY_LEN(value); values = RARRAY_PTR(value); @@ -280,37 +231,40 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow } // write null byte and fill in length - buffer_write_bytes(buffer, &zero, 1); - obj_length = buffer->position - start_position; - memcpy(buffer->buffer + length_location, &obj_length, 4); + SAFE_WRITE(buffer, &zero, 1); + obj_length = buffer_get_position(buffer) - start_position; + SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&obj_length, 4); break; } case T_STRING: { if (strcmp(rb_class2name(RBASIC(value)->klass), "Mongo::Code") == 0) { - int start_position, length_location, length, total_length; + buffer_position length_location, start_position, total_length; + int length; write_name_and_type(buffer, key, 0x0F); - start_position = buffer->position; - length_location = buffer_save_bytes(buffer, 4); + start_position = buffer_get_position(buffer); + length_location = buffer_save_space(buffer, 4); + if (length_location == -1) { + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c"); + } length = RSTRING_LEN(value) + 1; - buffer_write_bytes(buffer, (char*)&length, 4); - buffer_write_bytes(buffer, RSTRING_PTR(value), length - 1); - buffer_write_bytes(buffer, &zero, 1); + SAFE_WRITE(buffer, (char*)&length, 4); + SAFE_WRITE(buffer, RSTRING_PTR(value), length - 1); + SAFE_WRITE(buffer, &zero, 1); write_doc(buffer, rb_funcall(value, rb_intern("scope"), 0), Qfalse); - total_length = buffer->position - start_position; - memcpy(buffer->buffer + length_location, &total_length, 4); - + total_length = buffer_get_position(buffer) - start_position; + SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&total_length, 4); break; } else { int length = RSTRING_LEN(value) + 1; write_name_and_type(buffer, key, 0x02); - buffer_write_bytes(buffer, (char*)&length, 4); - buffer_write_bytes(buffer, RSTRING_PTR(value), length - 1); - buffer_write_bytes(buffer, &zero, 1); + SAFE_WRITE(buffer, (char*)&length, 4); + SAFE_WRITE(buffer, RSTRING_PTR(value), length - 1); + SAFE_WRITE(buffer, &zero, 1); break; } } @@ -319,8 +273,8 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow const char* str_value = rb_id2name(SYM2ID(value)); int length = strlen(str_value) + 1; write_name_and_type(buffer, key, 0x0E); - buffer_write_bytes(buffer, (char*)&length, 4); - buffer_write_bytes(buffer, str_value, length); + SAFE_WRITE(buffer, (char*)&length, 4); + SAFE_WRITE(buffer, str_value, length); break; } case T_OBJECT: @@ -336,14 +290,14 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow write_name_and_type(buffer, key, 0x05); if (subtype == 2) { const int other_length = length + 4; - buffer_write_bytes(buffer, (const char*)&other_length, 4); - buffer_write_bytes(buffer, &subtype, 1); + SAFE_WRITE(buffer, (const char*)&other_length, 4); + SAFE_WRITE(buffer, &subtype, 1); } - buffer_write_bytes(buffer, (const char*)&length, 4); + SAFE_WRITE(buffer, (const char*)&length, 4); if (subtype != 2) { - buffer_write_bytes(buffer, &subtype, 1); + SAFE_WRITE(buffer, &subtype, 1); } - buffer_write_bytes(buffer, RSTRING_PTR(string_data), length); + SAFE_WRITE(buffer, RSTRING_PTR(string_data), length); break; } if (strcmp(cls, "Mongo::ObjectID") == 0) { @@ -352,19 +306,22 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow write_name_and_type(buffer, key, 0x07); for (i = 0; i < 12; i++) { char byte = (char)FIX2INT(RARRAY_PTR(as_array)[i]); - buffer_write_bytes(buffer, &byte, 1); + SAFE_WRITE(buffer, &byte, 1); } break; } if (strcmp(cls, "Mongo::DBRef") == 0) { - int start_position, length_location, obj_length; + buffer_position length_location, start_position, obj_length; VALUE ns, oid; write_name_and_type(buffer, key, 0x03); - start_position = buffer->position; + start_position = buffer_get_position(buffer); // save space for length - length_location = buffer_save_bytes(buffer, 4); + length_location = buffer_save_space(buffer, 4); + if (length_location == -1) { + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c"); + } ns = rb_funcall(value, rb_intern("namespace"), 0); write_element(rb_str_new2("$ref"), ns, pack_extra(buffer, Qfalse)); @@ -372,9 +329,9 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow write_element(rb_str_new2("$id"), oid, pack_extra(buffer, Qfalse)); // write null byte and fill in length - buffer_write_bytes(buffer, &zero, 1); - obj_length = buffer->position - start_position; - memcpy(buffer->buffer + length_location, &obj_length, 4); + SAFE_WRITE(buffer, &zero, 1); + obj_length = buffer_get_position(buffer) - start_position; + SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&obj_length, 4); break; } } @@ -386,7 +343,7 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow double t = NUM2DBL(rb_funcall(value, rb_intern("to_f"), 0)); long long time_since_epoch = (long long)round(t * 1000); write_name_and_type(buffer, key, 0x09); - buffer_write_bytes(buffer, (const char*)&time_since_epoch, 8); + SAFE_WRITE(buffer, (const char*)&time_since_epoch, 8); break; } } @@ -399,30 +356,30 @@ static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow write_name_and_type(buffer, key, 0x0B); - buffer_write_bytes(buffer, pattern, length); - buffer_write_bytes(buffer, &zero, 1); + SAFE_WRITE(buffer, pattern, length); + SAFE_WRITE(buffer, &zero, 1); if (flags & IGNORECASE) { char ignorecase = 'i'; - buffer_write_bytes(buffer, &ignorecase, 1); + SAFE_WRITE(buffer, &ignorecase, 1); } if (flags & MULTILINE) { char multiline = 'm'; - buffer_write_bytes(buffer, &multiline, 1); + SAFE_WRITE(buffer, &multiline, 1); } if (flags & EXTENDED) { char extended = 'x'; - buffer_write_bytes(buffer, &extended, 1); + SAFE_WRITE(buffer, &extended, 1); } has_extra = rb_funcall(value, rb_intern("respond_to?"), 1, rb_str_new2("extra_options_str")); if (TYPE(has_extra) == T_TRUE) { VALUE extra = rb_funcall(value, rb_intern("extra_options_str"), 0); - int old_position = buffer->position; - buffer_write_bytes(buffer, RSTRING_PTR(extra), RSTRING_LEN(extra)); - qsort(buffer->buffer + old_position, RSTRING_LEN(extra), sizeof(char), cmp_char); + buffer_position old_position = buffer_get_position(buffer); + SAFE_WRITE(buffer, RSTRING_PTR(extra), RSTRING_LEN(extra)); + qsort(buffer_get_buffer(buffer) + old_position, RSTRING_LEN(extra), sizeof(char), cmp_char); } - buffer_write_bytes(buffer, &zero, 1); + SAFE_WRITE(buffer, &zero, 1); break; } @@ -439,13 +396,17 @@ static int write_element(VALUE key, VALUE value, VALUE extra) { return write_element_allow_id(key, value, extra, 0); } -static void write_doc(bson_buffer* buffer, VALUE hash, VALUE check_keys) { - int start_position = buffer->position; - int length_location = buffer_save_bytes(buffer, 4); - int length; - +static void write_doc(buffer_t buffer, VALUE hash, VALUE check_keys) { + buffer_position start_position = buffer_get_position(buffer); + buffer_position length_location = buffer_save_space(buffer, 4); + buffer_position length; VALUE id_str = rb_str_new2("_id"); VALUE id_sym = ID2SYM(rb_intern("_id")); + + if (length_location == -1) { + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c"); + } + if (rb_funcall(hash, rb_intern("has_key?"), 1, id_str) == Qtrue) { VALUE id = rb_hash_aref(hash, id_str); write_element_allow_id(id_str, id, pack_extra(buffer, check_keys), 1); @@ -469,20 +430,24 @@ static void write_doc(bson_buffer* buffer, VALUE hash, VALUE check_keys) { } // write null byte and fill in length - buffer_write_bytes(buffer, &zero, 1); - length = buffer->position - start_position; - memcpy(buffer->buffer + length_location, &length, 4); + SAFE_WRITE(buffer, &zero, 1); + length = buffer_get_position(buffer) - start_position; + SAFE_WRITE_AT_POS(buffer, length_location, &length, 4); } static VALUE method_serialize(VALUE self, VALUE doc, VALUE check_keys) { VALUE result; - bson_buffer* buffer = buffer_new(); - assert(buffer); + buffer_t buffer = buffer_new(); + if (buffer == NULL) { + rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c"); + } write_doc(buffer, doc, check_keys); - result = rb_str_new(buffer->buffer, buffer->position); - buffer_free(buffer); + result = rb_str_new(buffer_get_buffer(buffer), buffer_get_position(buffer)); + if (buffer_free(buffer) != 0) { + rb_raise(rb_eRuntimeError, "failed to free buffer"); + } return result; }