added additional checks for generate methods

4 files changed, 49 insertions, 25 deletions

diff --git a/lib/json/common.rb b/lib/json/common.rb
index 467c7b2..c7808fb 100644
--- a/lib/json/common.rb
+++ b/lib/json/common.rb
@@ -105,7 +105,7 @@ module JSON
 
   module_function
 
-  # Parse the JSON string _source_ into a Ruby data structure and return it.
+  # Parse the JSON document _source_ into a Ruby data structure and return it.
   #
   # _opts_ can have the following
   # keys:
@@ -122,9 +122,9 @@ module JSON
     JSON.parser.new(source, opts).parse
   end
 
-  # Parse the JSON string _source_ into a Ruby data structure and return it.
+  # Parse the JSON document _source_ into a Ruby data structure and return it.
   # The bang version of the parse method, defaults to the more dangerous values
-  # for the _opts_ hash, so be sure only to parse trusted _source_ strings.
+  # for the _opts_ hash, so be sure only to parse trusted _source_ documents.
   #
   # _opts_ can have the following keys:
   # * *max_nesting*: The maximum depth of nesting allowed in the parsed data
@@ -145,9 +145,8 @@ module JSON
     JSON.parser.new(source, opts).parse
   end
 
-  # Unparse the Ruby data structure _obj_ into a single line JSON string and
-  # return it. _state_ is
-  # * a JSON::State object,
+  # Generate a JSON document from the Ruby data structure _obj_ and return
+  # it. _state_ is * a JSON::State object,
   # * or a Hash like object (responding to to_hash),
   # * an object convertible into a hash by a to_h method,
   # that is used as or to configure a State object.
@@ -180,7 +179,11 @@ module JSON
     else
       state = State.new
     end
-    obj.to_json(state)
+    result = obj.to_json(state)
+    if result !~ /\A\s*(?:\[.*\]|\{.*\})\s*\Z/m
+      raise GeneratorError, "only generation of JSON objects or arrays allowed"
+    end
+    result
   end
 
   # :stopdoc:
@@ -190,14 +193,17 @@ module JSON
   module_function :unparse
   # :startdoc:
 
-  # Unparse the Ruby data structure _obj_ into a single line JSON string and
-  # return it. This method disables the checks for circles in Ruby objects, and
-  # also generates NaN, Infinity, and, -Infinity float values.
+  # Generate a JSON document from the Ruby data structure _obj_ and return it.
+  # This method disables the checks for circles in Ruby objects.
   #
   # *WARNING*: Be careful not to pass any Ruby data structures with circles as
   # _obj_ argument, because this will cause JSON to go into an infinite loop.
   def fast_generate(obj)
-    obj.to_json(nil)
+    result = obj.to_json(nil)
+    if result !~ /\A(?:\[.*\]|\{.*\})\Z/
+      raise GeneratorError, "only generation of JSON objects or arrays allowed"
+    end
+    result
   end
 
   # :stopdoc:
@@ -206,8 +212,9 @@ module JSON
   module_function :fast_unparse
   # :startdoc:
 
-  # Unparse the Ruby data structure _obj_ into a JSON string and return it. The
-  # returned string is a prettier form of the string returned by #unparse.
+  # Generate a JSON document from the Ruby data structure _obj_ and return it.
+  # The returned document is a prettier form of the document returned by
+  # #unparse.
   #
   # The _opts_ argument can be used to configure the generator, see the
   # generate method for a more detailed explanation.
@@ -229,7 +236,11 @@ module JSON
       end
       state.configure(opts)
     end
-    obj.to_json(state)
+    result = obj.to_json(state)
+    if result !~ /\A\s*(?:\[.*\]|\{.*\})\s*\Z/m
+      raise GeneratorError, "only generation of JSON objects or arrays allowed"
+    end
+    result
   end
 
   # :stopdoc:
@@ -270,7 +281,6 @@ module JSON
       proc.call result
     end
   end
-  module_function :recurse_proc
 
   alias restore load
   module_function :restore
diff --git a/tests/test_json.rb b/tests/test_json.rb
index 5307609..5d71d61 100755
--- a/tests/test_json.rb
+++ b/tests/test_json.rb
@@ -222,27 +222,27 @@ EOT
   def test_backslash
     data = [ '\\.(?i:gif|jpe?g|png)$' ]
     json = '["\\\\.(?i:gif|jpe?g|png)$"]'
-    assert_equal json, JSON.unparse(data)
+    assert_equal json, JSON.generate(data)
     assert_equal data, JSON.parse(json)
     #
     data = [ '\\"' ]
     json = '["\\\\\""]'
-    assert_equal json, JSON.unparse(data)
+    assert_equal json, JSON.generate(data)
     assert_equal data, JSON.parse(json)
     #
     json = '["/"]'
     data = JSON.parse(json)
     assert_equal ['/'], data
-    assert_equal json, JSON.unparse(data)
+    assert_equal json, JSON.generate(data)
     #
     json = '["\""]'
     data = JSON.parse(json)
     assert_equal ['"'], data
-    assert_equal json, JSON.unparse(data)
+    assert_equal json, JSON.generate(data)
     json = '["\\\'"]'
     data = JSON.parse(json)
     assert_equal ["'"], data
-    assert_equal '["\'"]', JSON.unparse(data)
+    assert_equal '["\'"]', JSON.generate(data)
   end
 
   def test_wrong_inputs
diff --git a/tests/test_json_generate.rb b/tests/test_json_generate.rb
index ae3ce0e..e725e6f 100755
--- a/tests/test_json_generate.rb
+++ b/tests/test_json_generate.rb
@@ -44,8 +44,8 @@ class TC_JSONGenerate < Test::Unit::TestCase
 EOT
   end
 
-  def test_unparse
-    json = unparse(@hash)
+  def test_generate
+    json = generate(@hash)
     assert_equal(JSON.parse(@json2), JSON.parse(json))
     parsed_json = parse(json)
     assert_equal(@hash, parsed_json)
@@ -53,10 +53,11 @@ EOT
     assert_equal('{"1":2}', json)
     parsed_json = parse(json)
     assert_equal({"1"=>2}, parsed_json)
+    assert_raise(GeneratorError) { generate(666) }
   end
 
-  def test_unparse_pretty
-    json = pretty_unparse(@hash)
+  def test_generate_pretty
+    json = pretty_generate(@hash)
     assert_equal(JSON.parse(@json3), JSON.parse(json))
     parsed_json = parse(json)
     assert_equal(@hash, parsed_json)
@@ -68,6 +69,19 @@ EOT
 EOT
     parsed_json = parse(json)
     assert_equal({"1"=>2}, parsed_json)
+    assert_raise(GeneratorError) { pretty_generate(666) }
+  end
+
+  def test_fast_generate
+    json = fast_generate(@hash)
+    assert_equal(JSON.parse(@json2), JSON.parse(json))
+    parsed_json = parse(json)
+    assert_equal(@hash, parsed_json)
+    json = fast_generate({1=>2})
+    assert_equal('{"1":2}', json)
+    parsed_json = parse(json)
+    assert_equal({"1"=>2}, parsed_json)
+    assert_raise(GeneratorError) { fast_generate(666) }
   end
 
   def test_states
diff --git a/tests/test_json_rails.rb b/tests/test_json_rails.rb
index 341d332..d33402d 100755
--- a/tests/test_json_rails.rb
+++ b/tests/test_json_rails.rb
@@ -141,6 +141,6 @@ EOT
   end
 
   def test_symbol
-    assert_equal '"foo"', JSON(:foo) #  we don't want an object here
+    assert_equal '"foo"', :foo.to_json #  we don't want an object here
   end
 end