From aaf590214cea92e5fb6315ecffaa3837e0d02f83 Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Wed, 16 Dec 2009 11:32:54 -0500
Subject: [PATCH] enforce 4MB limit when inserting BSON docs

---
 ext/cbson/cbson.c           | 8 ++++++++
 lib/mongo/errors.rb         | 5 ++++-
 lib/mongo/util/bson_ruby.rb | 3 +++
 test/test_bson.rb           | 7 +++++++
 4 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/ext/cbson/cbson.c b/ext/cbson/cbson.c
index 41152d6..15d002f 100644
--- a/ext/cbson/cbson.c
+++ b/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"));
 
diff --git a/lib/mongo/errors.rb b/lib/mongo/errors.rb
index a8b4260..b67f1ff 100644
--- a/lib/mongo/errors.rb
+++ b/lib/mongo/errors.rb
@@ -39,9 +39,12 @@ 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
-  
+
   # Raised when a database operation fails.
   class ConnectionFailure < MongoDBError; end
 
diff --git a/lib/mongo/util/bson_ruby.rb b/lib/mongo/util/bson_ruby.rb
index 44481ac..f15ab2a 100644
--- a/lib/mongo/util/bson_ruby.rb
+++ b/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
diff --git a/test/test_bson.rb b/test/test_bson.rb
index 22425de..6e52ce5 100644
--- a/test/test_bson.rb
+++ b/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'