diff --git a/lib/tainted_hash.rb b/lib/tainted_hash.rb index b216f39..5a53c76 100644 --- a/lib/tainted_hash.rb +++ b/lib/tainted_hash.rb @@ -29,7 +29,7 @@ def self.trigger_no_expose(hash) # # Returns a Hash. def original_hash - untaint_original_hash(@original_hash) + untaint_object(@original_hash) end # A Tainted Hash only exposes expected keys. You can either expose them @@ -241,29 +241,52 @@ def get_original_hash_value(key_s) end def set_original_hash_value(key_s, value) - if value.is_a?(Hash) && !value.is_a?(TaintedHash) - value = self.class.new(value, @new_class) - end + @original_hash[key_s] = taint_object(value) + end - @original_hash[key_s] = value + # Private: Returns an object that ensures that all elements directly + # accesible from the object will return a TaintedHash instead of a regular + # Hash. + # + # object - The object you want to taint + # + # + # Returns an object (in practice a TaintedHash, Array, or String) + def taint_object(object) + case object + when TaintedHash + object + when @new_class, Hash + self.class.new(object, @new_class) + when Array + object.map { |obj| taint_object(obj) } + else + object + end end - # Private: Returns a regular Hash, transforming all embedded TaintedHash - # objects into regular Hash objects with all keys exposed. + # Private: Returns an object that recursively turns all TaintedHash objects + # into Hash objects, with all keys and values exposed. # - # original_hash - The @original_hash you want to untaint + # object - The object you want to untaint # # - # Returns a Hash - def untaint_original_hash(original_hash) - hash = @new_class.new - original_hash.each do |key, value| - hash[key] = case value - when TaintedHash then untaint_original_hash(value.instance_variable_get :@original_hash) - else value - end + # Returns an object (in practice a Hash, Array, or String) + def untaint_object(object) + case object + when TaintedHash + untaint_object(object.original_hash) + when @new_class, Hash + hash = @new_class.new + object.each do |key, value| + hash[key] = untaint_object(value) + end + hash + when Array + object.map { |o| untaint_object(o) } + else + object end - hash end module RailsMethods diff --git a/test/tainted_hash_test.rb b/test/tainted_hash_test.rb index 306e5db..b2d76f8 100644 --- a/test/tainted_hash_test.rb +++ b/test/tainted_hash_test.rb @@ -61,6 +61,16 @@ def test_original_hash assert_equal @hash, @tainted.expose_all.original_hash end + def test_original_hash_with_mixed_nesting_of_hashes_and_arrays + hash = {'a' => 1, 'b' => [{'c' => [{'d' => 2, 'e' => 3}]}]} + tainted = TaintedHash.new hash + # Access a deeply nested value to force the inner Hash to be turned into a TaintedHash + assert_kind_of TaintedHash, tainted[:b][0][:c][0] + assert_equal hash, tainted.original_hash + assert_equal hash.merge('d' => 3), tainted.merge('d' => 3).original_hash + assert_equal hash, tainted.expose_all.original_hash + end + def test_expose_keys assert !@tainted.include?(:a) assert_equal [], @tainted.keys @@ -144,6 +154,15 @@ def test_slicing_nested_hashes assert_equal [], slice[:c].keys end + def test_mixed_nesting_of_hashes_and_arrays + hash = {'a' => 1, 'b' => [{'c' => [{'d' => 2, 'e' => 3}]}]} + tainted = TaintedHash.new hash + assert_kind_of TaintedHash, tainted[:b][0][:c][0] + assert_equal 2, tainted[:b][0][:c][0][:d] + tainted[:b][0][:c][0].expose(:d) + assert_equal %w(d), tainted[:b][0][:c][0].keys + end + def test_slicing_and_building_hashes hash = {'desc' => 'abc', 'files' => {'abc.txt' => 'abc'}} tainted = TaintedHash.new hash