Support deeply nested structures with TaintedHash

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

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