[rubygems/rubygems] Refactor Checksum classes and methods to reduce

code.
(https://github.com/rubygems/rubygems/pull/6917)

https://github.com/rubygems/rubygems/commit/2238bdaadc

8 files changed, 141 insertions, 270 deletions

diff --git a/lib/bundler/checksum.rb b/lib/bundler/checksum.rb
index 3b03935b57..fe8e73e727 100644
--- a/lib/bundler/checksum.rb
+++ b/lib/bundler/checksum.rb
@@ -2,258 +2,175 @@
 
 module Bundler
   class Checksum
-    class Store
-      attr_reader :store
-      protected :store
+    DEFAULT_BLOCK_SIZE = 16_384
+    private_constant :DEFAULT_BLOCK_SIZE
 
-      def initialize
-        @store = {}
-      end
+    class << self
+      def from_gem_source(source, digest_algorithms: %w[sha256])
+        raise ArgumentError, "not a valid gem source: #{source}" unless source.respond_to?(:with_read_io)
 
-      def initialize_copy(o)
-        @store = {}
-        o.store.each do |k, v|
-          @store[k] = v.dup
+        source.with_read_io do |io|
+          checksums = from_io(io, "#{source.path || source.inspect} gem source hexdigest", :digest_algorithms => digest_algorithms)
+          io.rewind
+          return checksums
         end
       end
 
-      def [](spec)
-        sums = @store[spec.full_name]
-
-        Checksum.new(spec.name, spec.version, spec.platform, sums&.transform_values(&:digest))
-      end
-
-      def register(spec, checksums)
-        register_full_name(spec.full_name, checksums)
-      end
-
-      def register_triple(name, version, platform, checksums)
-        register_full_name(GemHelpers.spec_full_name(name, version, platform), checksums)
-      end
-
-      def delete_full_name(full_name)
-        @store.delete(full_name)
-      end
-
-      def register_full_name(full_name, checksums)
-        sums = (@store[full_name] ||= {})
-
-        checksums.each do |checksum|
-          algo = checksum.algo
-          if multi = sums[algo]
-            multi.merge(checksum)
-          else
-            sums[algo] = Multi.new [checksum]
-          end
+      def from_io(io, source, digest_algorithms: %w[sha256])
+        digests = digest_algorithms.to_h do |algo|
+          [algo.to_s, Bundler::SharedHelpers.digest(algo.upcase).new]
         end
-      rescue SecurityError => e
-        raise e.exception(<<~MESSAGE)
-          Bundler found multiple different checksums for #{full_name}.
-          This means that there are multiple different `#{full_name}.gem` files.
-          This is a potential security issue, since Bundler could be attempting \
-          to install a different gem than what you expect.
-
-          #{e.message}
-          To resolve this issue:
-          1. delete any downloaded gems referenced above
-          2. run `bundle install`
-
-          If you are sure that the new checksum is correct, you can \
-          remove the `#{full_name}` entry under the lockfile `CHECKSUMS` \
-          section and rerun `bundle install`.
 
-          If you wish to continue installing the downloaded gem, and are certain it does not pose a \
-          security issue despite the mismatching checksum, do the following:
-          1. run `bundle config set --local disable_checksum_validation true` to turn off checksum verification
-          2. run `bundle install`
-        MESSAGE
-      end
+        until io.eof?
+          ret = io.read DEFAULT_BLOCK_SIZE
+          digests.each_value {|digest| digest << ret }
+        end
 
-      def use(other)
-        other.store.each do |k, v|
-          register_full_name k, v.values
+        digests.map do |algo, digest|
+          Checksum.new(algo, digest.hexdigest!, source)
         end
       end
     end
 
-    class Single
-      attr_reader :algo, :digest, :source
-      def initialize(algo, digest, source)
-        @algo = algo
-        @digest = digest
-        @source = source
-      end
-
-      def ==(other)
-        other.is_a?(Single) && other.digest == digest && other.algo == algo && source == other.source
-      end
+    attr_reader :algo, :digest, :sources
+    def initialize(algo, digest, source)
+      @algo = algo
+      @digest = digest
+      @sources = [source]
+    end
 
-      def hash
-        digest.hash
-      end
+    def ==(other)
+      match?(other) && other.sources == sources
+    end
 
-      alias_method :eql?, :==
+    alias_method :eql?, :==
 
-      def to_s
-        "#{algo}-#{digest} (from #{source})"
-      end
+    def match?(other)
+      other.is_a?(self.class) && other.digest == digest && other.algo == algo
     end
 
-    class Multi
-      attr_reader :algo, :digest, :checksums
-      protected :checksums
+    def hash
+      digest.hash
+    end
 
-      def initialize(checksums)
-        @checksums = checksums
+    def to_s
+      "#{to_lock} (from #{sources.first}#{", ..." if sources.size > 1})"
+    end
 
-        unless checksums && checksums.size > 0
-          raise ArgumentError, "must provide at least one checksum"
-        end
+    def to_lock
+      "#{algo}-#{digest}"
+    end
 
-        first = checksums.first
-        @algo = first.algo
-        @digest = first.digest
-      end
+    def merge!(other)
+      raise ArgumentError, "cannot merge checksums of different algorithms" unless algo == other.algo
 
-      def initialize_copy(o)
-        @checksums = o.checksums.dup
-        @algo = o.algo
-        @digest = o.digest
+      unless digest == other.digest
+        raise SecurityError, <<~MESSAGE
+          #{other}
+          #{to_lock} from:
+          * #{sources.join("\n* ")}
+        MESSAGE
       end
 
-      def merge(other)
-        raise ArgumentError, "cannot merge checksums of different algorithms" unless algo == other.algo
-        unless digest == other.digest
-          raise SecurityError, <<~MESSAGE
-            #{other}
-            #{self} from:
-            * #{sources.join("\n* ")}
-          MESSAGE
-        end
-
-        case other
-        when Single
-          @checksums << other
-        when Multi
-          @checksums.concat(other.checksums)
-        else
-          raise ArgumentError
-        end
-        @checksums.uniq!
+      @sources.concat(other.sources).uniq!
+      self
+    end
 
-        self
-      end
+    class Store
+      attr_reader :store
+      protected :store
 
-      def sources
-        @checksums.map(&:source)
+      def initialize
+        @store = {}
       end
 
-      def to_s
-        "#{algo}-#{digest}"
+      def initialize_copy(other)
+        @store = {}
+        other.store.each do |full_name, checksums|
+          store[full_name] = checksums.dup
+        end
       end
-    end
 
-    attr_reader :name, :version, :platform, :checksums
-
-    SHA256 = %r{\Asha256-([a-z0-9]{64}|[A-Za-z0-9+\/=]{44})\z}.freeze
-    private_constant :SHA256
-
-    def initialize(name, version, platform, checksums = {})
-      @name     = name
-      @version  = version
-      @platform = platform || Gem::Platform::RUBY
-      @checksums = checksums || {}
-
-      # can expand this validation when we support more hashing algos later
-      if !@checksums.is_a?(::Hash) || (@checksums.any? && !@checksums.key?("sha256"))
-        raise ArgumentError, "invalid checksums (#{@checksums.inspect})"
-      end
-      if @checksums.any? {|_, checksum| !checksum.is_a?(String) }
-        raise ArgumentError, "invalid checksums (#{@checksums})"
+      def checksums(full_name)
+        store[full_name]
       end
-    end
-
-    def self.digests_from_file_source(file_source, digest_algorithms: %w[sha256])
-      raise ArgumentError, "not a valid file source: #{file_source}" unless file_source.respond_to?(:with_read_io)
 
-      digests = digest_algorithms.map do |digest_algorithm|
-        [digest_algorithm.to_s, Bundler::SharedHelpers.digest(digest_algorithm.upcase).new]
-      end.to_h
-
-      file_source.with_read_io do |io|
-        until io.eof?
-          block = io.read(16_384)
-          digests.each_value {|digest| digest << block }
+      def register_gem_package(spec, source)
+        new_checksums = Checksum.from_gem_source(source)
+        new_checksums.each do |checksum|
+          register spec.full_name, checksum
         end
+      rescue SecurityError
+        expected = checksums(spec.full_name)
+        gem_lock_name = GemHelpers.lock_name(spec.name, spec.version, spec.platform)
+        raise SecurityError, <<~MESSAGE
+          Bundler cannot continue installing #{gem_lock_name}.
+          The checksum for the downloaded `#{spec.full_name}.gem` does not match \
+          the known checksum for the gem.
+          This means the contents of the downloaded \
+          gem is different from what was uploaded to the server \
+          or first used by your teammates, and could be a potential security issue.
 
-        io.rewind
-      end
+          To resolve this issue:
+          1. delete the downloaded gem located at: `#{source.path}`
+          2. run `bundle install`
 
-      digests
-    end
+          If you are sure that the new checksum is correct, you can \
+          remove the `#{gem_lock_name}` entry under the lockfile `CHECKSUMS` \
+          section and rerun `bundle install`.
 
-    def full_name
-      GemHelpers.spec_full_name(@name, @version, @platform)
-    end
+          If you wish to continue installing the downloaded gem, and are certain it does not pose a \
+          security issue despite the mismatching checksum, do the following:
+          1. run `bundle config set --local disable_checksum_validation true` to turn off checksum verification
+          2. run `bundle install`
 
-    def match_spec?(spec)
-      name == spec.name &&
-        version == spec.version &&
-        platform.to_s == spec.platform.to_s
-    end
+          #{expected.map do |checksum|
+            next unless actual = new_checksums.find {|c| c.algo == checksum.algo }
+            next if actual.digest == checksum.digest
 
-    def to_lock
-      out = String.new
-      out << "  #{GemHelpers.lock_name(name, version, platform)}"
-      checksums.sort_by(&:first).each_with_index do |(algo, checksum), idx|
-        out << (idx.zero? ? " " : ",")
-        out << algo << "-" << checksum
+            "(More info: The expected #{checksum.algo.upcase} checksum was #{checksum.digest.inspect}, but the " \
+            "checksum for the downloaded gem was #{actual.digest.inspect}. The expected checksum came from: #{checksum.sources.join(", ")})"
+          end.compact.join("\n")}
+        MESSAGE
       end
-      out << "\n"
 
-      out
-    end
+      def register(full_name, checksum)
+        return unless checksum
 
-    def match?(other)
-      return false unless match_spec?(other)
-      match_digests?(other.checksums)
-    end
+        sums = (store[full_name] ||= [])
+        sums.find {|c| c.algo == checksum.algo }&.merge!(checksum) || sums << checksum
+      rescue SecurityError => e
+        raise e.exception(<<~MESSAGE)
+          Bundler found multiple different checksums for #{full_name}.
+          This means that there are multiple different `#{full_name}.gem` files.
+          This is a potential security issue, since Bundler could be attempting \
+          to install a different gem than what you expect.
 
-    def match_digests?(digests)
-      return true if checksums.empty? && digests.empty?
+          #{e.message}
+          To resolve this issue:
+          1. delete any downloaded gems referenced above
+          2. run `bundle install`
 
-      common_algos = checksums.keys & digests.keys
-      return true if common_algos.empty?
+          If you are sure that the new checksum is correct, you can \
+          remove the `#{full_name}` entry under the lockfile `CHECKSUMS` \
+          section and rerun `bundle install`.
 
-      common_algos.all? do |algo|
-        checksums[algo] == digests[algo]
+          If you wish to continue installing the downloaded gem, and are certain it does not pose a \
+          security issue despite the mismatching checksum, do the following:
+          1. run `bundle config set --local disable_checksum_validation true` to turn off checksum verification
+          2. run `bundle install`
+        MESSAGE
       end
-    end
-
-    def merge!(other)
-      raise ArgumentError, "can't merge checksums for different specs" unless match_spec?(other)
 
-      merge_digests!(other.checksums)
-    end
-
-    def merge_digests!(digests)
-      if digests.any? {|_, checksum| !checksum.is_a?(String) }
-        raise ArgumentError, "invalid checksums (#{digests})"
+      def replace(full_name, checksum)
+        store[full_name] = checksum ? [checksum] : nil
       end
-      @checksums = @checksums.merge(digests) do |algo, ours, theirs|
-        if ours != theirs
-          raise ArgumentError, "Digest mismatch for #{algo}:\n\t* #{ours.inspect}\n\t* #{theirs.inspect}"
+
+      def register_store(other)
+        other.store.each do |full_name, checksums|
+          checksums.each {|checksum| register(full_name, checksum) }
         end
-        ours
       end
-
-      self
-    end
-
-    private
-
-    def sha256
-      @checksums.find {|c| c =~ SHA256 }
     end
   end
 end
diff --git a/lib/bundler/definition.rb b/lib/bundler/definition.rb
index 26e3126d84..5c605c47ac 100644
--- a/lib/bundler/definition.rb
+++ b/lib/bundler/definition.rb
@@ -753,7 +753,7 @@ module Bundler
         # has to be done separately, because we want to keep the locked checksum
         # store for a source, even when doing a full update
         if @locked_gems && locked_source = @locked_gems.sources.find {|s| s == source }
-          source.checksum_store&.use(locked_source.checksum_store)
+          source.checksum_store.register_store(locked_source.checksum_store)
         end
         # If the source is unlockable and the current command allows an unlock of
         # the source (for example, you are doing a `bundle update <foo>` of a git-pinned
diff --git a/lib/bundler/endpoint_specification.rb b/lib/bundler/endpoint_specification.rb
index 943e33be94..11c71d1ae7 100644
--- a/lib/bundler/endpoint_specification.rb
+++ b/lib/bundler/endpoint_specification.rb
@@ -135,7 +135,7 @@ module Bundler
           else
             raise ArgumentError, "The given checksum for #{full_name} (#{digest.inspect}) is not a valid SHA256 hexdigest nor base64digest"
           end
-          @checksum = Checksum::Single.new("sha256", digest, "API response from #{@spec_fetcher.uri}")
+          @checksum = Checksum.new("sha256", digest, "API response from #{@spec_fetcher.uri}")
         when "rubygems"
           @required_rubygems_version = Gem::Requirement.new(v)
         when "ruby"
diff --git a/lib/bundler/fetcher.rb b/lib/bundler/fetcher.rb
index 08d1ee3437..d493ca0064 100644
--- a/lib/bundler/fetcher.rb
+++ b/lib/bundler/fetcher.rb
@@ -140,14 +140,11 @@ module Bundler
       fetch_specs(gem_names).each do |name, version, platform, dependencies, metadata|
         spec = if dependencies
           EndpointSpecification.new(name, version, platform, self, dependencies, metadata).tap do |es|
-            unless index.local_search(es).empty?
-              # Duplicate spec.full_names, different spec.original_names
-              # index#<< ensures that the last one added wins, so if we're overriding
-              # here, make sure to also override the checksum, otherwise downloading the
-              # specs (even if that version is completely unused) will cause a SecurityError
-              source.checksum_store.delete_full_name(es.full_name)
-            end
-            source.checksum_store.register(es, [es.checksum]) if source && es.checksum
+            # Duplicate spec.full_names, different spec.original_names
+            # index#<< ensures that the last one added wins, so if we're overriding
+            # here, make sure to also override the checksum, otherwise downloading the
+            # specs (even if that version is completely unused) will cause a SecurityError
+            source.checksum_store.replace(es.full_name, es.checksum)
           end
         else
           RemoteSpecification.new(name, version, platform, self)
diff --git a/lib/bundler/lockfile_generator.rb b/lib/bundler/lockfile_generator.rb
index 4d2c2c2a86..8114c27917 100644
--- a/lib/bundler/lockfile_generator.rb
+++ b/lib/bundler/lockfile_generator.rb
@@ -69,10 +69,12 @@ module Bundler
     def add_checksums
       out << "\nCHECKSUMS\n"
 
-      empty_store = Checksum::Store.new
-
       definition.resolve.sort_by(&:full_name).each do |spec|
-        out << (spec.source.checksum_store || empty_store)[spec].to_lock
+        lock_name = GemHelpers.lock_name(spec.name, spec.version, spec.platform)
+        out << "  #{lock_name}"
+        checksums = spec.source.checksum_store.checksums(spec.full_name)
+        out << " #{checksums.map(&:to_lock).sort.join(",")}" if checksums
+        out << "\n"
       end
     end
 
diff --git a/lib/bundler/lockfile_parser.rb b/lib/bundler/lockfile_parser.rb
index 43d544fd32..9ffe5beffd 100644
--- a/lib/bundler/lockfile_parser.rb
+++ b/lib/bundler/lockfile_parser.rb
@@ -225,19 +225,16 @@ module Bundler
 
       version = Gem::Version.new(version)
       platform = platform ? Gem::Platform.new(platform) : Gem::Platform::RUBY
-      source = "#{@lockfile_path}:#{@pos} in the CHECKSUMS lockfile section"
-      checksums = checksums.split(",").map do |c|
-        algo, digest = c.split("-", 2)
-        Checksum::Single.new(algo, digest, source)
-      end
-
       full_name = GemHelpers.spec_full_name(name, version, platform)
-
       # Don't raise exception if there's a checksum for a gem that's not in the lockfile,
       # we prefer to heal invalid lockfiles
       return unless spec = @specs[full_name]
 
-      spec.source.checksum_store.register_full_name(full_name, checksums)
+      checksums.split(",").each do |c|
+        algo, digest = c.split("-", 2)
+        lock_name = GemHelpers.lock_name(spec.name, spec.version, spec.platform)
+        spec.source.checksum_store.register(full_name, Checksum.new(algo, digest, "#{@lockfile_path}:#{@pos} CHECKSUMS #{lock_name}"))
+      end
     end
 
     def parse_spec(line)
diff --git a/lib/bundler/rubygems_gem_installer.rb b/lib/bundler/rubygems_gem_installer.rb
index 0a289c416f..c381956fc3 100644
--- a/lib/bundler/rubygems_gem_installer.rb
+++ b/lib/bundler/rubygems_gem_installer.rb
@@ -117,54 +117,11 @@ module Bundler
 
     def validate_bundler_checksum(checksum_store)
       return true if Bundler.settings[:disable_checksum_validation]
-
       return true unless source = @package.instance_variable_get(:@gem)
       return true unless source.respond_to?(:with_read_io)
-      digests = Bundler::Checksum.digests_from_file_source(source).transform_values(&:hexdigest!)
-
-      checksum = checksum_store[spec]
-      unless checksum.match_digests?(digests)
-        expected = checksum_store.send(:store)[spec.full_name]
-
-        raise SecurityError, <<~MESSAGE
-          Bundler cannot continue installing #{spec.name} (#{spec.version}).
-          The checksum for the downloaded `#{spec.full_name}.gem` does not match \
-          the known checksum for the gem.
-          This means the contents of the downloaded \
-          gem is different from what was uploaded to the server \
-          or first used by your teammates, and could be a potential security issue.
-
-          To resolve this issue:
-          1. delete the downloaded gem located at: `#{source.path}`
-          2. run `bundle install`
-
-          If you are sure that the new checksum is correct, you can \
-          remove the `#{GemHelpers.lock_name spec.name, spec.version, spec.platform}` entry under the lockfile `CHECKSUMS` \
-          section and rerun `bundle install`.
-
-          If you wish to continue installing the downloaded gem, and are certain it does not pose a \
-          security issue despite the mismatching checksum, do the following:
-          1. run `bundle config set --local disable_checksum_validation true` to turn off checksum verification
-          2. run `bundle install`
-
-          #{expected.map do |algo, multi|
-            next unless actual = digests[algo]
-            next if actual == multi
-
-            "(More info: The expected #{algo.upcase} checksum was #{multi.digest.inspect}, but the " \
-            "checksum for the downloaded gem was #{actual.inspect}. The expected checksum came from: #{multi.sources.join(", ")})"
-          end.compact.join("\n")}
-          MESSAGE
-      end
-      register_digests(digests, checksum_store, source)
-      true
-    end
 
-    def register_digests(digests, checksum_store, source)
-      checksum_store.register(
-        spec,
-        digests.map {|algo, digest| Checksum::Single.new(algo, digest, "downloaded gem @ `#{source.path}`") }
-      )
+      checksum_store.register_gem_package spec, source
+      true
     end
   end
 end
diff --git a/lib/bundler/source/git.rb b/lib/bundler/source/git.rb
index adbce5fce4..24d871bd11 100644
--- a/lib/bundler/source/git.rb
+++ b/lib/bundler/source/git.rb
@@ -11,6 +11,7 @@ module Bundler
 
       def initialize(options)
         @options = options
+        @checksum_store = Checksum::Store.new
         @glob = options["glob"] || DEFAULT_GLOB
 
         @allow_cached = false