JRuby: validate key names and only move _id when necessary.

ext/java/src/org/jbson/RubyBSONEncoder.java

@@ -58,13 +58,18 @@ public class RubyBSONEncoder extends BSONEncoder {
     private RubyString _idAsString;
     private RubyString _tfAsString;
 
+    private boolean _check_keys;
+    private boolean _move_id;
+
     private static final int BIT_SIZE = 64;
     private static final long MAX = (1L << (BIT_SIZE - 1)) - 1;
     private static final BigInteger LONG_MAX = BigInteger.valueOf(MAX);
     private static final BigInteger LONG_MIN = BigInteger.valueOf(-MAX - 1);
 
 
-    public RubyBSONEncoder(Ruby runtime){
+    public RubyBSONEncoder(Ruby runtime, boolean check_keys, boolean move_id){
+      _check_keys = check_keys;
+      _move_id = move_id;
       _runtime = runtime;
       _rbclsByteBuffer = _lookupConstant( _runtime, "BSON::ByteBuffer" );
       _rbclsDBRef = _lookupConstant( _runtime, "BSON::DBRef" );
@@ -140,11 +145,10 @@ public class RubyBSONEncoder extends BSONEncoder {
         }
 
         final int sizePos = _buf.getPosition();
-        _buf.writeInt( 0 ); // leaving space for sthis.  set it at the end
+        _buf.writeInt( 0 ); // leaving space for this. set it at the end.
 
         List transientFields = null;
-        boolean rewriteID = ( myType == OBJECT && name == null );
+        boolean rewriteID = _move_id && ( myType == OBJECT && name == null );
-
 
         if ( myType == OBJECT ) {
 
@@ -175,17 +179,13 @@ public class RubyBSONEncoder extends BSONEncoder {
                  if( hashKey instanceof String) {
                      str = hashKey.toString();
                  }
-
                  else if (hashKey instanceof RubyString) {
                      str = ((RubyString)hashKey).asJavaString();
                  }
                  else if (hashKey instanceof RubySymbol) {
                      str = ((RubySymbol)hashKey).asJavaString();
                  }
-
+                // If we're rewriting the _id, we can move on.
-                 testNull(str);
-
-                 // If we're rewriting the _id, we can move on.
                  if ( rewriteID && str.equals( "_id" ) )
                     continue;
 
@@ -208,6 +208,10 @@ public class RubyBSONEncoder extends BSONEncoder {
     }
 
     protected void _putObjectField( String name , Object val ) {
+        if( _check_keys )
+            testValidKey( name );
+        else
+            testNull( name );
 
         if ( name.equals( "_transientFields" ) )
             return;
@@ -395,6 +399,22 @@ public class RubyBSONEncoder extends BSONEncoder {
        }
     }
 
+    // Make sure that name contains no null bytes, '.'s
+    // and doesn't start with a '$'.
+    private void testValidKey(String str) {
+       byte[] bytes = str.getBytes();
+
+       if( bytes[0] == 36 )
+           _rbRaise( (RubyClass)_rbclsInvalidKeyName, "$ not allowed in key name.");
+
+       for(int j = 0; j < bytes.length; j++ ) {
+         if(bytes[j] == '\u0000')
+             _rbRaise( (RubyClass)_rbclsInvalidDocument, "Null not allowed");
+         if(bytes[j] == 46)
+             _rbRaise( (RubyClass)_rbclsInvalidKeyName, ". not allowed in key name.");
+       }
+    }
+
     private void putIterable( String name , Iterable l ){
         _put( ARRAY , name );
         final int sizePos = _buf.getPosition();

lib/bson/bson_java.rb

@@ -6,7 +6,7 @@ module BSON
     # we don't create a new one on each call to #serialize.
     def self.serialize(obj, check_keys=false, move_id=false)
       raise InvalidDocument, "BSON_JAVA.serialize takes a Hash" unless obj.is_a?(Hash)
-      enc = Java::OrgJbson::RubyBSONEncoder.new(JRuby.runtime)
+      enc = Java::OrgJbson::RubyBSONEncoder.new(JRuby.runtime, check_keys, move_id)
       ByteBuffer.new(enc.encode(obj))
     end
 

test/bson/bson_test.rb

@@ -484,11 +484,9 @@ class BSONTest < Test::Unit::TestCase
                  @encoder.serialize(a, false, true).to_s
 
     # Java doesn't support this. Isn't actually necessary.
-    if !(RUBY_PLATFORM =~ /java/)
       assert_equal ")\000\000\000\002text\000\004\000\000\000abc\000\002key" +
                    "\000\004\000\000\000abc\000\020_id\000\001\000\000\000\000",
                    @encoder.serialize(a, false, false).to_s
-    end
   end
 
   def test_move_id_with_nested_doc

test/db_api_test.rb

@@ -621,7 +621,6 @@ class DBAPITest < Test::Unit::TestCase
     assert_equal("mike", @@coll.find_one()["hello"])
   end
 
-  if !RUBY_PLATFORM =~ /java/
   def test_invalid_key_names
     @@coll.remove
 
@@ -658,7 +657,6 @@ class DBAPITest < Test::Unit::TestCase
       @@coll.insert({"hello" => {"hel.lo" => "world"}})
     end
   end
-  end
 
   def test_collection_names
     assert_raise TypeError do