move buffer stuff out to buffer.c/h and be safer about allocs

This commit is contained in:
Mike Dirolf 2009-12-02 15:22:16 -05:00
parent 72a166622f
commit c893b9269f
3 changed files with 283 additions and 132 deletions

134
ext/cbson/buffer.c Normal file
View File

@ -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 <stdlib.h>
#include <string.h>
#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;
}

52
ext/cbson/buffer.h Normal file
View File

@ -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

View File

@ -36,14 +36,23 @@
#include "regex.h"
#endif
#include "version.h"
#include <assert.h>
#include <string.h>
#include <math.h>
#include <unistd.h>
#include <time.h>
#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;
}