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

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

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

1 files changed, 121 insertions, 204 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