enforce 4MB limit when inserting BSON docs

ext/cbson/cbson.c

@@ -64,6 +64,7 @@ static VALUE RegexpOfHolding;
 static VALUE OrderedHash;
 static VALUE InvalidName;
 static VALUE InvalidStringEncoding;
+static VALUE InvalidDocument;
 static VALUE DigestMD5;
 
 #if HAVE_RUBY_ENCODING_H
@@ -455,6 +456,12 @@ static void write_doc(buffer_t buffer, VALUE hash, VALUE 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) {
+      rb_raise(InvalidDocument, "Document too large: BSON documents are limited to 4MB.");
+      return;
+    }
     SAFE_WRITE_AT_POS(buffer, length_location, &length, 4);
 }
 
@@ -799,6 +806,7 @@ void Init_cbson() {
     rb_require("mongo/errors");
     InvalidName = rb_const_get(mongo, rb_intern("InvalidName"));
     InvalidStringEncoding = rb_const_get(mongo, rb_intern("InvalidStringEncoding"));
+    InvalidDocument = rb_const_get(mongo, rb_intern("InvalidDocument"));
     rb_require("mongo/util/ordered_hash");
     OrderedHash = rb_const_get(rb_cObject, rb_intern("OrderedHash"));
 

lib/mongo/errors.rb

@@ -39,6 +39,9 @@ module Mongo
   # Raised on failures in connection to the database server.
   class ConnectionTimeoutError < MongoRubyError; end
 
+  # Raised when trying to insert a document that exceeds the 4MB limit.
+  class InvalidDocument < MongoDBError; end
+
   # Raised when a database operation fails.
   class OperationFailure < MongoDBError; end
 

lib/mongo/util/bson_ruby.rb

@@ -107,6 +107,9 @@ class BSON_RUBY
     obj.each {|k, v| serialize_key_value(k, v, check_keys) unless k == '_id' || k == :_id }
 
     serialize_eoo_element(@buf)
+    if @buf.size > 4 * 1024 * 1024
+      raise InvalidDocument, "Document is too large (#{@buf.size}). BSON documents are limited to 4MB (#{4 * 1024 * 1024})."
+    end
     @buf.put_int(@buf.size, 0)
     self
   end

test/test_bson.rb

@@ -23,6 +23,13 @@ class BSONTest < Test::Unit::TestCase
     assert_equal doc, BSON.deserialize(bson)
   end
 
+  def test_document_length
+    doc = {'name' => 'a' * 5 * 1024 * 1024}
+    assert_raise InvalidDocument do
+      assert BSON.serialize(doc)
+    end
+  end
+
   # In 1.8 we test that other string encodings raise an exception.
   # In 1.9 we test that they get auto-converted.
   if RUBY_VERSION < '1.9'