Resolve #136 and control growth of containers

5 files changed, 66 insertions, 21 deletions

diff --git a/History.md b/History.md
index f85ad87..6f5f364 100644
--- a/History.md
+++ b/History.md
@@ -23,6 +23,15 @@
         Set objects so they can be reduced to a single Set. [#117][], [#127][],
         [#134][].
 
+    *   Fixed an uncontrolled growth bug in MIME::Types::Container where a key
+        miss would create a new entry with an empty Set in the container. This
+        was working as designed (this particular feature was heavily used
+        during MIME::Type registry construction), but the design was flawed in
+        that it did not have any way of determining the difference between
+        construction and querying. This would mean that, if you have a function
+        in your web app that queries the MIME::Types registry by extension, the
+        extension registry would grow uncontrollably. [#136][]
+
 *   Deprecations:
 
     *   Lazy loading (`$RUBY_MIME_TYPES_LAZY_LOAD`) has been deprecated.
@@ -163,16 +172,24 @@
             The internal format of the columnar store has changed; many of the
             boolean flags are now loaded from a single file. Resolves [#85][].
 
-[#112]: https://github.com/mime-types/ruby-mime-types/pull/112
-[#117]: https://github.com/mime-types/ruby-mime-types/pull/117
-[#118]: https://github.com/mime-types/ruby-mime-types/pull/118
-[#132]: https://github.com/mime-types/ruby-mime-types/pull/132
-[#135]: https://github.com/mime-types/ruby-mime-types/pull/135
 [#79]: https://github.com/mime-types/ruby-mime-types/pull/79
 [#84]: https://github.com/mime-types/ruby-mime-types/pull/84
 [#85]: https://github.com/mime-types/ruby-mime-types/pull/85
 [#95]: https://github.com/mime-types/ruby-mime-types/pull/95
 [#97]: https://github.com/mime-types/ruby-mime-types/pull/97
+[#112]: https://github.com/mime-types/ruby-mime-types/pull/112
+[#117]: https://github.com/mime-types/ruby-mime-types/issues/117
+[#118]: https://github.com/mime-types/ruby-mime-types/pull/118
+[#125]: https://github.com/mime-types/ruby-mime-types/pull/125
+[#126]: https://github.com/mime-types/ruby-mime-types/pull/126
+[#127]: https://github.com/mime-types/ruby-mime-types/issues/127
+[#129]: https://github.com/mime-types/ruby-mime-types/pull/129
+[#130]: https://github.com/mime-types/ruby-mime-types/pull/130
+[#127]: https://github.com/mime-types/ruby-mime-types/issues/127
+[#132]: https://github.com/mime-types/ruby-mime-types/pull/132
+[#134]: https://github.com/mime-types/ruby-mime-types/issues/134
+[#135]: https://github.com/mime-types/ruby-mime-types/pull/135
+[#136]: https://github.com/mime-types/ruby-mime-types/issues/136
 [Code-of-Conduct.md]: Code-of-Conduct_md.html
 [Contributor Covenant]: http://contributor-covenant.org
 [mime-types-data]: https://github.com/mime-types/mime-types-data
diff --git a/lib/mime/types.rb b/lib/mime/types.rb
index fd56908..9cfd52e 100644
--- a/lib/mime/types.rb
+++ b/lib/mime/types.rb
@@ -197,7 +197,7 @@ Type #{type} is already registered as a variant of #{type.simplified}.
   private
 
   def add_type_variant!(mime_type)
-    @type_variants[mime_type.simplified] << mime_type
+    @type_variants.add(mime_type.simplified, mime_type)
   end
 
   def reindex_extensions!(mime_type)
@@ -206,7 +206,7 @@ Type #{type} is already registered as a variant of #{type.simplified}.
   end
 
   def index_extensions!(mime_type)
-    mime_type.extensions.each { |ext| @extension_index[ext] << mime_type }
+    mime_type.extensions.each { |ext| @extension_index.add(ext, mime_type) }
   end
 
   def prune_matches(matches, complete, registered)
diff --git a/lib/mime/types/container.rb b/lib/mime/types/container.rb
index abccc2e..e56ee2a 100644
--- a/lib/mime/types/container.rb
+++ b/lib/mime/types/container.rb
@@ -1,32 +1,58 @@
 # frozen_string_literal: true
 
 require 'set'
+require 'forwardable'
+
+# MIME::Types requires a serializable keyed container that returns an empty Set
+# on a key miss. Hash#default_value cannot be used because, while it traverses
+# the Marshal format correctly, it won't survive any other serialization
+# format (plus, a default of a mutable object resuls in a shared mess).
+# Hash#default_proc cannot be used without a wrapper because it prevents
+# Marshal serialization (and doesn't survive the round-trip).
+class MIME::Types::Container
+  extend Forwardable
 
-# MIME::Types requires a container Hash with a default values for keys
-# resulting in an empty Set, but this cannot be dumped through Marshal because
-# of the presence of that default Proc. This class exists solely to satisfy
-# that need.
-class MIME::Types::Container < Hash # :nodoc:
   def initialize
-    super
-    self.default_proc = ->(h, k) { h[k] = Set.new }
+    @container = {}
+  end
+
+  def [](key)
+    container[key] || EMPTY_SET
+  end
+
+  def_delegators :@container,
+    :count,
+    :each_value,
+    :keys,
+    :merge,
+    :merge!,
+    :select,
+    :values
+
+  def add(key, value)
+    (container[key] ||= Set.new).add(value)
   end
 
   def marshal_dump
-    {}.merge(self)
+    {}.merge(container)
   end
 
   def marshal_load(hash)
-    self.default_proc = ->(h, k) { h[k] = Set.new }
-    merge!(hash)
+    @container = hash
   end
 
   def encode_with(coder)
-    each { |k, v| coder[k] = v.to_a }
+    container.each { |k, v| coder[k] = v.to_a }
   end
 
   def init_with(coder)
-    self.default_proc = ->(h, k) { h[k] = Set.new }
-    coder.map.each { |k, v| self[k] = Set[*v] }
+    coder.map.each { |k, v| container[k] = Set[*v] }
   end
+
+  private
+
+  attr_reader :container
+
+  EMPTY_SET = Set.new.freeze
+  private_constant :EMPTY_SET
 end
diff --git a/test/test_mime_types.rb b/test/test_mime_types.rb
index 111f771..f4528a3 100644
--- a/test/test_mime_types.rb
+++ b/test/test_mime_types.rb
@@ -151,7 +151,9 @@ describe MIME::Types do
     end
 
     it 'does not find unknown extensions' do
+      keys = mime_types.instance_variable_get(:@extension_index).keys
       assert_empty mime_types.type_for('zzz')
+      assert_equal keys, mime_types.instance_variable_get(:@extension_index).keys
     end
 
     it 'modifying type extensions causes reindexing' do
diff --git a/test/test_mime_types_cache.rb b/test/test_mime_types_cache.rb
index d8cfb3a..2209cab 100644
--- a/test/test_mime_types_cache.rb
+++ b/test/test_mime_types_cache.rb
@@ -104,7 +104,7 @@ end
 describe MIME::Types::Container do
   it 'marshals and unmarshals correctly' do
     container = MIME::Types::Container.new
-    container['xyz'] << 'abc'
+    container.add('xyz', 'abc')
 
     # default proc should return Set[]
     assert_equal(Set[], container['abc'])