/* add a lastError message on the end of the buffer. * returns 0 on failure */ static int add_last_error(buffer_t buffer, int request_id, PyObject* args) { int message_start; int document_start; int message_length; int document_length; PyObject* key; PyObject* value; Py_ssize_t pos = 0; PyObject* one; message_start = buffer_save_space(buffer, 4); if (message_start == -1) { PyErr_NoMemory(); return 0; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00" /* responseTo */ "\xd4\x07\x00\x00" /* opcode */ "\x00\x00\x00\x00" /* options */ "admin.$cmd\x00" /* collection name */ "\x00\x00\x00\x00" /* skip */ "\xFF\xFF\xFF\xFF", /* limit (-1) */ 31)) { return 0; } /* save space for length */ document_start = buffer_save_space(buffer, 4); if (document_start == -1) { PyErr_NoMemory(); return 0; } /* getlasterror: 1 */ one = PyLong_FromLong(1); if (!write_pair(buffer, "getlasterror", 12, one, 0, 4, 1)) { Py_DECREF(one); return 0; } Py_DECREF(one); /* getlasterror options */ while (PyDict_Next(args, &pos, &key, &value)) { if (!decode_and_write_pair(buffer, key, value, 0, 4, 0)) { return 0; } } /* EOD */ if (!buffer_write_bytes(buffer, "\x00", 1)) { return 0; } message_length = buffer_get_position(buffer) - message_start; document_length = buffer_get_position(buffer) - document_start; memcpy(buffer_get_buffer(buffer) + message_start, &message_length, 4); memcpy(buffer_get_buffer(buffer) + document_start, &document_length, 4); return 1; }
static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) { /* NOTE just using a random number as the request_id */ int request_id = rand(); char* collection_name = NULL; int collection_name_length; int num_to_return; long long cursor_id; buffer_t buffer; int length_location, message_length; PyObject* result; if (!PyArg_ParseTuple(args, "et#iL", "utf-8", &collection_name, &collection_name_length, &num_to_return, &cursor_id)) { return NULL; } buffer = buffer_new(); if (!buffer) { PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } // save space for message length length_location = buffer_save_space(buffer, 4); if (length_location == -1) { PyMem_Free(collection_name); PyErr_NoMemory(); return NULL; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00" "\xd5\x07\x00\x00" "\x00\x00\x00\x00", 12) || !buffer_write_bytes(buffer, collection_name, collection_name_length + 1) || !buffer_write_bytes(buffer, (const char*)&num_to_return, 4) || !buffer_write_bytes(buffer, (const char*)&cursor_id, 8)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } PyMem_Free(collection_name); message_length = buffer_get_position(buffer) - length_location; memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4); /* objectify buffer */ result = Py_BuildValue("i" BYTES_FORMAT_STRING, request_id, buffer_get_buffer(buffer), buffer_get_position(buffer)); buffer_free(buffer); return result; }
static buffer_t _command_buffer_new(char* ns, int ns_len) { buffer_t buffer; if (!(buffer = buffer_new())) { PyErr_NoMemory(); return NULL; } /* Save space for message length and request id */ if ((buffer_save_space(buffer, 8)) == -1) { PyErr_NoMemory(); buffer_free(buffer); return NULL; } if (!buffer_write_bytes(buffer, "\x00\x00\x00\x00" /* responseTo */ "\xd4\x07\x00\x00" /* opcode */ "\x00\x00\x00\x00", /* flags */ 12) || !buffer_write_bytes(buffer, ns, ns_len + 1) || /* namespace */ !buffer_write_bytes(buffer, "\x00\x00\x00\x00" /* skip */ "\xFF\xFF\xFF\xFF", /* limit (-1) */ 8)) { buffer_free(buffer); return NULL; } return buffer; }
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); } else if (rb_funcall(hash, rb_intern("has_key?"), 1, id_sym) == Qtrue) { VALUE id = rb_hash_aref(hash, id_sym); write_element_allow_id(id_sym, id, pack_extra(buffer, check_keys), 1); } // we have to check for an OrderedHash and handle that specially if (strcmp(rb_class2name(RBASIC(hash)->klass), "OrderedHash") == 0) { VALUE keys = rb_funcall(hash, rb_intern("keys"), 0); int i; for(i = 0; i < RARRAY_LEN(keys); i++) { VALUE key = RARRAY_PTR(keys)[i]; VALUE value = rb_hash_aref(hash, key); write_element(key, value, pack_extra(buffer, check_keys)); } } else { rb_hash_foreach(hash, write_element, pack_extra(buffer, check_keys)); } // write null byte and fill in length SAFE_WRITE(buffer, &zero, 1); length = buffer_get_position(buffer) - start_position; // make sure that length doesn't exceed 4MB if (length > 4 * 1024 * 1024) { buffer_free(buffer); rb_raise(InvalidDocument, "Document too large: BSON documents are limited to 4MB."); return; } SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&length, 4); }
static int init_insert_buffer(buffer_t buffer, int request_id, int options, const char* coll_name, int coll_name_len) { /* Save space for message length */ int length_location = buffer_save_space(buffer, 4); if (length_location == -1) { PyErr_NoMemory(); return length_location; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00" "\xd2\x07\x00\x00", 8) || !buffer_write_bytes(buffer, (const char*)&options, 4) || !buffer_write_bytes(buffer, coll_name, coll_name_len + 1)) { return -1; } return length_location; }
static void write_doc(buffer_t buffer, VALUE hash, VALUE check_keys, VALUE move_id) { buffer_position start_position = buffer_get_position(buffer); buffer_position length_location = buffer_save_space(buffer, 4); buffer_position length; int allow_id; int (*write_function)(VALUE, VALUE, VALUE) = NULL; 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"); } // write '_id' first if move_id is true. then don't allow an id to be written. if(move_id == Qtrue) { allow_id = 0; if (rb_funcall(hash, rb_intern("has_key?"), 1, id_str) == Qtrue) { VALUE id = rb_hash_aref(hash, id_str); write_element_with_id(id_str, id, pack_extra(buffer, check_keys)); } else if (rb_funcall(hash, rb_intern("has_key?"), 1, id_sym) == Qtrue) { VALUE id = rb_hash_aref(hash, id_sym); write_element_with_id(id_sym, id, pack_extra(buffer, check_keys)); } } else { allow_id = 1; // Ensure that hash doesn't contain both '_id' and :_id if ((rb_obj_classname(hash), "Hash") == 0) { if ((rb_funcall(hash, rb_intern("has_key?"), 1, id_str) == Qtrue) && (rb_funcall(hash, rb_intern("has_key?"), 1, id_sym) == Qtrue)) { VALUE oid_sym = rb_hash_delete(hash, id_sym); rb_funcall(hash, rb_intern("[]="), 2, id_str, oid_sym); } } } if(allow_id == 1) { write_function = write_element_with_id; } else { write_function = write_element_without_id; } // we have to check for an OrderedHash and handle that specially if (strcmp(rb_obj_classname(hash), "BSON::OrderedHash") == 0) { VALUE keys = rb_funcall(hash, rb_intern("keys"), 0); int i; for(i = 0; i < RARRAY_LEN(keys); i++) { VALUE key = rb_ary_entry(keys, i); VALUE value = rb_hash_aref(hash, key); write_function(key, value, pack_extra(buffer, check_keys)); } } else if (rb_obj_is_kind_of(hash, RB_HASH) == Qtrue) { rb_hash_foreach(hash, write_function, pack_extra(buffer, check_keys)); } else { bson_buffer_free(buffer); rb_raise(InvalidDocument, "BSON.serialize takes a Hash but got a %s", rb_obj_classname(hash)); } // write null byte and fill in length SAFE_WRITE(buffer, &zero, 1); length = buffer_get_position(buffer) - start_position; // make sure that length doesn't exceed 4MB if (length > max_bson_size) { bson_buffer_free(buffer); rb_raise(InvalidDocument, "Document too large: BSON documents are limited to %d bytes.", max_bson_size); return; } SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&length, 4); }
static int write_element(VALUE key, VALUE value, VALUE extra, int allow_id) { buffer_t buffer = (buffer_t)NUM2LL(rb_ary_entry(extra, 0)); VALUE check_keys = rb_ary_entry(extra, 1); if (TYPE(key) == T_SYMBOL) { // TODO better way to do this... ? key = rb_str_new2(rb_id2name(SYM2ID(key))); } if (TYPE(key) != T_STRING) { bson_buffer_free(buffer); rb_raise(rb_eTypeError, "keys must be strings or symbols"); } if (allow_id == 0 && strcmp("_id", RSTRING_PTR(key)) == 0) { return ST_CONTINUE; } if (check_keys == Qtrue) { int i; if (RSTRING_LEN(key) > 0 && RSTRING_PTR(key)[0] == '$') { bson_buffer_free(buffer); rb_raise(InvalidKeyName, "%s - key must not start with '$'", RSTRING_PTR(key)); } for (i = 0; i < RSTRING_LEN(key); i++) { if (RSTRING_PTR(key)[i] == '.') { bson_buffer_free(buffer); rb_raise(InvalidKeyName, "%s - key must not contain '.'", RSTRING_PTR(key)); } } } switch(TYPE(value)) { case T_BIGNUM: { if (rb_funcall(value, gt_operator, 1, LL2NUM(9223372036854775807LL)) == Qtrue || rb_funcall(value, lt_operator, 1, LL2NUM(-9223372036854775808ULL)) == Qtrue) { bson_buffer_free(buffer); rb_raise(rb_eRangeError, "MongoDB can only handle 8-byte ints"); } } // NOTE: falls through to T_FIXNUM code case T_FIXNUM: { long long ll_value; ll_value = NUM2LL(value); if (ll_value > 2147483647LL || ll_value < -2147483648LL) { write_name_and_type(buffer, key, 0x12); SAFE_WRITE(buffer, (char*)&ll_value, 8); } else { int int_value; write_name_and_type(buffer, key, 0x10); int_value = (int)ll_value; SAFE_WRITE(buffer, (char*)&int_value, 4); } break; } case T_TRUE: { write_name_and_type(buffer, key, 0x08); SAFE_WRITE(buffer, &one, 1); break; } case T_FALSE: { write_name_and_type(buffer, key, 0x08); SAFE_WRITE(buffer, &zero, 1); break; } case T_FLOAT: { double d = NUM2DBL(value); write_name_and_type(buffer, key, 0x01); SAFE_WRITE(buffer, (char*)&d, 8); break; } case T_NIL: { write_name_and_type(buffer, key, 0x0A); break; } case T_HASH: { write_name_and_type(buffer, key, 0x03); write_doc(buffer, value, check_keys, Qfalse); break; } case T_ARRAY: { buffer_position length_location, start_position, obj_length; int items, i; VALUE* values; write_name_and_type(buffer, key, 0x04); start_position = buffer_get_position(buffer); // save space for length length_location = buffer_save_space(buffer, 4); if (length_location == -1) { rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c"); } items = RARRAY_LENINT(value); for(i = 0; i < items; i++) { char* name; VALUE key; INT2STRING(&name, i); key = rb_str_new2(name); write_element_with_id(key, rb_ary_entry(value, i), pack_extra(buffer, check_keys)); FREE_INTSTRING(name); } // write null byte and fill in length 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: { int length; write_name_and_type(buffer, key, 0x02); length = RSTRING_LENINT(value) + 1; SAFE_WRITE(buffer, (char*)&length, 4); write_utf8(buffer, value, 0); SAFE_WRITE(buffer, &zero, 1); break; } case T_SYMBOL: { const char* str_value = rb_id2name(SYM2ID(value)); int length = (int)strlen(str_value) + 1; write_name_and_type(buffer, key, 0x0E); SAFE_WRITE(buffer, (char*)&length, 4); SAFE_WRITE(buffer, str_value, length); break; } case T_OBJECT: { // TODO there has to be a better way to do these checks... const char* cls = rb_obj_classname(value); if (strcmp(cls, "BSON::Binary") == 0 || strcmp(cls, "ByteBuffer") == 0) { const char subtype = strcmp(cls, "ByteBuffer") ? (const char)FIX2INT(rb_funcall(value, rb_intern("subtype"), 0)) : 2; VALUE string_data = rb_funcall(value, rb_intern("to_s"), 0); int length = RSTRING_LENINT(string_data); write_name_and_type(buffer, key, 0x05); if (subtype == 2) { const int other_length = length + 4; SAFE_WRITE(buffer, (const char*)&other_length, 4); SAFE_WRITE(buffer, &subtype, 1); } SAFE_WRITE(buffer, (const char*)&length, 4); if (subtype != 2) { SAFE_WRITE(buffer, &subtype, 1); } SAFE_WRITE(buffer, RSTRING_PTR(string_data), length); break; } if (strcmp(cls, "BSON::ObjectId") == 0) { VALUE as_array = rb_funcall(value, rb_intern("to_a"), 0); int i; write_name_and_type(buffer, key, 0x07); for (i = 0; i < 12; i++) { char byte = (char)FIX2INT(rb_ary_entry(as_array, i)); SAFE_WRITE(buffer, &byte, 1); } break; } if (strcmp(cls, "BSON::DBRef") == 0) { buffer_position length_location, start_position, obj_length; VALUE ns, oid; write_name_and_type(buffer, key, 0x03); start_position = buffer_get_position(buffer); // save space for length 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_with_id(rb_str_new2("$ref"), ns, pack_extra(buffer, Qfalse)); oid = rb_funcall(value, rb_intern("object_id"), 0); write_element_with_id(rb_str_new2("$id"), oid, pack_extra(buffer, Qfalse)); // write null byte and fill in length 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; } if (strcmp(cls, "BSON::Code") == 0) { buffer_position length_location, start_position, total_length; int length; VALUE code_str; write_name_and_type(buffer, key, 0x0F); 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"); } code_str = rb_funcall(value, rb_intern("code"), 0); length = RSTRING_LENINT(code_str) + 1; SAFE_WRITE(buffer, (char*)&length, 4); SAFE_WRITE(buffer, RSTRING_PTR(code_str), length - 1); SAFE_WRITE(buffer, &zero, 1); write_doc(buffer, rb_funcall(value, rb_intern("scope"), 0), Qfalse, Qfalse); total_length = buffer_get_position(buffer) - start_position; SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&total_length, 4); break; } if (strcmp(cls, "BSON::MaxKey") == 0) { write_name_and_type(buffer, key, 0x7f); break; } if (strcmp(cls, "BSON::MinKey") == 0) { write_name_and_type(buffer, key, 0xff); break; } if (strcmp(cls, "BSON::Timestamp") == 0) { write_name_and_type(buffer, key, 0x11); int seconds = FIX2INT( rb_funcall(value, rb_intern("seconds"), 0)); int increment = FIX2INT( rb_funcall(value, rb_intern("increment"), 0)); SAFE_WRITE(buffer, (const char*)&increment, 4); SAFE_WRITE(buffer, (const char*)&seconds, 4); break; } if (strcmp(cls, "DateTime") == 0 || strcmp(cls, "Date") == 0 || strcmp(cls, "ActiveSupport::TimeWithZone") == 0) { bson_buffer_free(buffer); rb_raise(InvalidDocument, "%s is not currently supported; use a UTC Time instance instead.", cls); break; } if(strcmp(cls, "Complex") == 0 || strcmp(cls, "Rational") == 0 || strcmp(cls, "BigDecimal") == 0) { bson_buffer_free(buffer); rb_raise(InvalidDocument, "Cannot serialize the Numeric type %s as BSON; only Bignum, Fixnum, and Float are supported.", cls); break; } bson_buffer_free(buffer); rb_raise(InvalidDocument, "Cannot serialize an object of class %s into BSON.", cls); break; } case T_DATA: { const char* cls = rb_obj_classname(value); if (strcmp(cls, "Time") == 0) { 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); SAFE_WRITE(buffer, (const char*)&time_since_epoch, 8); break; } if(strcmp(cls, "BigDecimal") == 0) { bson_buffer_free(buffer); rb_raise(InvalidDocument, "Cannot serialize the Numeric type %s as BSON; only Bignum, Fixnum, and Float are supported.", cls); break; } bson_buffer_free(buffer); rb_raise(InvalidDocument, "Cannot serialize an object of class %s into BSON.", cls); break; } case T_REGEXP: { VALUE pattern = RREGEXP_SRC(value); long flags = RREGEXP_OPTIONS(value); VALUE has_extra; write_name_and_type(buffer, key, 0x0B); write_utf8(buffer, pattern, 1); SAFE_WRITE(buffer, &zero, 1); if (flags & IGNORECASE) { char ignorecase = 'i'; SAFE_WRITE(buffer, &ignorecase, 1); } if (flags & MULTILINE) { char multiline = 'm'; char dotall = 's'; SAFE_WRITE(buffer, &multiline, 1); SAFE_WRITE(buffer, &dotall, 1); } if (flags & EXTENDED) { char extended = 'x'; 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); buffer_position old_position = buffer_get_position(buffer); SAFE_WRITE(buffer, RSTRING_PTR(extra), RSTRING_LENINT(extra)); qsort(buffer_get_buffer(buffer) + old_position, RSTRING_LEN(extra), sizeof(char), cmp_char); } SAFE_WRITE(buffer, &zero, 1); break; } default: { const char* cls = rb_obj_classname(value); bson_buffer_free(buffer); rb_raise(InvalidDocument, "Cannot serialize an object of class %s (type %d) into BSON.", cls, TYPE(value)); break; } } return ST_CONTINUE; }
/* add a lastError message on the end of the buffer. * returns 0 on failure */ static int add_last_error(PyObject* self, buffer_t buffer, int request_id, char* ns, int nslen, PyObject* args) { struct module_state *state = GETSTATE(self); int message_start; int document_start; int message_length; int document_length; PyObject* key; PyObject* value; Py_ssize_t pos = 0; PyObject* one; char *p = strchr(ns, '.'); /* Length of the database portion of ns. */ nslen = p ? (int)(p - ns) : nslen; message_start = buffer_save_space(buffer, 4); if (message_start == -1) { PyErr_NoMemory(); return 0; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00" /* responseTo */ "\xd4\x07\x00\x00" /* opcode */ "\x00\x00\x00\x00", /* options */ 12) || !buffer_write_bytes(buffer, ns, nslen) || /* database */ !buffer_write_bytes(buffer, ".$cmd\x00" /* collection name */ "\x00\x00\x00\x00" /* skip */ "\xFF\xFF\xFF\xFF", /* limit (-1) */ 14)) { return 0; } /* save space for length */ document_start = buffer_save_space(buffer, 4); if (document_start == -1) { PyErr_NoMemory(); return 0; } /* getlasterror: 1 */ if (!(one = PyLong_FromLong(1))) return 0; if (!write_pair(state->_cbson, buffer, "getlasterror", 12, one, 0, 4, 1)) { Py_DECREF(one); return 0; } Py_DECREF(one); /* getlasterror options */ while (PyDict_Next(args, &pos, &key, &value)) { if (!decode_and_write_pair(state->_cbson, buffer, key, value, 0, 4, 0)) { return 0; } } /* EOD */ if (!buffer_write_bytes(buffer, "\x00", 1)) { return 0; } message_length = buffer_get_position(buffer) - message_start; document_length = buffer_get_position(buffer) - document_start; memcpy(buffer_get_buffer(buffer) + message_start, &message_length, 4); memcpy(buffer_get_buffer(buffer) + document_start, &document_length, 4); return 1; }
static PyObject* _cbson_query_message(PyObject* self, PyObject* args) { /* NOTE just using a random number as the request_id */ struct module_state *state = GETSTATE(self); int request_id = rand(); unsigned int options; char* collection_name = NULL; int collection_name_length; int begin, cur_size, max_size = 0; int num_to_skip; int num_to_return; PyObject* query; PyObject* field_selector = Py_None; unsigned char uuid_subtype = 3; buffer_t buffer; int length_location, message_length; PyObject* result; if (!PyArg_ParseTuple(args, "Iet#iiO|Ob", &options, "utf-8", &collection_name, &collection_name_length, &num_to_skip, &num_to_return, &query, &field_selector, &uuid_subtype)) { return NULL; } buffer = buffer_new(); if (!buffer) { PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } // save space for message length length_location = buffer_save_space(buffer, 4); if (length_location == -1) { PyMem_Free(collection_name); PyErr_NoMemory(); return NULL; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) || !buffer_write_bytes(buffer, (const char*)&options, 4) || !buffer_write_bytes(buffer, collection_name, collection_name_length + 1) || !buffer_write_bytes(buffer, (const char*)&num_to_skip, 4) || !buffer_write_bytes(buffer, (const char*)&num_to_return, 4)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } begin = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, query, 0, uuid_subtype, 1)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } max_size = buffer_get_position(buffer) - begin; if (field_selector != Py_None) { begin = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, field_selector, 0, uuid_subtype, 1)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } cur_size = buffer_get_position(buffer) - begin; max_size = (cur_size > max_size) ? cur_size : max_size; } PyMem_Free(collection_name); message_length = buffer_get_position(buffer) - length_location; memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4); /* objectify buffer */ result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id, buffer_get_buffer(buffer), buffer_get_position(buffer), max_size); buffer_free(buffer); return result; }
static PyObject* _cbson_update_message(PyObject* self, PyObject* args) { /* NOTE just using a random number as the request_id */ struct module_state *state = GETSTATE(self); int request_id = rand(); char* collection_name = NULL; int collection_name_length; int before, cur_size, max_size = 0; PyObject* doc; PyObject* spec; unsigned char multi; unsigned char upsert; unsigned char safe; unsigned char check_keys; unsigned char uuid_subtype; PyObject* last_error_args; int options; buffer_t buffer; int length_location, message_length; PyObject* result; if (!PyArg_ParseTuple(args, "et#bbOObObb", "utf-8", &collection_name, &collection_name_length, &upsert, &multi, &spec, &doc, &safe, &last_error_args, &check_keys, &uuid_subtype)) { return NULL; } options = 0; if (upsert) { options += 1; } if (multi) { options += 2; } buffer = buffer_new(); if (!buffer) { PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } // save space for message length length_location = buffer_save_space(buffer, 4); if (length_location == -1) { PyMem_Free(collection_name); PyErr_NoMemory(); return NULL; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00" "\xd1\x07\x00\x00" "\x00\x00\x00\x00", 12) || !buffer_write_bytes(buffer, collection_name, collection_name_length + 1) || !buffer_write_bytes(buffer, (const char*)&options, 4)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } before = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, spec, 0, uuid_subtype, 1)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } max_size = buffer_get_position(buffer) - before; before = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } cur_size = buffer_get_position(buffer) - before; max_size = (cur_size > max_size) ? cur_size : max_size; message_length = buffer_get_position(buffer) - length_location; memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4); if (safe) { if (!add_last_error(self, buffer, request_id, collection_name, collection_name_length, last_error_args)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } } PyMem_Free(collection_name); /* objectify buffer */ result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id, buffer_get_buffer(buffer), buffer_get_position(buffer), max_size); buffer_free(buffer); return result; }
static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) { /* NOTE just using a random number as the request_id */ struct module_state *state = GETSTATE(self); int request_id = rand(); char* collection_name = NULL; int collection_name_length; PyObject* docs; PyObject* doc; PyObject* iterator; int before, cur_size, max_size = 0; int options = 0; unsigned char check_keys; unsigned char safe; unsigned char continue_on_error; unsigned char uuid_subtype; PyObject* last_error_args; buffer_t buffer; int length_location, message_length; PyObject* result; if (!PyArg_ParseTuple(args, "et#ObbObb", "utf-8", &collection_name, &collection_name_length, &docs, &check_keys, &safe, &last_error_args, &continue_on_error, &uuid_subtype)) { return NULL; } if (continue_on_error) { options += 1; } buffer = buffer_new(); if (!buffer) { PyErr_NoMemory(); PyMem_Free(collection_name); return NULL; } // save space for message length length_location = buffer_save_space(buffer, 4); if (length_location == -1) { PyMem_Free(collection_name); PyErr_NoMemory(); return NULL; } if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) || !buffer_write_bytes(buffer, "\x00\x00\x00\x00" "\xd2\x07\x00\x00", 8) || !buffer_write_bytes(buffer, (const char*)&options, 4) || !buffer_write_bytes(buffer, collection_name, collection_name_length + 1)) { PyMem_Free(collection_name); buffer_free(buffer); return NULL; } iterator = PyObject_GetIter(docs); if (iterator == NULL) { PyObject* InvalidOperation = _error("InvalidOperation"); if (InvalidOperation) { PyErr_SetString(InvalidOperation, "input is not iterable"); Py_DECREF(InvalidOperation); } buffer_free(buffer); PyMem_Free(collection_name); return NULL; } while ((doc = PyIter_Next(iterator)) != NULL) { before = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) { Py_DECREF(doc); Py_DECREF(iterator); buffer_free(buffer); PyMem_Free(collection_name); return NULL; } Py_DECREF(doc); cur_size = buffer_get_position(buffer) - before; max_size = (cur_size > max_size) ? cur_size : max_size; } Py_DECREF(iterator); if (PyErr_Occurred()) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } if (!max_size) { PyObject* InvalidOperation = _error("InvalidOperation"); if (InvalidOperation) { PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert"); Py_DECREF(InvalidOperation); } buffer_free(buffer); PyMem_Free(collection_name); return NULL; } message_length = buffer_get_position(buffer) - length_location; memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4); if (safe) { if (!add_last_error(self, buffer, request_id, collection_name, collection_name_length, last_error_args)) { buffer_free(buffer); PyMem_Free(collection_name); return NULL; } } PyMem_Free(collection_name); /* objectify buffer */ result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id, buffer_get_buffer(buffer), buffer_get_position(buffer), max_size); buffer_free(buffer); return result; }
static PyObject* _cbson_do_batched_write_command(PyObject* self, PyObject* args) { struct module_state *state = GETSTATE(self); long max_bson_size; long max_cmd_size; long max_write_batch_size; long idx_offset = 0; int idx = 0; int cmd_len_loc; int lst_len_loc; int ns_len; int ordered; char *ns = NULL; PyObject* max_bson_size_obj; PyObject* max_write_batch_size_obj; PyObject* command; PyObject* doc; PyObject* docs; PyObject* ctx; PyObject* iterator; PyObject* result; PyObject* results; PyObject* to_publish = NULL; unsigned char op; unsigned char check_keys; codec_options_t options; unsigned char empty = 1; unsigned char errors = 0; buffer_t buffer; if (!PyArg_ParseTuple(args, "et#bOObO&O", "utf-8", &ns, &ns_len, &op, &command, &docs, &check_keys, convert_codec_options, &options, &ctx)) { return NULL; } max_bson_size_obj = PyObject_GetAttrString(ctx, "max_bson_size"); #if PY_MAJOR_VERSION >= 3 max_bson_size = PyLong_AsLong(max_bson_size_obj); #else max_bson_size = PyInt_AsLong(max_bson_size_obj); #endif Py_XDECREF(max_bson_size_obj); if (max_bson_size == -1) { destroy_codec_options(&options); PyMem_Free(ns); return NULL; } /* * Max BSON object size + 16k - 2 bytes for ending NUL bytes * XXX: This should come from the server - SERVER-10643 */ max_cmd_size = max_bson_size + 16382; max_write_batch_size_obj = PyObject_GetAttrString(ctx, "max_write_batch_size"); #if PY_MAJOR_VERSION >= 3 max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj); #else max_write_batch_size = PyInt_AsLong(max_write_batch_size_obj); #endif Py_XDECREF(max_write_batch_size_obj); if (max_write_batch_size == -1) { destroy_codec_options(&options); PyMem_Free(ns); return NULL; } /* Default to True */ ordered = !((PyDict_GetItemString(command, "ordered")) == Py_False); if (!(results = PyList_New(0))) { destroy_codec_options(&options); PyMem_Free(ns); return NULL; } if (!(to_publish = PyList_New(0))) { destroy_codec_options(&options); PyMem_Free(ns); Py_DECREF(results); return NULL; } if (!(buffer = _command_buffer_new(ns, ns_len))) { destroy_codec_options(&options); PyMem_Free(ns); Py_DECREF(results); Py_DECREF(to_publish); return NULL; } PyMem_Free(ns); /* Position of command document length */ cmd_len_loc = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, command, 0, &options, 0)) { goto cmdfail; } /* Write type byte for array */ *(buffer_get_buffer(buffer) + (buffer_get_position(buffer) - 1)) = 0x4; switch (op) { case _INSERT: { if (!buffer_write_bytes(buffer, "documents\x00", 10)) goto cmdfail; break; } case _UPDATE: { /* MongoDB does key validation for update. */ check_keys = 0; if (!buffer_write_bytes(buffer, "updates\x00", 8)) goto cmdfail; break; } case _DELETE: { /* Never check keys in a delete command. */ check_keys = 0; if (!buffer_write_bytes(buffer, "deletes\x00", 8)) goto cmdfail; break; } default: { PyObject* InvalidOperation = _error("InvalidOperation"); if (InvalidOperation) { PyErr_SetString(InvalidOperation, "Unknown command"); Py_DECREF(InvalidOperation); } goto cmdfail; } } /* Save space for list document */ lst_len_loc = buffer_save_space(buffer, 4); if (lst_len_loc == -1) { PyErr_NoMemory(); goto cmdfail; } iterator = PyObject_GetIter(docs); if (iterator == NULL) { PyObject* InvalidOperation = _error("InvalidOperation"); if (InvalidOperation) { PyErr_SetString(InvalidOperation, "input is not iterable"); Py_DECREF(InvalidOperation); } goto cmdfail; } while ((doc = PyIter_Next(iterator)) != NULL) { int sub_doc_begin = buffer_get_position(buffer); int cur_doc_begin; int cur_size; int enough_data = 0; int enough_documents = 0; char key[16]; empty = 0; INT2STRING(key, idx); if (!buffer_write_bytes(buffer, "\x03", 1) || !buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) { goto cmditerfail; } cur_doc_begin = buffer_get_position(buffer); if (!write_dict(state->_cbson, buffer, doc, check_keys, &options, 1)) { goto cmditerfail; } /* We have enough data, maybe send this batch. */ enough_data = (buffer_get_position(buffer) > max_cmd_size); enough_documents = (idx >= max_write_batch_size); if (enough_data || enough_documents) { buffer_t new_buffer; cur_size = buffer_get_position(buffer) - cur_doc_begin; /* This single document is too large for the command. */ if (!idx) { if (op == _INSERT) { _set_document_too_large(cur_size, max_bson_size); } else { PyObject* DocumentTooLarge = _error("DocumentTooLarge"); if (DocumentTooLarge) { /* * There's nothing intelligent we can say * about size for update and remove. */ PyErr_SetString(DocumentTooLarge, "command document too large"); Py_DECREF(DocumentTooLarge); } } goto cmditerfail; } if (!(new_buffer = buffer_new())) { PyErr_NoMemory(); goto cmditerfail; } /* New buffer including the current overflow document */ if (!buffer_write_bytes(new_buffer, (const char*)buffer_get_buffer(buffer), lst_len_loc + 5) || !buffer_write_bytes(new_buffer, "0\x00", 2) || !buffer_write_bytes(new_buffer, (const char*)buffer_get_buffer(buffer) + cur_doc_begin, cur_size)) { buffer_free(new_buffer); goto cmditerfail; } /* * Roll the existing buffer back to the beginning * of the last document encoded. */ buffer_update_position(buffer, sub_doc_begin); if (!buffer_write_bytes(buffer, "\x00\x00", 2)) { buffer_free(new_buffer); goto cmditerfail; } result = _send_write_command(ctx, buffer, lst_len_loc, cmd_len_loc, &errors, to_publish); buffer_free(buffer); buffer = new_buffer; if (!result) goto cmditerfail; #if PY_MAJOR_VERSION >= 3 result = Py_BuildValue("NN", PyLong_FromLong(idx_offset), result); #else result = Py_BuildValue("NN", PyInt_FromLong(idx_offset), result); #endif if (!result) goto cmditerfail; if (PyList_Append(results, result) < 0) { Py_DECREF(result); goto cmditerfail; } Py_DECREF(result); if (errors && ordered) { destroy_codec_options(&options); Py_DECREF(iterator); buffer_free(buffer); return results; } idx_offset += idx; idx = 0; Py_DECREF(to_publish); if (!(to_publish = PyList_New(0))) { goto cmditerfail; } } if (PyList_Append(to_publish, doc) < 0) { goto cmditerfail; } Py_CLEAR(doc); idx += 1; } Py_DECREF(iterator); if (PyErr_Occurred()) { goto cmdfail; } if (empty) { PyObject* InvalidOperation = _error("InvalidOperation"); if (InvalidOperation) { PyErr_SetString(InvalidOperation, "cannot do an empty bulk write"); Py_DECREF(InvalidOperation); } goto cmdfail; } if (!buffer_write_bytes(buffer, "\x00\x00", 2)) goto cmdfail; result = _send_write_command(ctx, buffer, lst_len_loc, cmd_len_loc, &errors, to_publish); if (!result) goto cmdfail; #if PY_MAJOR_VERSION >= 3 result = Py_BuildValue("NN", PyLong_FromLong(idx_offset), result); #else result = Py_BuildValue("NN", PyInt_FromLong(idx_offset), result); #endif if (!result) goto cmdfail; buffer_free(buffer); if (PyList_Append(results, result) < 0) { Py_DECREF(result); goto cmdfail; } Py_DECREF(result); Py_DECREF(to_publish); destroy_codec_options(&options); return results; cmditerfail: Py_XDECREF(doc); Py_DECREF(iterator); cmdfail: destroy_codec_options(&options); Py_DECREF(results); Py_XDECREF(to_publish); buffer_free(buffer); return NULL; }
static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow_id) { buffer_t buffer = (buffer_t)NUM2LL(rb_ary_entry(extra, 0)); VALUE check_keys = rb_ary_entry(extra, 1); if (TYPE(key) == T_SYMBOL) { // TODO better way to do this... ? key = rb_str_new2(rb_id2name(SYM2ID(key))); } if (TYPE(key) != T_STRING) { buffer_free(buffer); rb_raise(rb_eTypeError, "keys must be strings or symbols"); } if (!allow_id && strcmp("_id", RSTRING_PTR(key)) == 0) { return ST_CONTINUE; } if (check_keys == Qtrue) { int i; if (RSTRING_LEN(key) > 0 && RSTRING_PTR(key)[0] == '$') { buffer_free(buffer); rb_raise(InvalidName, "key must not start with '$'"); } for (i = 0; i < RSTRING_LEN(key); i++) { if (RSTRING_PTR(key)[i] == '.') { buffer_free(buffer); rb_raise(InvalidName, "key must not contain '.'"); } } } switch(TYPE(value)) { case T_BIGNUM: case T_FIXNUM: { if (rb_funcall(value, rb_intern(">"), 1, LL2NUM(9223372036854775807LL)) == Qtrue || rb_funcall(value, rb_intern("<"), 1, LL2NUM(-9223372036854775808LL)) == Qtrue) { buffer_free(buffer); rb_raise(rb_eRangeError, "MongoDB can only handle 8-byte ints"); } if (rb_funcall(value, rb_intern(">"), 1, INT2NUM(2147483647L)) == Qtrue || rb_funcall(value, rb_intern("<"), 1, INT2NUM(-2147483648L)) == Qtrue) { long long ll_value; write_name_and_type(buffer, key, 0x12); ll_value = NUM2LL(value); SAFE_WRITE(buffer, (char*)&ll_value, 8); } else { int int_value; write_name_and_type(buffer, key, 0x10); int_value = NUM2LL(value); SAFE_WRITE(buffer, (char*)&int_value, 4); } break; } case T_TRUE: { write_name_and_type(buffer, key, 0x08); SAFE_WRITE(buffer, &one, 1); break; } case T_FALSE: { write_name_and_type(buffer, key, 0x08); SAFE_WRITE(buffer, &zero, 1); break; } case T_FLOAT: { double d = NUM2DBL(value); write_name_and_type(buffer, key, 0x01); SAFE_WRITE(buffer, (char*)&d, 8); break; } case T_NIL: { write_name_and_type(buffer, key, 0x0A); break; } case T_HASH: { write_name_and_type(buffer, key, 0x03); write_doc(buffer, value, check_keys); break; } case T_ARRAY: { buffer_position length_location, start_position, obj_length; int items, i; VALUE* values; write_name_and_type(buffer, key, 0x04); start_position = buffer_get_position(buffer); // save space for length 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); for(i = 0; i < items; i++) { char* name; VALUE key; INT2STRING(&name, i); key = rb_str_new2(name); write_element(key, values[i], pack_extra(buffer, check_keys)); free(name); } // write null byte and fill in length 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) { buffer_position length_location, start_position, total_length; int length; write_name_and_type(buffer, key, 0x0F); 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; 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_get_position(buffer) - start_position; SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&total_length, 4); break; } else { int length; write_name_and_type(buffer, key, 0x02); value = TO_UTF8(value); length = RSTRING_LEN(value) + 1; SAFE_WRITE(buffer, (char*)&length, 4); write_utf8(buffer, value); SAFE_WRITE(buffer, &zero, 1); break; } } case T_SYMBOL: { const char* str_value = rb_id2name(SYM2ID(value)); int length = strlen(str_value) + 1; write_name_and_type(buffer, key, 0x0E); SAFE_WRITE(buffer, (char*)&length, 4); SAFE_WRITE(buffer, str_value, length); break; } case T_OBJECT: { // TODO there has to be a better way to do these checks... const char* cls = rb_class2name(RBASIC(value)->klass); if (strcmp(cls, "Mongo::Binary") == 0 || strcmp(cls, "ByteBuffer") == 0) { const char subtype = strcmp(cls, "ByteBuffer") ? (const char)FIX2INT(rb_funcall(value, rb_intern("subtype"), 0)) : 2; VALUE string_data = rb_funcall(value, rb_intern("to_s"), 0); int length = RSTRING_LEN(string_data); write_name_and_type(buffer, key, 0x05); if (subtype == 2) { const int other_length = length + 4; SAFE_WRITE(buffer, (const char*)&other_length, 4); SAFE_WRITE(buffer, &subtype, 1); } SAFE_WRITE(buffer, (const char*)&length, 4); if (subtype != 2) { SAFE_WRITE(buffer, &subtype, 1); } SAFE_WRITE(buffer, RSTRING_PTR(string_data), length); break; } if (strcmp(cls, "Mongo::ObjectID") == 0) { VALUE as_array = rb_funcall(value, rb_intern("to_a"), 0); int i; write_name_and_type(buffer, key, 0x07); for (i = 0; i < 12; i++) { char byte = (char)FIX2INT(RARRAY_PTR(as_array)[i]); SAFE_WRITE(buffer, &byte, 1); } break; } if (strcmp(cls, "Mongo::DBRef") == 0) { buffer_position length_location, start_position, obj_length; VALUE ns, oid; write_name_and_type(buffer, key, 0x03); start_position = buffer_get_position(buffer); // save space for length 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)); oid = rb_funcall(value, rb_intern("object_id"), 0); write_element(rb_str_new2("$id"), oid, pack_extra(buffer, Qfalse)); // write null byte and fill in length 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_DATA: { // TODO again, is this really the only way to do this? const char* cls = rb_class2name(RBASIC(value)->klass); if (strcmp(cls, "Time") == 0) { 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); SAFE_WRITE(buffer, (const char*)&time_since_epoch, 8); break; } } case T_REGEXP: { int length = RREGEXP_SRC_LEN(value); char* pattern = (char*)RREGEXP_SRC_PTR(value); long flags = RREGEXP(value)->ptr->options; VALUE has_extra; write_name_and_type(buffer, key, 0x0B); SAFE_WRITE(buffer, pattern, length); SAFE_WRITE(buffer, &zero, 1); if (flags & IGNORECASE) { char ignorecase = 'i'; SAFE_WRITE(buffer, &ignorecase, 1); } if (flags & MULTILINE) { char multiline = 'm'; SAFE_WRITE(buffer, &multiline, 1); } if (flags & EXTENDED) { char extended = 'x'; 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); 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); } SAFE_WRITE(buffer, &zero, 1); break; } default: { buffer_free(buffer); rb_raise(rb_eTypeError, "no c encoder for this type yet (%d)", TYPE(value)); break; } } return ST_CONTINUE; }